Update app.py

app.py

@@ -12,109 +12,91 @@ os.makedirs("models", exist_ok=True)
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 print(f"Using device: {device}")
 
+# Use smaller YOLOv5n model instead of x-large
 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)
+if not model_path.exists():
+    print("Downloading and caching YOLOv5n...")
+    torch.hub.download_url_to_file("https://github.com/ultralytics/yolov5/releases/download/v6.2/yolov5n.pt", "models/yolov5n.pt")
+
+# Optimized model loading
+model = torch.hub.load("ultralytics/yolov5", "custom", path=str(model_path), autoshape=False).to(device)
 
-# Model configurations
-model.conf = 0.6 
-model.iou = 0.6   
-model.classes = None  
+# Model optimizations
+model.conf = 0.5  # Slightly lower confidence threshold
+model.iou = 0.45  # Lower IoU threshold for faster NMS
+model.classes = None  # Detect all classes
 
+# Precision optimizations
 if device.type == "cuda":
-    model.half()
+    model.half()  # FP16 inference
+    torch.backends.cudnn.benchmark = True  # Better CUDA performance
 else:
-    torch.set_num_threads(os.cpu_count())
+    model.float()
+    torch.set_num_threads(2)  # Limit CPU threads for better resource management
 
 model.eval()
 
-np.random.seed(42)
-colors = np.random.uniform(0, 255, size=(len(model.names), 3))
+# Simplified color generation
+colors = np.random.rand(len(model.names), 3) * 255
 
-total_inference_time = 0
-inference_count = 0
+total_time = 0
+frame_count = 0
 
 def detect_objects(image):
-    global total_inference_time, inference_count
+    global total_time, frame_count
     
     if image is None:
         return None
     
-    start_time = time.time()
-    output_image = image.copy()
-    input_size = 640
+    start = time.perf_counter()
     
-    with torch.no_grad():
-        results = model(image, size=input_size)
+    # Reduce input size and use optimized preprocessing
+    input_size = 320  # Reduced from 640
+    im = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+    im = cv2.resize(im, (input_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
+    with torch.no_grad():
+        if device.type == "cuda":
+            im = torch.from_numpy(im).to(device).half().permute(2, 0, 1).unsqueeze(0) / 255
+        else:
+            im = torch.from_numpy(im).to(device).float().permute(2, 0, 1).unsqueeze(0) / 255
+        
+        pred = model(im, augment=False)[0]
     
-    detections = results.pred[0].cpu().numpy()
+    # Faster post-processing
+    pred = pred.float() if device.type == "cpu" else pred.half()
+    pred = non_max_suppression(pred, model.conf, model.iou, agnostic=False)[0]
     
-    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)
+    # Optimized visualization
+    output = image.copy()
+    if pred is not None and len(pred):
+        pred[:, :4] = scale_coords(im.shape[2:], pred[:, :4], output.shape).round()
+        for *xyxy, conf, cls in pred:
+            x1, y1, x2, y2 = map(int, xyxy)
+            cv2.rectangle(output, (x1, y1), (x2, y2), colors[int(cls)].tolist(), 2)
     
-    fps = 1 / inference_time
+    # FPS calculation
+    dt = time.perf_counter() - start
+    total_time += dt
+    frame_count += 1
+    fps = 1 / dt
+    avg_fps = frame_count / total_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)
+    # Simplified FPS display
+    cv2.putText(output, f"FPS: {fps:.1f}", (10, 30), 
+                cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 0), 2)
     
-    return output_image
+    return output
 
-example_images = ["spring_street_after.jpg", "pexels-hikaique-109919.jpg"]
-os.makedirs("examples", exist_ok=True)
+# Use smaller example images
+example_images = ["pexels-hikaique-109919.jpg", "spring_street_after.jpg"]
 
-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.Blocks(title="Optimized YOLOv5") as demo:
+    gr.Markdown("# Real-Time YOLOv5 Object 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(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])
+        input_img = gr.Image(label="Input", source="webcam" if os.getenv('SPACE_ID') else None)
+        output_img = gr.Image(label="Output")
+    gr.Examples(examples=example_images, inputs=input_img, outputs=output_img, fn=detect_objects)
+    input_img.change(fn=detect_objects, inputs=input_img, outputs=output_img)
 
-demo.launch()
+demo.launch()