Update app.py

app.py

@@ -1,38 +1,34 @@
 import gradio as gr
 from ultralytics import YOLO
-import torch
+import numpy as np
 
-model = YOLO("best.pt")  # Make sure best.pt exists in the same folder
+# Load YOLOv8 model
+model = YOLO("best.pt")  # Make sure best.pt is in the repo root
 
 def predict(image):
-    try:
-        results = model(image)[0]
+    # image is numpy array from Gradio
+    results = model(image)[0]  # YOLO prediction
 
-        detections = []
-        for box in results.boxes:
-            xmin, ymin, xmax, ymax = box.xyxy[0].tolist()
-            cls = int(box.cls)
-            conf = float(box.conf)
+    annotated_image = results.plot()  # returns numpy array
+    detections = []
 
+    if results.boxes is not None:
+        for box, cls, conf in zip(results.boxes.xyxy, results.boxes.cls, results.boxes.conf):
             detections.append({
-                "label": model.names[cls],
-                "confidence": conf,
-                "box": [xmin, ymin, xmax, ymax]
+                "label": model.names[int(cls)],
+                "confidence": float(conf),
+                "box": [float(coord) for coord in box]
             })
 
-        return image, detections
-    except Exception as e:
-        print("ERROR:", e)
-        return image, []
+    return annotated_image, detections
 
+# Gradio interface
 demo = gr.Interface(
     fn=predict,
-    inputs=gr.Image(type="numpy", label="Upload Image"),
-    outputs=[
-        gr.Image(type="numpy", label="Processed Image"),
-        gr.JSON(label="Detections")
-    ],
-    title="YOLO Helmet Detection"
+    inputs=gr.Image(type="numpy"),
+    outputs=[gr.Image(type="numpy"), gr.JSON()],
+    title="Helmet Detection YOLOv8",
+    description="Upload an image and detect helmet / head using YOLOv8"
 )
 
 demo.launch()