Update app.py

app.py

@@ -2,63 +2,59 @@ import torch
 import cv2
 import numpy as np
 import gradio as gr
+from PIL import Image
 import random
 
 # Load YOLOv5 model
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-model = torch.hub.load('ultralytics/yolov5', 'yolov5x', pretrained=True).to(device)
-model.eval()
+model = torch.hub.load('ultralytics/yolov5', 'yolov5x', pretrained=True).to(device) 
 
-# Use half-precision if CUDA is available
-if device.type == 'cuda':
-    model.half()
-
-# Get class names
+# Get class names from the model
 CLASS_NAMES = model.names  
 
-# Assign random colors for each class
+# Generate consistent colors for each class
+random.seed(42)  # Fix the seed for consistent colors
 CLASS_COLORS = {cls: (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255)) for cls in CLASS_NAMES}
 
-def detect_objects(image):
-    """Detect objects in an image using YOLOv5 with optimized inference speed."""
-    image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)  # Convert to BGR for OpenCV
-    img_resized = cv2.resize(image, (640, 640))  # Resize for faster processing
-    img_tensor = torch.from_numpy(img_resized).to(device).float() / 255.0  # Normalize
-    img_tensor = img_tensor.permute(2, 0, 1).unsqueeze(0)  # Convert to batch format
-
-    if device.type == 'cuda':
-        img_tensor = img_tensor.half()  # Use half precision for speed
+def preprocess_image(image):
+    """Convert numpy image to PIL format for YOLOv5 processing."""
+    image = Image.fromarray(image)
+    image = image.convert("RGB")
+    return image
 
-    # Run model inference
-    with torch.no_grad():
-        results = model(img_tensor)
-
-    detections = results.xyxy[0].cpu().numpy()
+def detect_objects(image):
+    """Detect objects in the image and draw bounding boxes with consistent colors."""
+    image = preprocess_image(image)
+    results = model([image])  # YOLOv5 inference
+    image = np.array(image)  # Convert PIL image back to numpy for OpenCV
 
-    for x1, y1, x2, y2, conf, cls in detections:
-        x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
-        class_name = CLASS_NAMES[int(cls)]
-        confidence = conf * 100
+    for *box, conf, cls in results.xyxy[0]:  
+        x1, y1, x2, y2 = map(int, box)
+        class_name = CLASS_NAMES[int(cls)]  
+        confidence = conf.item() * 100  
 
-        color = CLASS_COLORS[class_name]
+        color = CLASS_COLORS[class_name]  # Use pre-generated consistent color
 
         # Draw bounding box
-        cv2.rectangle(image, (x1, y1), (x2, y2), color, 3)
+        cv2.rectangle(image, (x1, y1), (x2, y2), color, 4)  
 
-        # Label
+        # Display class label with confidence score
         label = f"{class_name} ({confidence:.1f}%)"
-        cv2.putText(image, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.8, color, 2)
+        cv2.putText(image, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 
+                    1, color, 3, cv2.LINE_AA)  
+
+    return image
 
-    return cv2.cvtColor(image, cv2.COLOR_BGR2RGB)  # Convert back to RGB for Gradio
 
-# Gradio Interface
+# Create Gradio Interface
 iface = gr.Interface(
     fn=detect_objects,
     inputs=gr.Image(type="numpy", label="Upload Image"),
     outputs=gr.Image(type="numpy", label="Detected Objects"),
-    title="Fast Object Detection with YOLOv5",
-    description="Use webcam or upload an image for object detection results.",
-    allow_flagging="never"
+    title="Object Detection with YOLOv5",
+    description="Use webcam or upload an image to detect objects.",
+    allow_flagging="never",
+    examples=["spring_street_after.jpg", "pexels-hikaique-109919.jpg"]
 )
 
 # Launch the app