Update app.py

app.py

@@ -3,13 +3,15 @@ import cv2
 import numpy as np
 import gradio as gr
 from PIL import Image
+import random
 
-# Load YOLOv5 model from Ultralytics' official repo
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 model = torch.hub.load('ultralytics/yolov5', 'yolov5x', pretrained=True).to(device)  # Load YOLOv5x model
 
-# Load COCO class labels
-CLASS_NAMES = model.names  # YOLOv5's built-in class names
+CLASS_NAMES = model.names  
+
+random.seed(42)
+CLASS_COLORS = {cls: (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255)) for cls in CLASS_NAMES}
 
 def preprocess_image(image):
     image = Image.fromarray(image)
@@ -19,28 +21,27 @@ def preprocess_image(image):
 def detect_objects(image):
     image = preprocess_image(image)
 
-    # Run inference using YOLOv5
     results = model(image)
 
-    # Convert results to bounding box format
     image = np.array(image)
     
-    for *box, conf, cls in results.xyxy[0]:  # YOLOv5 format: [x1, y1, x2, y2, conf, class]
+    for *box, conf, cls in results.xyxy[0]:  
         x1, y1, x2, y2 = map(int, box)
-        class_name = CLASS_NAMES[int(cls)]  # Get class name
-        confidence = conf.item() * 100  # Convert confidence to percentage
+        class_name = CLASS_NAMES[int(cls)]  
+        confidence = conf.item() * 100  
+
+        color = CLASS_COLORS[class_name]
 
-        # Draw a bolder bounding box
-        cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 4)  # Increased thickness
+        cv2.rectangle(image, (x1, y1), (x2, y2), color, 4)  
 
-        # Larger text for class label
+        
         label = f"{class_name} ({confidence:.1f}%)"
         cv2.putText(image, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 
-                    1, (0, 255, 0), 3, cv2.LINE_AA)  # Larger text
+                    1, color, 3, cv2.LINE_AA)  # Larger text
 
     return image
 
-# Gradio UI with Submit button
+
 iface = gr.Interface(
     fn=detect_objects,
     inputs=gr.Image(type="numpy", label="Upload Image"),