Update app.py

app.py

@@ -7,24 +7,43 @@ from ultralytics import YOLO
 # Load the YOLOv8 model
 model = YOLO('best.pt')
 
+# Define a list of colors for different classes
+colors = [
+    (255, 0, 0),     # Red
+    (0, 255, 0),     # Green
+    (0, 0, 255),     # Blue
+    (255, 255, 0),   # Cyan
+    (255, 0, 255),   # Magenta
+    (0, 255, 255),   # Yellow
+    (128, 0, 0),     # Maroon
+    (0, 128, 0),     # Olive
+    (128, 128, 0),   # Teal
+    (0, 0, 128),     # Navy
+    (128, 0, 128),   # Purple
+    (0, 128, 128)    # Aqua
+]
+
 def detect_objects(image):
     # Convert the input image to a format YOLO can work with
     image = np.array(image)
     
     # Perform detection
-    results = model(image,line_width=1)[0]
+    results = model(image)[0]
     
     # Draw bounding boxes on the image
     for box in results.boxes.data.cpu().numpy():
         x1, y1, x2, y2, score, class_id = box
         x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
         
-        # Draw the bounding box
-        cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 2)
+        # Select color for the class id
+        color = colors[int(class_id) % len(colors)]
+        
+        # Draw the bounding box with a thinner line
+        cv2.rectangle(image, (x1, y1), (x2, y2), color, 1)  # Line width set to 1
         
         # Put the class name above the bounding box
         class_name = model.model.names[int(class_id)]
-        cv2.putText(image, f'{class_name} {score:.2f}', (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
+        cv2.putText(image, f'{class_name} {score:.2f}', (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 1)
     
     # Convert back to PIL image
     return Image.fromarray(image)