Update app.py

app.py

@@ -1,120 +1,83 @@
+import cv2
 import torch
 import numpy as np
 import gradio as gr
-import cv2
 import time
 import os
 from pathlib import Path
+import onnxruntime as ort
 
-# 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  
+# Set device for ONNX Runtime
+providers = ['CUDAExecutionProvider', 'CPUExecutionProvider'] if torch.cuda.is_available() else ['CPUExecutionProvider']
+session = ort.InferenceSession("models/yolov5n.onnx", providers=providers)
 
-if device.type == "cuda":
-    model.half()
-else:
-    torch.set_num_threads(os.cpu_count())
-
-model.eval()
+# Load model class names
+class_names = ["person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck", "boat", "traffic light"]  # Modify based on model
 
+# Generate random colors for classes
 np.random.seed(42)
-colors = np.random.uniform(0, 255, size=(len(model.names), 3))
+colors = np.random.uniform(0, 255, size=(len(class_names), 3))
 
-total_inference_time = 0
-inference_count = 0
+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
 
 def detect_objects(image):
-    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
-    
-    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)
+    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)
     
+    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
 
-example_images = ["spring_street_after.jpg", "pexels-hikaique-109919.jpg"]
-os.makedirs("examples", exist_ok=True)
+def real_time_detection():
+    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)
+    
+    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="Optimized YOLOv5 Object Detection") as demo:
+with gr.Blocks(title="YOLOv5 Real-Time Object Detection") as demo:
     gr.Markdown("""
-    # Optimized YOLOv5 Object Detection
-    Detects objects using YOLOv5 with enhanced visualization and FPS tracking.
+    # Real-Time Object Detection with YOLOv5
+    **Upload an image or run real-time detection**
     """)
     
     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():
+            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.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
-    )
+        with gr.Column():
+            output_image = gr.Image(label="Detection Results", type="numpy")
     
-    submit_button.click(fn=detect_objects, 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)
 
 demo.launch()