Create app.py

app.py

@@ -0,0 +1,55 @@
+import os
+import torch
+import gradio as gr
+from PIL import Image
+
+# Step 1: Search for best.pt in the training directory
+base_path = "yolov5/runs/train/"
+best_path = None
+
+# Search through the directory structure to find best.pt
+for root, dirs, files in os.walk(base_path):
+    if "best.pt" in files:
+        best_path = os.path.join(root, "best.pt")
+        break
+
+# Step 2: If best.pt is not found, use pre-trained weights
+if best_path is None:
+    print("Trained weights (best.pt) not found.")
+    print("Using pre-trained YOLOv5 weights (yolov5s.pt) instead.")
+
+    # Fallback to a pre-trained model if best.pt is not available
+    model = torch.hub.load('ultralytics/yolov5', 'yolov5s')  # Load pre-trained weights
+
+else:
+    print(f"Model weights found at: {best_path}")
+    # Load YOLOv5 model with the correct path
+    model = torch.hub.load('ultralytics/yolov5', 'custom', path=best_path)
+
+# Step 3: Detection Function
+def detect_weapons(image):
+    results = model(image)
+    detected_classes = results.pandas().xyxy[0]['name'].unique()
+    
+    # Check for threats
+    threat_message = "Threat detected: Be careful" if len(detected_classes) > 0 else "No threat detected"
+    return threat_message, Image.fromarray(results.render()[0])
+
+# Step 4: Gradio Interface
+def inference(image):
+    threat, detected_image = detect_weapons(image)
+    return threat, detected_image
+
+iface = gr.Interface(
+    fn=inference,
+    inputs=gr.Image(type="numpy", label="Upload Image"),
+    outputs=[
+        gr.Textbox(label="Threat Detection"),
+        gr.Image(label="Detected Image"),
+    ],
+    title="Weapon Detection AI",
+    description="Upload an image to detect weapons like bombs, guns, and pistols."
+)
+
+# Step 5: Launch Gradio App
+iface.launch()