Update app.py

app.py

@@ -138,38 +138,42 @@ def detect_image(image):
 
 def detect_video(video_path):
     video = cv2.VideoCapture(video_path)
+    frame_rate = video.get(cv2.CAP_PROP_FPS)  # Get the frame rate of the input video
+
+    # Get the dimensions of the video frames
+    width = int(video.get(cv2.CAP_PROP_FRAME_WIDTH))
+    height = int(video.get(cv2.CAP_PROP_FRAME_HEIGHT))
+
+    # Create a VideoWriter object to write the output video
+    video_output_path = 'output_video.mp4'  # Replace with your desired output video path
+    fourcc = cv2.VideoWriter_fourcc(*'mp4v')  # You can change the codec as needed
+    video_writer = cv2.VideoWriter(video_output_path, fourcc, frame_rate, (width, height))
+
     frame_count = 0
     while True:
         success, frame = video.read()
         if not success:
             break
-        frame = model(frame)[0].numpy().astype('uint8')
-        frame_output_path = f'frame_{frame_count}.jpg'  # Replace with your desired output path
-        cv2.imwrite(frame_output_path, frame)
+        
+        # Process the frame using your model
+        results = model(frame)
+        # Access the first detection result
+        detection_result = results.pandas().xyxy.iloc[0]
+        # Extract the bounding box coordinates
+        xmin, ymin, xmax, ymax = detection_result["xmin"], detection_result["ymin"], detection_result["xmax"], detection_result["ymax"]
+        # Draw the bounding box on the frame
+        frame = cv2.rectangle(frame, (xmin, ymin), (xmax, ymax), (0, 255, 0), 2)
+        
+        # Write the frame to the output video
+        video_writer.write(frame)
+
         frame_count += 1
-    video.release()
-    cv2.destroyAllWindows()
 
-    frame_rate = video.get(cv2.CAP_PROP_FPS)
-    
-    image_dir = '.'  # Replace with the directory containing the image files (frames)
-    image_files = sorted(os.listdir(image_dir))
-    
-    image_path = os.path.join(image_dir, image_files[0])
-    frame = cv2.imread(image_path)
-    height, width, _ = frame.shape
-    
-    video_output_path = 'output_video.mp4'  # Replace with your desired output video path
-    fourcc = cv2.VideoWriter_fourcc(*'mp4v')  # You can change the codec as needed
-    video_writer = cv2.VideoWriter(video_output_path, fourcc, frame_rate, (width, height))
-    
-    for image_file in image_files:
-        image_path = os.path.join(image_dir, image_file)
-        frame = cv2.imread(image_path)
-        video_writer.write(frame)
-    
+    video.release()
     video_writer.release()
-    return video_writer
+
+    return video_output_path
+
 
 # Create Gradio interfaces for different modes
 img_interface = gr.Interface(