Update app.py

app.py

@@ -12,7 +12,7 @@ os.makedirs("models", exist_ok=True)
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 print(f"Using device: {device}")
 
-# Load YOLOv5 Nano model
+# Load YOLOv5 Model
 model_path = Path("models/yolov5n.pt")
 if model_path.exists():
     print(f"Loading model from cache: {model_path}")
@@ -22,54 +22,57 @@ else:
     model = torch.hub.load("ultralytics/yolov5", "yolov5n", pretrained=True).to(device)
     torch.save(model.state_dict(), model_path)
 
-# Optimize model for speed
-model.conf = 0.3  # Lower confidence threshold
-model.iou = 0.3   # Non-Maximum Suppression IoU threshold
-model.classes = None  # Detect all classes
-
+# Configure model
+model.conf = 0.5
+model.iou = 0.5
+model.classes = None
 if device.type == "cuda":
-    model.half()  # Use FP16 for faster inference
+    model.half()
 else:
     torch.set_num_threads(os.cpu_count())
-
 model.eval()
 
-# Pre-generate colors for bounding boxes
+# Generate colors for bounding boxes
 np.random.seed(42)
 colors = np.random.uniform(0, 255, size=(len(model.names), 3))
 
+def detect_objects(image):
+    if image is None:
+        return None
+    
+    output_image = image.copy()
+    results = model(image, size=640)
+    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()
+        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
+
 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))
-
-    total_frames = 0
-    total_time = 0
-
+    
     while cap.isOpened():
         ret, frame = cap.read()
         if not ret:
-            break  # Break if no more frames
-        
-        start_time = time.time()
+            break
         
-        # Convert frame for YOLOv5
         img = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
         results = model(img, size=640)
-        
-        inference_time = time.time() - start_time
-        total_time += inference_time
-        total_frames += 1
-        
         detections = results.pred[0].cpu().numpy()
 
         for *xyxy, conf, cls in detections:
@@ -79,28 +82,27 @@ def process_video(video_path):
             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)
-
-        # Calculate FPS
-        avg_fps = total_frames / total_time if total_time > 0 else 0
-        cv2.putText(frame, f"FPS: {avg_fps:.2f}", (20, 40), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
-
+        
         out.write(frame)
-
+    
     cap.release()
     out.release()
-    
     return output_path
 
 # Gradio Interface
-with gr.Blocks(title="Real-Time YOLOv5 Video Detection") as demo:
-    gr.Markdown("# Real-Time YOLOv5 Video Detection (30+ FPS)")
-
-    with gr.Row():
-        video_input = gr.Video(label="Upload Video")
-        process_button = gr.Button("Process Video")
-
-    video_output = gr.Video(label="Processed Video")
-
-    process_button.click(fn=process_video, inputs=video_input, outputs=video_output)
-
-demo.launch()
+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()