Update app.py

app.py

@@ -31,7 +31,7 @@ from sam2.build_sam import build_sam2_video_predictor
 
 from moviepy.editor import ImageSequenceClip
 
-def get_video_fps(video_path):
+def get_video_properties(video_path):
     # Open the video file
     cap = cv2.VideoCapture(video_path)
     
@@ -41,8 +41,15 @@ def get_video_fps(video_path):
     
     # Get the FPS of the video
     fps = cap.get(cv2.CAP_PROP_FPS)
+    
+    # Get the width and height of the video
+    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+
+    # Release the video capture object
+    cap.release()
 
-    return fps
+    return fps, width, height
 
 def clear_points(image):
     # we clean all
@@ -395,7 +402,7 @@ def propagate_to_all(video_in, checkpoint, stored_inference_state, stored_frame_
         return gr.update(value=jpeg_images), gr.update(value=None), gr.update(choices=available_frames_to_check, value=working_frame, visible=True), available_frames_to_check, gr.update(visible=True), None
     elif vis_frame_type == "render":
         # Create a video clip from the image sequence
-        original_fps = get_video_fps(video_in)
+        original_fps, o_width, o_height = get_video_properties(video_in)
         fps = original_fps  # Frames per second
         total_frames = len(jpeg_images)
         clip = ImageSequenceClip(jpeg_images, fps=fps)
@@ -413,20 +420,22 @@ def propagate_to_all(video_in, checkpoint, stored_inference_state, stored_frame_
         # Create a video from the masks_images array
         mask_video_filename = "final_masks_video.mp4"
         
-        # Get the dimensions of the first mask image
-        frame = cv2.imread(masks_images[0])
-        height, width, _ = frame.shape
-        
-        # Define the video writer
+        # Define the video writer with the original video's dimensions
         fourcc = cv2.VideoWriter_fourcc(*"mp4v")
-        fps = original_fps  # Frames per second
-        video_writer = cv2.VideoWriter(mask_video_filename, fourcc, fps, (width, height))
+        video_writer = cv2.VideoWriter(mask_video_filename, fourcc, fps, (o_width, o_height))
         
         # Write each mask image to the video
         for mask_path in masks_images:
+            # Read the mask frame
             frame = cv2.imread(mask_path)
-            video_writer.write(frame)
+            
+            # Resize the mask frame to the original video's dimensions
+            frame_resized = cv2.resize(frame, (o_width, o_height), interpolation=cv2.INTER_NEAREST)
+            
+            # Write the resized frame to the video
+            video_writer.write(frame_resized)
         
+        # Release the video writer
         video_writer.release()
         print(f"Mask Video saved at {mask_video_filename}")