Create app-vid.py

app-vid.py

@@ -0,0 +1,57 @@
+## https://github.com/BlueMirrors/Yolov5-ONNX/blob/main/detect.py
+import gradio as gr
+import cv2
+import numpy as np
+from cvu.detector.yolov5 import Yolov5 as Yolov5Onnx
+from vidsz.opencv import Reader, Writer
+
+def process_input(input_file, weight='./jtbz_opt.onnx', device='cpu'):
+    # Determine if the input is an image or a video
+    if input_file.name.endswith(('.jpg', '.jpeg', '.png')):
+        # load model
+        model = Yolov5Onnx(classes="coco",
+                           backend="onnx",
+                           weight=weight,
+                           device=device)
+
+        # read image
+        image = cv2.imread(input_file.name)
+
+        # inference
+        preds = model(image)
+
+        # draw image
+        preds.draw(image)
+
+        # return image
+        return image
+
+    else:
+        model = Yolov5Onnx(classes="coco",
+                           backend="onnx",
+                           weight=weight,
+                           device=device)
+        reader = Reader(input_file.name)
+        writer = Writer(reader, name="output.mp4")
+
+        # warmup
+        warmup = np.random.randint(0, 255, reader.read().shape).astype("float")
+        for i in range(100):
+            model(warmup)
+
+        inference_time = 0
+        for frame in reader:
+            # inference
+            start = time.time()
+            preds = model(frame)
+            inference_time += time.time() - start
+
+            # draw on frame
+            preds.draw(frame)
+            writer.write(frame)
+
+        writer.release()
+        reader.release()
+
+        # return video
+        return gr.outputs.Video("output.mp4")