Update app.py

app.py

@@ -12,6 +12,7 @@ from torchvision.transforms import functional as F
 from PIL import Image
 import cv2
 import gradio as gr
+import numpy as np  # Add this import
 
 from yolov5.models.yolo import Model
 from yolov5.utils.general import non_max_suppression
@@ -22,6 +23,7 @@ model.eval()
 
 def preprocess_image(image):
     try:
+        image = image.convert("RGB")  # Ensure image is in RGB mode
         image_tensor = F.to_tensor(image)
         return image_tensor.unsqueeze(0).to(device)
     except Exception as e:
@@ -34,12 +36,13 @@ def draw_boxes(image, outputs, threshold=0.3):
         h, w, _ = image.shape
 
         for box in outputs:
-            score, label, x1, y1, x2, y2 = box[4].item(), int(box[5].item()), box[0].item(), box[1].item(), box[2].item(), box[3].item()
-            if score > threshold:
-                x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
-                cv2.rectangle(image, (x1, y1), (x2, y2), (255, 0, 0), 2)
-                text = f"{model.names[label]:s}: {score:.2f}"
-                cv2.putText(image, text, (x1, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 2)
+            if box is not None:
+                x1, y1, x2, y2, score, label = box[:6]
+                if score > threshold:
+                    x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
+                    cv2.rectangle(image, (x1, y1), (x2, y2), (255, 0, 0), 2)
+                    text = f"{model.names[int(label)]:s}: {score:.2f}"
+                    cv2.putText(image, text, (x1, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 2)
 
         return cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
     except Exception as e:
@@ -51,8 +54,8 @@ def detect_objects(image):
     if image_tensor is None:
         return image
     try:
-        outputs = model(image_tensor)
-        outputs = non_max_suppression(outputs)[0]
+        outputs = model(image_tensor)[0]  # Get the first element of the output
+        outputs = non_max_suppression(outputs, conf_thres=0.25, iou_thres=0.45)[0]  # Apply NMS
         result_image = draw_boxes(image, outputs)
         return result_image
     except Exception as e: