Update app.py

app.py

@@ -1,35 +1,120 @@
-import cv2
 import torch
-from ultralytics import YOLO
+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)
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print(f"Using device: {device}")
+
+model_path = Path("models/yolov5n.pt")
+if model_path.exists():
+    print(f"Loading model from cache: {model_path}")
+    model = torch.hub.load("ultralytics/yolov5", "yolov5x", pretrained=True, source="local", path=str(model_path)).to(device)
+else:
+    print("Downloading YOLOv5n model and caching...")
+    model = torch.hub.load("ultralytics/yolov5", "yolov5x", pretrained=True).to(device)
+    torch.save(model.state_dict(), model_path)
+
+# Model configurations
+model.conf = 0.6 
+model.iou = 0.6   
+model.classes = None  
 
-# Load YOLOv5 model with GPU support if available
-device = "cuda" if torch.cuda.is_available() else "cpu"
-model = YOLO("yolov5s.pt").to(device)  # You can use yolov5m.pt for better accuracy
+if device.type == "cuda":
+    model.half()
+else:
+    torch.set_num_threads(os.cpu_count())
 
-# Initialize video capture (Webcam)
-cap = cv2.VideoCapture(0)
-cap.set(cv2.CAP_PROP_FPS, 30)  # Ensure 30+ FPS
-cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640)  # Set width
-cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)  # Set height
+model.eval()
 
-while cap.isOpened():
-    ret, frame = cap.read()
-    if not ret:
-        break
+np.random.seed(42)
+colors = np.random.uniform(0, 255, size=(len(model.names), 3))
 
-    # Perform object detection
-    results = model(frame)
+total_inference_time = 0
+inference_count = 0
 
-    # Plot the results (draw bounding boxes)
-    result_img = results[0].plot()
+def detect_objects(image):
+    global total_inference_time, inference_count
+    
+    if image is None:
+        return None
+    
+    start_time = time.time()
+    output_image = image.copy()
+    input_size = 640
+    
+    with torch.no_grad():
+        results = model(image, size=input_size)
+    
+    inference_time = time.time() - start_time
+    total_inference_time += inference_time
+    inference_count += 1
+    avg_inference_time = total_inference_time / inference_count
+    
+    detections = results.pred[0].cpu().numpy()
+    
+    for *xyxy, conf, cls in detections:
+        x1, y1, x2, y2 = map(int, xyxy)
+        class_id = int(cls)
+        color = colors[class_id].tolist()
+        
+        # Thicker bounding boxes
+        cv2.rectangle(output_image, (x1, y1), (x2, y2), color, 3, lineType=cv2.LINE_AA)
+        
+        label = f"{model.names[class_id]} {conf:.2f}"
+        font_scale, font_thickness = 0.9, 2  
+        (w, h), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, font_scale, font_thickness)
+        
+        cv2.rectangle(output_image, (x1, y1 - h - 10), (x1 + w + 10, y1), color, -1)
+        cv2.putText(output_image, label, (x1 + 5, y1 - 5),
+                    cv2.FONT_HERSHEY_SIMPLEX, font_scale, (255, 255, 255), font_thickness, lineType=cv2.LINE_AA)
+    
+    fps = 1 / inference_time
+    
+    # Stylish FPS display
+    overlay = output_image.copy()
+    cv2.rectangle(overlay, (10, 10), (300, 80), (0, 0, 0), -1)
+    output_image = cv2.addWeighted(overlay, 0.6, output_image, 0.4, 0)
+    cv2.putText(output_image, f"FPS: {fps:.2f}", (20, 40),
+                cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2, lineType=cv2.LINE_AA)
+    cv2.putText(output_image, f"Avg FPS: {1/avg_inference_time:.2f}", (20, 70),
+                cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2, lineType=cv2.LINE_AA)
+    
+    return output_image
 
-    # Display the output in a window
-    cv2.imshow("YOLOv5 Live Object Detection", result_img)
+example_images = ["spring_street_after.jpg", "pexels-hikaique-109919.jpg"]
+os.makedirs("examples", exist_ok=True)
 
-    # Break loop with 'q' key
-    if cv2.waitKey(1) & 0xFF == ord("q"):
-        break
+with gr.Blocks(title="Optimized YOLOv5 Object Detection") as demo:
+    gr.Markdown("""
+    # Optimized YOLOv5 Object Detection
+    Detects objects using YOLOv5 with enhanced visualization and FPS tracking.
+    """)
+    
+    with gr.Row():
+        with gr.Column(scale=1):
+            input_image = gr.Image(label="Input Image", type="numpy")
+            submit_button = gr.Button("Submit", variant="primary")
+            clear_button = gr.Button("Clear")
+        
+        with gr.Column(scale=1):
+            output_image = gr.Image(label="Detected Objects", type="numpy")
+    
+    gr.Examples(
+        examples=example_images,
+        inputs=input_image,
+        outputs=output_image,
+        fn=detect_objects,
+        cache_examples=True
+    )
+    
+    submit_button.click(fn=detect_objects, inputs=input_image, outputs=output_image)
+    clear_button.click(lambda: (None, None), None, [input_image, output_image])
 
-# Release resources
-cap.release()
-cv2.destroyAllWindows()
+demo.launch()