Update app.py

app.py

@@ -1,83 +1,274 @@
-import cv2
 import torch
 import numpy as np
 import gradio as gr
+import cv2
 import time
 import os
+import threading
+from queue import Queue
 from pathlib import Path
-import onnxruntime as ort
 
-# Set device for ONNX Runtime
-providers = ['CUDAExecutionProvider', 'CPUExecutionProvider'] if torch.cuda.is_available() else ['CPUExecutionProvider']
-session = ort.InferenceSession("models/yolov5n.onnx", providers=providers)
+# 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/yolov5x.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 YOLOv5x model and caching...")
+    model = torch.hub.load("ultralytics/yolov5", "yolov5x", pretrained=True).to(device)
+    torch.save(model.state_dict(), model_path)
 
-# Load model class names
-class_names = ["person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck", "boat", "traffic light"]  # Modify based on model
+# Model configurations for better performance
+model.conf = 0.5  # Slightly lower confidence threshold for real-time
+model.iou = 0.45  # Slightly lower IOU threshold for real-time
+model.classes = None  # Detect all classes
+model.max_det = 20  # Limit detections for speed
 
-# Generate random colors for classes
+if device.type == "cuda":
+    model.half()  # Half precision for CUDA
+else:
+    torch.set_num_threads(os.cpu_count())
+
+model.eval()
+
+# Precompute colors for bounding boxes
 np.random.seed(42)
-colors = np.random.uniform(0, 255, size=(len(class_names), 3))
+colors = np.random.uniform(0, 255, size=(len(model.names), 3))
+
+# Performance tracking
+total_inference_time = 0
+inference_count = 0
+fps_queue = Queue(maxsize=30)  # Store last 30 FPS values for smoothing
 
-def preprocess(image):
-    image = cv2.resize(image, (640, 640))
-    image = image.transpose((2, 0, 1)) / 255.0  # Normalize
-    image = np.expand_dims(image, axis=0).astype(np.float32)
-    return image
+# Threading variables
+processing_lock = threading.Lock()
+stop_event = threading.Event()
+frame_queue = Queue(maxsize=2)  # Small queue to avoid lag
+result_queue = Queue(maxsize=2)
 
 def detect_objects(image):
+    """Process a single image for object detection"""
+    global total_inference_time, inference_count
+    
+    if image is None:
+        return None
+    
     start_time = time.time()
-    image_input = preprocess(image)
-    outputs = session.run(None, {session.get_inputs()[0].name: image_input})
-    detections = outputs[0][0]
     output_image = image.copy()
+    input_size = 640
+    
+    # Optimize input for inference
+    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()
+    
+    # Draw detections
+    for *xyxy, conf, cls in detections:
+        x1, y1, x2, y2 = map(int, xyxy)
+        class_id = int(cls)
+        color = colors[class_id].tolist()
+        
+        # Bounding box
+        cv2.rectangle(output_image, (x1, y1), (x2, y2), color, 3, lineType=cv2.LINE_AA)
+        
+        # Label with class name and confidence
+        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
     
-    for det in detections:
-        x1, y1, x2, y2, conf, cls = map(int, det[:6])
-        if conf > 0.6:  # Confidence threshold
-            color = colors[cls].tolist()
-            cv2.rectangle(output_image, (x1, y1), (x2, y2), color, 2)
-            label = f"{class_names[cls]} {conf:.2f}"
-            cv2.putText(output_image, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 2)
+    # 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)
     
-    fps = 1 / (time.time() - start_time)
-    cv2.putText(output_image, f"FPS: {fps:.2f}", (20, 40), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
     return output_image
 
-def real_time_detection():
+def process_frame_thread():
+    """Background thread for processing frames"""
+    while not stop_event.is_set():
+        if not frame_queue.empty():
+            frame = frame_queue.get()
+            
+            # Skip if there's a processing lock (from image upload)
+            if processing_lock.locked():
+                result_queue.put(frame)  # Return unprocessed frame
+                continue
+                
+            # Process the frame
+            with torch.no_grad():  # Ensure no gradients for inference
+                input_size = 384  # Smaller size for real-time processing
+                results = model(frame, size=input_size)
+            
+            # Calculate FPS
+            inference_time = time.time() - frame.get('timestamp', time.time())
+            current_fps = 1 / inference_time if inference_time > 0 else 30
+            
+            # Update rolling FPS average
+            fps_queue.put(current_fps)
+            avg_fps = sum(list(fps_queue.queue)) / fps_queue.qsize()
+            
+            # Draw detections
+            output = frame['image'].copy()
+            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()
+                
+                # Draw rectangle and label
+                cv2.rectangle(output, (x1, y1), (x2, y2), color, 2, lineType=cv2.LINE_AA)
+                
+                label = f"{model.names[class_id]} {conf:.2f}"
+                font_scale, font_thickness = 0.6, 1  # Smaller for real-time
+                (w, h), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, font_scale, font_thickness)
+                
+                cv2.rectangle(output, (x1, y1 - h - 5), (x1 + w + 5, y1), color, -1)
+                cv2.putText(output, label, (x1 + 3, y1 - 3),
+                            cv2.FONT_HERSHEY_SIMPLEX, font_scale, (255, 255, 255), font_thickness, lineType=cv2.LINE_AA)
+            
+            # Add FPS counter
+            cv2.rectangle(output, (10, 10), (210, 80), (0, 0, 0), -1)
+            cv2.putText(output, f"FPS: {current_fps:.1f}", (20, 40),
+                        cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 0), 2, lineType=cv2.LINE_AA)
+            cv2.putText(output, f"Avg FPS: {avg_fps:.1f}", (20, 70),
+                        cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 0), 2, lineType=cv2.LINE_AA)
+            
+            # Put the processed frame in the result queue
+            result_queue.put({'image': output, 'fps': current_fps})
+        else:
+            time.sleep(0.001)  # Small sleep to prevent CPU spinning
+
+def webcam_feed():
+    """Generator function for webcam feed"""
+    # Start the processing thread if not already running
+    if not any(thread.name == "frame_processor" for thread in threading.enumerate()):
+        stop_event.clear()
+        processor = threading.Thread(target=process_frame_thread, name="frame_processor", daemon=True)
+        processor.start()
+    
+    # Open webcam
     cap = cv2.VideoCapture(0)
-    cap.set(cv2.CAP_PROP_FRAME_WIDTH, 1280)
-    cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 720)
-    cap.set(cv2.CAP_PROP_FPS, 60)
+    cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
+    cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)
     
-    while cap.isOpened():
-        start_time = time.time()
-        ret, frame = cap.read()
-        if not ret:
-            break
-        output_frame = detect_objects(frame)
-        cv2.imshow("Real-Time Object Detection", output_frame)
-        if cv2.waitKey(1) & 0xFF == ord('q'):
-            break
-        print(f"FPS: {1 / (time.time() - start_time):.2f}")
-    cap.release()
-    cv2.destroyAllWindows()
-
-with gr.Blocks(title="YOLOv5 Real-Time Object Detection") as demo:
+    try:
+        while True:
+            success, frame = cap.read()
+            if not success:
+                break
+                
+            # Put frame in queue for processing
+            if not frame_queue.full():
+                frame_queue.put({'image': frame, 'timestamp': time.time()})
+            
+            # Get processed frame from result queue
+            if not result_queue.empty():
+                result = result_queue.get()
+                yield result['image']
+            else:
+                # If no processed frame is available, yield the raw frame
+                yield frame
+                
+            # Control frame rate to not overwhelm the system
+            time.sleep(0.01)
+    finally:
+        cap.release()
+
+def process_uploaded_image(image):
+    """Process an uploaded image (this will be separate from real-time)"""
+    with processing_lock:  # Acquire lock to pause real-time processing
+        return detect_objects(image)
+
+# Setup Gradio interface
+example_images = ["spring_street_after.jpg", "pexels-hikaique-109919.jpg"]
+os.makedirs("examples", exist_ok=True)
+
+with gr.Blocks(title="YOLOv5 Object Detection - Real-time & Image Upload") as demo:
     gr.Markdown("""
-    # Real-Time Object Detection with YOLOv5
-    **Upload an image or run real-time detection**
+    # YOLOv5 Object Detection
+    ## Real-time webcam detection and image upload processing
     """)
     
-    with gr.Row():
-        with gr.Column():
-            input_image = gr.Image(label="Upload Image", type="numpy")
-            detect_button = gr.Button("Detect Objects")
-            start_rt_button = gr.Button("Start Real-Time Detection")
+    with gr.Tabs():
+        with gr.TabItem("Real-time Detection"):
+            gr.Markdown("""
+            ### Real-time Object Detection
+            Using your webcam for continuous object detection at 30+ FPS.
+            """)
+            webcam_output = gr.Image(label="Real-time Detection", type="numpy")
         
-        with gr.Column():
-            output_image = gr.Image(label="Detection Results", type="numpy")
+        with gr.TabItem("Image Upload"):
+            gr.Markdown("""
+            ### Image Upload Detection
+            Upload an image to detect objects.
+            """)
+            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=process_uploaded_image,
+                cache_examples=True
+            )
+    
+    # Set up event handlers
+    submit_button.click(fn=process_uploaded_image, inputs=input_image, outputs=output_image)
+    clear_button.click(lambda: (None, None), None, [input_image, output_image])
     
-    detect_button.click(detect_objects, inputs=input_image, outputs=output_image)
-    start_rt_button.click(lambda: real_time_detection(), None, None)
+    # Connect webcam feed
+    demo.load(lambda: None, None, webcam_output, _js="""
+        () => {
+            // Keep the webcam tab refreshing at high frequency
+            setInterval(() => {
+                if (document.querySelector('.tabitem:first-child').style.display !== 'none') {
+                    const webcamImg = document.querySelector('.tabitem:first-child img');
+                    if (webcamImg) {
+                        const src = webcamImg.src;
+                        webcamImg.src = src.includes('?') ? src.split('?')[0] + '?t=' + Date.now() : src + '?t=' + Date.now();
+                    }
+                }
+            }, 33); // ~30 FPS refresh rate
+            return [];
+        }
+    """)
+    
+    # Start webcam feed
+    webcam_output.update(webcam_feed)
+
+# Cleanup function to stop threads when app closes
+def cleanup():
+    stop_event.set()
+    print("Cleaning up threads...")
 
-demo.launch()
+demo.close = cleanup
+demo.launch()