Update app.py

app.py

@@ -4,6 +4,7 @@ import yolov5
 from PIL import Image
 from huggingface_hub import hf_hub_download
 import os 
+import cv2 
 
 
 app_title = "Detect san pham VSK"
@@ -40,7 +41,55 @@ def predict(image, threshold=0.25, model_id=None):
     output_image = Image.fromarray(numpy_image)
     return output_image, count_result
 
-gr.Interface(
+def show_pred_vid(
+    video_path: str = None,
+    model_path: str = None,
+    image_size: 640,
+    conf_threshold: float = 0.25,
+    iou_threshold: float = 0.45,
+):
+    cap = cv2.VideoCapture(video_path)
+
+    while cap.isOpened():
+        success, frame = cap.read()
+
+        if success:
+            model = YOLO(model_path)
+            model.overrides['conf'] = conf_threshold
+            model.overrides['iou'] = iou_threshold
+            model.overrides['agnostic_nms'] = False
+            model.overrides['max_det'] = 1000
+            results = model.predict(frame)
+            annotated_frame = results[0].plot()
+
+            if cv2.waitKey(1) & 0xFF == ord("q"):
+                break
+        else:
+            break
+    cap.release()
+    cv2.destroyAllWindows()
+
+inputs_vid = [
+    gr.components.Video(type="filepath", label="Input Video"),
+    gr.inputs.Dropdown(["linhcuem/yolov5_chamdiem_ver13"], default="linhcuem/yolov5_chamdiem_ver13", label="Model"),
+    gr.inputs.Slider(minimum=320, maximum=1280, default=640, step=32, label= "Image Size"),
+    gr.inputs.Slider(minimum=0.0, maximum=1.0, default=0.25, step=0.05, label="Confidence Threshold"),
+    gr.inputs.Slider(minimum=0.0, maximum=1.0, default=0.45, step=0.05, label="IOU Threshold"),
+]
+
+outputs_vid = gr.outputs.Image(type="filepath", label="Output Video")
+interface_vid = gr.Interface(
+    fn=show_pred_vid,
+    inputs = inputs_vid,
+    outputs = outputs_vid,
+    title = app_title,
+    description=description,
+    cache_examples=False,
+    theme='huggingface'
+)
+
+    
+interface_image = gr.Interface(
     title=app_title,
     description="DO ANH DAT",
     fn=predict,
@@ -52,4 +101,9 @@ gr.Interface(
     outputs=[gr.Image(type="pil"),gr.Textbox(show_label=False)],
     examples=examples,
     cache_examples=True if examples else Fale,
-).launch(enable_queue=True)
+)
+
+gr.TabbedInterface(
+    [interface_image, interface_vid],
+    tab_names=['Image inferece', 'Video inference']
+).launch()