Update app.py

app.py

@@ -3,106 +3,64 @@ import numpy as np
 import gradio as gr
 import cv2
 import time
-import os
-from pathlib import Path
-
-# Create cache directory for models
-os.makedirs("models", exist_ok=True)
-
+# Check device availability
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-print(f"Using device: {device}")
-
-# Load YOLOv5 Model
-model_path = Path("models/yolov5n.pt")
-if model_path.exists():
-    print(f"Loading model from cache: {model_path}")
-    model = torch.hub.load("ultralytics/yolov5", "custom", path=str(model_path), source="local").to(device)
-else:
-    print("Downloading YOLOv5n model and caching...")
-    model = torch.hub.load("ultralytics/yolov5", "yolov5n", pretrained=True).to(device)
-    torch.save(model.state_dict(), model_path)
-
-# Configure model
-model.conf = 0.5
-model.iou = 0.5
-model.classes = None
+# Load smaller YOLOv5 model
+model = torch.hub.load("ultralytics/yolov5", "yolov5x", pretrained=True).to(device)
+# Optimization configurations
+model.conf = 0.3  # Confidence threshold
+model.iou = 0.3   # NMS IoU threshold
 if device.type == "cuda":
-    model.half()
-else:
-    torch.set_num_threads(os.cpu_count())
-model.eval()
-
-# Generate colors for bounding boxes
-np.random.seed(42)
+    model.half().to(device)  # Use FP16 for performance boost
+model.eval()  # Set model to evaluation mode
+# Assign fixed colors to each class for bounding boxes
 colors = np.random.uniform(0, 255, size=(len(model.names), 3))
-
 def detect_objects(image):
-    if image is None:
-        return None
-    
+    start_time = time.time()
+
+    # Convert BGR to RGB (if needed, Gradio might already provide RGB)
+    # image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+
+    # Perform inference
+    with torch.no_grad():
+        results = model(image, size=640)  # Fixed inference size
+
+    # Process results directly on numpy array
     output_image = image.copy()
-    results = model(image, size=640)
+
+    # Extract detections
     detections = results.pred[0].cpu().numpy()
-    
+
     for *xyxy, conf, cls in detections:
         x1, y1, x2, y2 = map(int, xyxy)
         class_id = int(cls)
+
+        # Draw bounding box
         color = colors[class_id].tolist()
-        cv2.rectangle(output_image, (x1, y1), (x2, y2), color, 3, lineType=cv2.LINE_AA)
-        label = f"{model.names[class_id]} {conf:.2f}"
-        cv2.putText(output_image, label, (x1, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (255, 255, 255), 2)
-    
-    return output_image
+        cv2.rectangle(output_image, (x1, y1), (x2, y2), color, 2)
 
-def process_video(video_path):
-    cap = cv2.VideoCapture(video_path)
-    if not cap.isOpened():
-        return "Error: Could not open video file."
-    
-    frame_width = int(cap.get(3))
-    frame_height = int(cap.get(4))
-    fps = cap.get(cv2.CAP_PROP_FPS)
-    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
-    output_path = "output_video.mp4"
-    out = cv2.VideoWriter(output_path, fourcc, fps, (frame_width, frame_height))
-    
-    while cap.isOpened():
-        ret, frame = cap.read()
-        if not ret:
-            break
-        
-        img = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
-        results = model(img, size=640)
-        detections = results.pred[0].cpu().numpy()
+        # Create label
+        label = f"{model.names[class_id]} {conf:.2f}"
 
-        for *xyxy, conf, cls in detections:
-            x1, y1, x2, y2 = map(int, xyxy)
-            class_id = int(cls)
-            color = colors[class_id].tolist()
-            cv2.rectangle(frame, (x1, y1), (x2, y2), color, 3, lineType=cv2.LINE_AA)
-            label = f"{model.names[class_id]} {conf:.2f}"
-            cv2.putText(frame, label, (x1, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (255, 255, 255), 2)
-        
-        out.write(frame)
-    
-    cap.release()
-    out.release()
-    return output_path
+        # Draw label background
+        (w, h), * = cv2.getTextSize(label, cv2.FONT*HERSHEY_SIMPLEX, 0.5, 1)
+        cv2.rectangle(output_image, (x1, y1 - 20), (x1 + w, y1), color, -1)
 
-# Gradio Interface
-with gr.Blocks(title="YOLOv5 Object Detection") as demo:
-    gr.Markdown("# YOLOv5 Object Detection (Image & Video)")
-    
-    with gr.Tab("Image Detection"):
-        img_input = gr.Image(label="Upload Image", type="numpy")
-        img_output = gr.Image(label="Detected Objects", type="numpy")
-        img_submit = gr.Button("Detect Objects")
-        img_submit.click(fn=detect_objects, inputs=img_input, outputs=img_output)
-    
-    with gr.Tab("Video Detection"):
-        vid_input = gr.Video(label="Upload Video")
-        vid_output = gr.Video(label="Processed Video")
-        vid_submit = gr.Button("Process Video")
-        vid_submit.click(fn=process_video, inputs=vid_input, outputs=vid_output)
-    
-    demo.launch()
+        # Draw label text
+        cv2.putText(output_image, label, (x1, y1 - 5),
+                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)
+    # Calculate FPS
+    fps = 1 / (time.time() - start_time)
+    print(f"FPS: {fps:.2f}")
+    return output_image
+# 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="Optimized Object Detection with YOLOv5",
+    description="Faster detection using YOLOv5s with FP16 and optimized processing",
+    allow_flagging="never",
+    examples=["spring_street_after.jpg", "pexels-hikaique-109919.jpg"],
+)
+iface.launch()