Update app.py

app.py

@@ -138,40 +138,56 @@ 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_rate = video.get(cv2.CAP_PROP_FPS)
+    
+    # Create a directory to store the frames
+    frames_dir = 'frames'
+    os.makedirs(frames_dir, exist_ok=True)
 
     frame_count = 0
     while True:
         success, frame = video.read()
         if not success:
             break
-        
-        # 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_output_path = os.path.join(frames_dir, f'frame_{frame_count:04d}.jpg')
+        cv2.imwrite(frame_output_path, frame)
         frame_count += 1
 
     video.release()
+    cv2.destroyAllWindows()
+
+    # Process the frames with object detection and save the results
+    results_dir = 'results'
+    os.makedirs(results_dir, exist_ok=True)
+
+    for i in range(frame_count):
+        frame_path = os.path.join(frames_dir, f'frame_{i:04d}.jpg')
+        frame = cv2.imread(frame_path)
+        results = model(frame)
+        results_output_path = os.path.join(results_dir, f'results_{i:04d}.jpg')
+        cv2.imwrite(results_output_path, results.render()[0])
+
+    # Create the output video from the processed frames
+    frame_files = sorted(os.listdir(results_dir))
+    frame_path = os.path.join(results_dir, frame_files[0])
+    frame = cv2.imread(frame_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 frame_file in frame_files:
+        frame_path = os.path.join(results_dir, frame_file)
+        frame = cv2.imread(frame_path)
+        video_writer.write(frame)
+
     video_writer.release()
 
+    # Clean up the temporary directories
+    os.rmdir(frames_dir)
+    os.rmdir(results_dir)
+
     return video_output_path
 
 
@@ -196,4 +212,4 @@ interfaces = [img_interface, vid_interface]
 tabbed_interface = gr.TabbedInterface(interfaces, ["Image", "Video"])
 
 # Launch the tabbed interface 
-tabbed_interface.launch(debug=True)
+tabbed_interface.launch(debug=True)