Update app.py

app.py

@@ -7,12 +7,12 @@ import os
 from PIL import Image
 import numpy as np
 import threading
+import skvideo.io
 import cv2
 
 should_continue = True
 
-
-@spaces.GPU(duration=120)    
+@spaces.GPU(duration=120)
 def yolov9_inference(model_id, img_path=None, vid_path=None, tracking_algorithm = None):
     global should_continue
     img_extensions = ['.jpg', '.jpeg', '.png', '.gif']  # Add more image extensions if needed
@@ -33,10 +33,40 @@ def yolov9_inference(model_id, img_path=None, vid_path=None, tracking_algorithm
         print(input_path)
         output_path = run(weights=model_id, imgsz=(image_size,image_size), conf_thres=conf_threshold, iou_thres=iou_threshold, source=input_path, device='0', hide_conf= True)
     elif vid_path is not None:
-        #_, vid_extension = os.path.splitext(vid_path)
-        #if vid_extension.lower() in vid_extensions:
-        input_path = vid_path
-        print(input_path)
+        vid_name = 'output.mp4'
+
+        # Create a VideoCapture object
+        cap = cv2.VideoCapture(vid_path)
+
+        # Check if video opened successfully
+        if not cap.isOpened():
+            print("Error opening video file")
+
+        # Read the video frame by frame
+        frames = []
+        while cap.isOpened():
+            ret, frame = cap.read()
+            if ret:
+                frames.append(frame)
+            else:
+                break
+
+        # Release the VideoCapture object
+        cap.release()
+
+        # Convert the list of frames to a numpy array
+        vid_data = np.array(frames)
+
+        # Create a VideoWriter object
+        out = cv2.VideoWriter(vid_name, cv2.VideoWriter_fourcc(*'mp4v'), 30, (frames[0].shape[1], frames[0].shape[0]))
+
+        # Write the frames to the output video file
+        for frame in frames:
+            out.write(frame)
+
+        # Release the VideoWriter object
+        out.release()
+        input_path = vid_name
         if tracking_algorithm == 'deep_sort':
             output_path = run_deepsort(weights=model_id, imgsz=(image_size,image_size), conf_thres=conf_threshold, iou_thres=iou_threshold, source=input_path, device='0', draw_trails=True)
         elif tracking_algorithm == 'strong_sort':
@@ -66,8 +96,10 @@ iface = gr.Interface(
         gr.Dropdown(
             label="Model",
             choices=[
-                "yolov9_e_trained-converted.pt",
-                "best_model-converted.pt"
+                "last_best_model.pt",
+                "best_model-converted.pt",
+                "yolov9_e_trained.pt",
+                "best_model-converted-reparams.pt"
             ],
             value="last_best_model.pt"
         ),
@@ -89,8 +121,8 @@ iface = gr.Interface(
         gr.Textbox(label="Output path")
     ],
     examples=[
-        ["best_model-converted.pt", "camera1_A_133.png", None, "deep_sort"],
-        ["best_model-converted.pt", None, "test.mp4", "strong_sort"]
+        ["last_best_model.pt", "camera1_A_133.png", None, "deep_sort"],
+        ["last_best_model.pt", None, "test.mp4", "strong_sort"]
     ],
     title='YOLOv9: Real-time Object Detection',
     description='This is a real-time object detection system using YOLOv9.',