Update app.py

app.py

@@ -1,4 +1,4 @@
-from ultralytics import YOLO  # Use Ultralytics' YOLO module
+from ultralytics import YOLO  
 import torch
 import cv2
 import numpy as np
@@ -11,6 +11,9 @@ model = YOLO("yolov5s.pt")  # Load pre-trained YOLOv5s model
 model.to(device)
 model.eval()
 
+# Load COCO class labels
+CLASS_NAMES = model.names  # YOLOv5's built-in class names
+
 def preprocess_image(image):
     image = Image.fromarray(image)
     image = image.convert("RGB")
@@ -21,16 +24,18 @@ def detect_objects(image):
     results = model.predict(image)  # Run YOLOv5 inference
 
     # Convert results to bounding box format
-    detections = []
+    image = np.array(image)
     for result in results:
-        for box in result.boxes.xyxy:
+        for box, cls in zip(result.boxes.xyxy, result.boxes.cls):
             x1, y1, x2, y2 = map(int, box[:4])
-            detections.append([x1, y1, x2, y2])
+            class_name = CLASS_NAMES[int(cls)]  # Get class name
 
-    # Draw bounding boxes
-    image = np.array(image)
-    for x1, y1, x2, y2 in detections:
-        cv2.rectangle(image, (x1, y1), (x2, y2), (255, 0, 0), 2)
+            # Draw bounding box
+            cv2.rectangle(image, (x1, y1), (x2, y2), (255, 0, 0), 2)
+
+            # Put class label
+            cv2.putText(image, class_name, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 
+                        0.5, (255, 0, 0), 2, cv2.LINE_AA)
 
     return image