Update app.py

app.py

@@ -3,58 +3,112 @@ import numpy as np
 import gradio as gr
 import cv2
 import time
-from PIL import Image
 
+# Check device availability
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print(f"Using device: {device}")
 
-# Load YOLOv5 model
+# Load YOLOv5x model (larger model for better accuracy)
 model = torch.hub.load("ultralytics/yolov5", "yolov5x", pretrained=True).to(device)
 
+# Optimization configurations
+model.conf = 0.3  # Confidence threshold of 0.3 as specified
+model.iou = 0.3   # NMS IoU threshold of 0.3 as specified
+model.classes = None  # Detect all 80+ COCO classes
+
+# Enable half-precision for GPU acceleration
 if device.type == "cuda":
     model.half()  # Use FP16 for performance boost
 
-# Print available object classes
-print(f"Model loaded with {len(model.names)} classes: {model.names}")
+# Set model to evaluation mode for inference
+model.eval()
 
-# Assign random colors to each class for bounding boxes
-colors = {i: [int(c) for c in np.random.randint(0, 255, 3)] for i in range(len(model.names))}
+# Assign fixed colors to each class for consistent visualization
+np.random.seed(42)  # For reproducible colors
+colors = np.random.uniform(0, 255, size=(len(model.names), 3))
 
 def detect_objects(image):
-    start_time = time.time()  # Start FPS measurement
-
-   
-    img_tensor = torch.from_numpy(image).permute(2, 0, 1).float().to(device) / 255.0
-    img_tensor = img_tensor.unsqueeze(0)  
-
+    """
+    Process input image for object detection using YOLOv5
+    
+    Args:
+        image: Input image as numpy array
+        
+    Returns:
+        output_image: Image with detection results visualized
+    """
+    start_time = time.time()
+    
+    # Convert image to RGB if it's in BGR format
+    if image.shape[2] == 3 and image[0,0,0] == image[0,0,2]:
+        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+    
+    # Create a copy for drawing results
+    output_image = image.copy()
+    
+    # Resize input to 640x640 for optimal processing speed
+    input_size = 640
+    
+    # Perform inference with no gradient calculation
     with torch.no_grad():
-        results = model(img_tensor)  
+        results = model(image, size=input_size)
     
-    detections = results.xyxy[0].cpu().numpy()  
-
-    img_cv = image.copy() 
-
-    for det in detections:
-        x1, y1, x2, y2, conf, cls = map(int, det[:6])
-        label = f"{model.names[cls]}: {conf:.2f}"
-
-        cv2.rectangle(img_cv, (x1, y1), (x2, y2), colors[cls], 2)
-        cv2.putText(img_cv, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, colors[cls], 2)
-
-    # FPS Calculation
-    end_time = time.time()
-    fps = 1 / (end_time - start_time)
-    print(f"FPS: {fps:.2f}")
-
-    return img_cv
+    # Extract detections from first (and only) image
+    detections = results.pred[0].cpu().numpy()
+    
+    # Draw each detection on the output image
+    for *xyxy, conf, cls in detections:
+        # Extract coordinates and convert to integers
+        x1, y1, x2, y2 = map(int, xyxy)
+        class_id = int(cls)
+        
+        # Get color for this class
+        color = colors[class_id].tolist()
+        
+        # Draw bounding box
+        cv2.rectangle(output_image, (x1, y1), (x2, y2), color, 2)
+        
+        # Create label with class name and confidence score
+        label = f"{model.names[class_id]} {conf:.2f}"
+        
+        # Calculate text size for background rectangle
+        (w, h), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)
+        
+        # Draw label background
+        cv2.rectangle(output_image, (x1, y1 - 20), (x1 + w, y1), color, -1)
+        
+        # Draw label text
+        cv2.putText(output_image, label, (x1, y1 - 5),
+                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)
+    
+    # Calculate and display FPS
+    fps = 1 / (time.time() - start_time)
+    
+    # Add FPS counter to the image
+    cv2.putText(output_image, f"FPS: {fps:.2f}", (10, 30),
+                cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
+    
+    print(f"Detection complete - FPS: {fps:.2f}")
+    
+    return output_image
 
-# 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="Object Detection with YOLOv5",
-    description="Optimized for 30+ FPS real-time object detection!",
+    title="Optimized Object Detection with YOLOv5x",
+    description="""
+    This system utilizes YOLOv5x to detect 80+ object types from the COCO dataset.
+    - Processing speed: Optimized for 30+ FPS at 640x640 resolution
+    - Confidence threshold: 0.3
+    - IoU threshold: 0.3
+    - Color-coded bounding boxes with confidence scores
+    """,
     allow_flagging="never",
+    examples=["spring_street_after.jpg", "pexels-hikaique-109919.jpg"],
 )
 
-iface.launch()
+# Launch the interface
+if __name__ == "__main__":
+    iface.launch()