Upload 2 files

app.py

@@ -3,7 +3,7 @@ from PIL import Image
 import gradio as gr
 from huggingface_hub import snapshot_download
 import os
-
+import cv2
 
 def load_model(repo_id):
     download_dir = snapshot_download(repo_id)
@@ -14,22 +14,32 @@ def load_model(repo_id):
     return detection_model
 
 
-def predict(pilimg):
-
-    source = pilimg
-    # x = np.asarray(pilimg)
-    # print(x.shape)
-    result = detection_model.predict(source, conf=0.5, iou=0.6)
-    img_bgr = result[0].plot()
-    out_pilimg = Image.fromarray(img_bgr[..., ::-1])  # RGB-order PIL image
-    
-    return out_pilimg
+def predict(input_data):
+    if isinstance(input_data, Image.Image): 
+        source = input_data
+        result = detection_model.predict(source, conf=0.5, iou=0.6)
+        img_bgr = result[0].plot()
+        out_pilimg = Image.fromarray(img_bgr[..., ::-1]) 
+        return out_pilimg
+    elif isinstance(input_data, str) and input_data.endswith('.mp4'): 
+        cap = cv2.VideoCapture(input_data)
+        ret, frame = cap.read() 
+        if ret:
+            result = detection_model.predict(frame, conf=0.5, iou=0.6)
+            img_bgr = result[0].plot()
+            out_pilimg = Image.fromarray(img_bgr[..., ::-1])
+            cap.release()
+            return out_pilimg
+        cap.release()
+    return None 
 
 
 REPO_ID = "dexpyw/model"
 detection_model = load_model(REPO_ID)
 
-gr.Interface(fn=predict,
-             inputs=gr.Image(type="pil"),
-             outputs=gr.Image(type="pil")
-             ).launch(share=True)
+gr.Interface(
+    fn=predict,
+    inputs=[gr.Image(type="pil", label="Upload Image"), gr.Video(type="file", label="Upload Video")],
+    outputs=gr.Image(type="pil", label="Predicted Output"),
+    live=True
+).launch(share=True)

requirements.txt

@@ -1,2 +1,3 @@
 ultralytics
-huggingface_hub
+huggingface_hub
+opencv-python