Update app.py

app.py

@@ -9,7 +9,15 @@ model_path = 'best.pt'  # Replace with the path to your trained .pt file
 model = YOLO(model_path)
 
 # Function to perform inference on an image
-#@spaces.GPU
+colors = {
+    0: (255, 0, 0),    # Red for category 0
+    1: (0, 255, 0),    # Green for category 1
+    2: (0, 0, 255),    # Blue for category 2
+    3: (255, 255, 0),  # Cyan for category 3
+    4: (255, 0, 255)   # Magenta for category 4
+}
+
+# Function to perform inference on an image
 def infer_image(image):
     # Convert the image from BGR to RGB
     image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
@@ -24,11 +32,14 @@ def infer_image(image):
             cls = int(box.cls[0])
             conf = float(box.conf[0])
             
+            # Get the color for the current class
+            color = colors.get(cls, (0, 255, 0))  # Default to green if class not found
+            
             # Draw bounding box
-            cv2.rectangle(image, (int(x1), int(y1)), (int(x2), int(y2)), (0, 255, 0), 2)
+            cv2.rectangle(image, (int(x1), int(y1)), (int(x2), int(y2)), color, 2)
             # Draw label
             label = f'{model.names[cls]} {conf:.2f}'
-            cv2.putText(image, label, (int(x1), int(y1) - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
+            cv2.putText(image, label, (int(x1), int(y1) - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
     
     return image