Update app.py

app.py

@@ -6,27 +6,27 @@ model = YOLO("best.pt")
 
 def predict(img):
     results = model(img)
-    # This draws the boxes and labels on the image
     annotated_img = results[0].plot() 
     
-    # Extract the labels and confidences to send back as data
     detections = []
-    for box in results[0].boxes:
-        class_id = int(box.cls[0])
-        label = model.names[class_id]
-        conf = float(box.conf[0])
-        detections.append({
-            "label": label, 
-            "confidence": conf
-        })
+    # Safely extract boxes if they exist
+    if len(results) > 0 and len(results[0].boxes) > 0:
+        for box in results[0].boxes:
+            # Use .item() to safely convert PyTorch tensors to standard Python numbers
+            class_id = int(box.cls.item())
+            label = model.names[class_id]
+            conf = float(box.conf.item())
+            
+            detections.append({
+                "label": label, 
+                "confidence": conf
+            })
         
-    # We now return BOTH the image and a JSON string containing the detection data
     return annotated_img, json.dumps(detections)
 
 demo = gr.Interface(
     fn=predict, 
     inputs=gr.Image(type="pil"), 
-    # Update the outputs to expect both an Image and a Textbox (for our JSON)
     outputs=[
         gr.Image(type="pil"), 
         gr.Textbox(label="Detections")