modified app.py

app.py

@@ -2,7 +2,6 @@ import gradio as gr
 import cv2
 import requests
 import os
-
 from ultralytics import YOLO
 
 file_urls = [
@@ -12,36 +11,29 @@ file_urls = [
 ]
 
 def download_file(url, save_name):
-    url = url
     if not os.path.exists(save_name):
         file = requests.get(url)
-        open(save_name, 'wb').write(file.content)
+        with open(save_name, 'wb') as f:
+            f.write(file.content)
 
 for i, url in enumerate(file_urls):
-    if 'mp4' in file_urls[i]:
-        download_file(
-            file_urls[i],
-            f"video.mp4"
-        )
+    if url.endswith('.mp4') or '.mp4?' in url:
+        download_file(url, 'video.mp4')
     else:
-        download_file(
-            file_urls[i],
-            f"image_{i}.jpg"
-        )
+        download_file(url, f'image_{i}.jpg')
 
 model = YOLO('best.pt')
-path  = [['image_0.jpg'], ['image_1.jpg']]
+path = [['image_0.jpg'], ['image_1.jpg']]
 video_path = [['video.mp4']]
 
 def show_preds_image(image_path):
     image = cv2.imread(image_path)
-    outputs = model.predict(source=image_path)
-    results = outputs[0].cpu().numpy()
-    for i, det in enumerate(results.boxes.xyxy):
+    results = model(image_path)[0]
+    for box in results.boxes.xyxy:
         cv2.rectangle(
             image,
-            (int(det[0]), int(det[1])),
-            (int(det[2]), int(det[3])),
+            (int(box[0]), int(box[1])),
+            (int(box[2]), int(box[3])),
             color=(0, 0, 255),
             thickness=2,
             lineType=cv2.LINE_AA
@@ -49,10 +41,10 @@ def show_preds_image(image_path):
     return cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
 
 inputs_image = [
-    gr.components.Image(type="filepath", label="Input Image"),
+    gr.Image(type="filepath", label="Input Image"),
 ]
 outputs_image = [
-    gr.components.Image(type="numpy", label="Output Image"),
+    gr.Image(type="numpy", label="Output Image"),
 ]
 interface_image = gr.Interface(
     fn=show_preds_image,
@@ -65,29 +57,29 @@ interface_image = gr.Interface(
 
 def show_preds_video(video_path):
     cap = cv2.VideoCapture(video_path)
-    while(cap.isOpened()):
+    while cap.isOpened():
         ret, frame = cap.read()
-        if ret:
-            frame_copy = frame.copy()
-            outputs = model.predict(source=frame)
-            results = outputs[0].cpu().numpy()
-            for i, det in enumerate(results.boxes.xyxy):
-                cv2.rectangle(
-                    frame_copy,
-                    (int(det[0]), int(det[1])),
-                    (int(det[2]), int(det[3])),
-                    color=(0, 0, 255),
-                    thickness=2,
-                    lineType=cv2.LINE_AA
-                )
-            yield cv2.cvtColor(frame_copy, cv2.COLOR_BGR2RGB)
+        if not ret:
+            break
+        frame_copy = frame.copy()
+        results = model(frame_copy)[0]
+        for box in results.boxes.xyxy:
+            cv2.rectangle(
+                frame_copy,
+                (int(box[0]), int(box[1])),
+                (int(box[2]), int(box[3])),
+                color=(0, 0, 255),
+                thickness=2,
+                lineType=cv2.LINE_AA
+            )
+        yield cv2.cvtColor(frame_copy, cv2.COLOR_BGR2RGB)
+    cap.release()
 
 inputs_video = [
-    gr.components.Video(type="filepath", label="Input Video"),
-
+    gr.Video(label="Input Video"),
 ]
 outputs_video = [
-    gr.components.Image(type="numpy", label="Output Image"),
+    gr.Image(type="numpy", label="Output Frame"),
 ]
 interface_video = gr.Interface(
     fn=show_preds_video,
@@ -101,4 +93,4 @@ interface_video = gr.Interface(
 gr.TabbedInterface(
     [interface_image, interface_video],
     tab_names=['Image inference', 'Video inference']
-).queue().launch()
+).queue().launch()