Update app.py

app.py

@@ -316,6 +316,7 @@ def propagate_to_all(video_in, checkpoint, stored_inference_state, stored_frame_
     
     # Initialize a list to store file paths of saved images
     jpeg_images = []
+    masks_frames = []
 
     # run propagation throughout the video and collect the results in a dict
     video_segments = {}  # video_segments contains the per-frame segmentation results
@@ -352,11 +353,18 @@ def propagate_to_all(video_in, checkpoint, stored_inference_state, stored_frame_
         if f"frame_{out_frame_idx}.jpg" not in available_frames_to_check:
             available_frames_to_check.append(f"frame_{out_frame_idx}.jpg")
 
+        # Save the raw binary mask as a separate image
+        mask_filename = os.path.join(frames_output_dir, f"mask_{out_frame_idx}.jpg")
+        binary_mask = (out_mask * 255).astype(np.uint8)  # Scale mask to 0-255
+        mask_image = Image.fromarray(binary_mask)
+        mask_image.save(mask_filename)  # Save the mask as a JPEG
+        masks_frames.append(mask_filename)  # Append to the list of masks
+
     torch.cuda.empty_cache()
     print(f"JPEG_IMAGES: {jpeg_images}")
 
     if vis_frame_type == "check":
-        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)
+        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)
@@ -371,8 +379,12 @@ def propagate_to_all(video_in, checkpoint, stored_inference_state, stored_frame_
             final_vid_output_path,
             codec='libx264'
         )
+
+        mask_clip = ImageSequenceClip(masks_frames, fps=fps)
+        # Write the result to a file
+        mask_final_vid_output_path = "mask_output_video.mp4"
         
-        return gr.update(value=None), gr.update(value=final_vid_output_path), working_frame, available_frames_to_check, gr.update(visible=True)
+        return gr.update(value=None), gr.update(value=final_vid_output_path), working_frame, available_frames_to_check, gr.update(visible=True), mask_final_vid_output_path
 
 def update_ui(vis_frame_type):
     if vis_frame_type == "check":
@@ -478,6 +490,7 @@ with gr.Blocks(css=css) as demo:
                 reset_prpgt_brn = gr.Button("Reset", visible=False)
                 output_propagated = gr.Gallery(label="Propagated Mask samples gallery", columns=4, visible=False)
                 output_video = gr.Video(visible=False)
+                mask_final_output = gr.Video(label="Mask Video")
                 # output_result_mask = gr.Image()
     
     
@@ -581,7 +594,7 @@ with gr.Blocks(css=css) as demo:
     ).then(
         fn = propagate_to_all,
         inputs = [video_in, checkpoint, stored_inference_state, stored_frame_names, video_frames_dir, vis_frame_type, available_frames_to_check, working_frame],
-        outputs = [output_propagated, output_video, working_frame, available_frames_to_check, reset_prpgt_brn]
+        outputs = [output_propagated, output_video, working_frame, available_frames_to_check, reset_prpgt_brn, mask_final_output]
     )
 
 demo.launch(show_api=False, show_error=True)