Update app.py

app.py

@@ -3,31 +3,60 @@ from PIL import Image
 import gradio as gr
 from huggingface_hub import snapshot_download
 import os
+import tempfile
+import cv2
 
 def load_model(repo_id):
     download_dir = snapshot_download(repo_id)
-    print(download_dir)
     path = os.path.join(download_dir, "best_int8_openvino_model")
-    print(path)
-    detection_model = YOLO(path, task='detect')
-    return detection_model
+    return YOLO(path, task='detect')
 
-def predict(pilimg, conf_threshold, iou_threshold):
-    source = pilimg
-    result = detection_model.predict(source, conf=conf_threshold, iou=iou_threshold)
+def predict_image(pilimg, conf_threshold, iou_threshold):
+    result = detection_model.predict(pilimg, conf=conf_threshold, iou=iou_threshold)
     img_bgr = result[0].plot()
     out_pilimg = Image.fromarray(img_bgr[..., ::-1])  # Convert BGR to RGB for PIL
     return out_pilimg
 
+def predict_video(video_path, conf_threshold, iou_threshold):
+    cap = cv2.VideoCapture(video_path)
+    width  = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+    fps    = cap.get(cv2.CAP_PROP_FPS)
+
+    out_path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
+    out_writer = cv2.VideoWriter(out_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (width, height))
+
+    while cap.isOpened():
+        ret, frame = cap.read()
+        if not ret:
+            break
+        result = detection_model.predict(frame, conf=conf_threshold, iou=iou_threshold)
+        annotated = result[0].plot()
+        out_writer.write(annotated)
+
+    cap.release()
+    out_writer.release()
+    return out_path
+
+def predict(file, conf_threshold, iou_threshold):
+    ext = os.path.splitext(file.name)[1].lower()
+    if ext in ['.jpg', '.jpeg', '.png']:
+        img = Image.open(file).convert("RGB")
+        return predict_image(img, conf_threshold, iou_threshold)
+    elif ext in ['.mp4', '.mov', '.avi']:
+        return predict_video(file.name, conf_threshold, iou_threshold)
+    else:
+        return "Unsupported file type. Please upload an image or video."
+
 REPO_ID = "sensura/belisha-beacon-zebra-crossing-yoloV8"
 detection_model = load_model(REPO_ID)
 
 gr.Interface(
     fn=predict,
     inputs=[
-        gr.Image(type="pil", label="Upload Image"),
-        gr.Slider(minimum=0.1, maximum=1.0, value=0.5, step=0.05, label="Confidence Threshold"),
-        gr.Slider(minimum=0.1, maximum=1.0, value=0.6, step=0.05, label="IoU Threshold")
+        gr.File(label="Upload Image or Video"),
+        gr.Slider(0.1, 1.0, 0.5, step=0.05, label="Confidence Threshold"),
+        gr.Slider(0.1, 1.0, 0.6, step=0.05, label="IoU Threshold")
     ],
-    outputs=gr.Image(type="pil", label="Detection Output")
+    outputs=gr.outputs.Image(type="pil", label="Detected Image or Video") | gr.outputs.Video(label="Detected Video")
 ).launch(share=True)