Update app.py

app.py

@@ -5,6 +5,7 @@ import cv2
 import time
 import os
 from pathlib import Path
+from PIL import Image
 
 # Create cache directory for models
 os.makedirs("models", exist_ok=True)
@@ -91,16 +92,37 @@ def process_video(video_path):
     
     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)")
+def process_image(image):
+    img = np.array(image)
+    results = model(img, size=640)
 
-    with gr.Row():
-        video_input = gr.Video(label="Upload Video")
-        process_button = gr.Button("Process Video")
+    detections = results.pred[0].cpu().numpy()
 
-    video_output = gr.Video(label="Processed Video")
+    for *xyxy, conf, cls in detections:
+        x1, y1, x2, y2 = map(int, xyxy)
+        class_id = int(cls)
+        color = colors[class_id].tolist()
+        cv2.rectangle(img, (x1, y1), (x2, y2), color, 3, lineType=cv2.LINE_AA)
+        label = f"{model.names[class_id]} {conf:.2f}"
+        cv2.putText(img, label, (x1, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (255, 255, 255), 2)
 
-    process_button.click(fn=process_video, inputs=video_input, outputs=video_output)
+    return Image.fromarray(img)
 
-demo.launch()
+# Gradio Interface
+with gr.Blocks(title="Real-Time YOLOv5 Video & Image Detection") as demo:
+    gr.Markdown("# Real-Time YOLOv5 Object Detection")
+
+    with gr.Tabs():
+        with gr.TabItem("Video Detection"):
+            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)
+
+        with gr.TabItem("Image Detection"):
+            image_input = gr.Image(type="pil", label="Upload Image")
+            image_output = gr.Image(label="Detected Objects")
+            image_input.change(fn=process_image, inputs=image_input, outputs=image_output)
+
+demo.launch()