Update app.py

app.py

@@ -256,13 +256,15 @@ def sam_process(input_first_frame_image, checkpoint, tracking_points, trackings_
 
     return "output_first_frame.jpg", frame_names, inference_state
 
-def propagate_to_all(checkpoint, stored_inference_state, stored_frame_names, video_frames_dir):   
+def propagate_to_all(checkpoint, stored_inference_state, stored_frame_names, video_frames_dir, vis_frame_type):   
     #### PROPAGATION ####
     sam2_checkpoint, model_cfg = load_model(checkpoint)
     predictor = build_sam2_video_predictor(model_cfg, sam2_checkpoint)
+    
     inference_state = stored_inference_state
     frame_names = stored_frame_names
     video_dir = video_frames_dir
+    
     # Define a directory to save the JPEG images
     frames_output_dir = "frames_output_images"
     os.makedirs(frames_output_dir, exist_ok=True)
@@ -279,7 +281,10 @@ def propagate_to_all(checkpoint, stored_inference_state, stored_frame_names, vid
         }
     
     # render the segmentation results every few frames
-    vis_frame_stride = 15
+    if vis_frame_type == "check":
+        vis_frame_stride = 15
+    elif vis_frame_type == "render":
+        vis_frame_stride = 1
     plt.close("all")
     for out_frame_idx in range(0, len(frame_names), vis_frame_stride):
         plt.figure(figsize=(6, 4))
@@ -298,7 +303,11 @@ def propagate_to_all(checkpoint, stored_inference_state, stored_frame_names, vid
         # Close the plot
         plt.close()
 
-    return jpeg_images
+    if vis_frame_type == "check":
+        return gr.update(value=jpeg_images, visible=True), gr.update(visible=False, value=None)
+    elif vis_frame_type == "render":
+        return gr.update(visible=False, value=None), gr.update(value=final_vid, visible=True)
+    
 
 with gr.Blocks() as demo:
     first_frame_path = gr.State()
@@ -323,7 +332,7 @@ with gr.Blocks() as demo:
                 points_map = gr.Image(
                     label="points map", 
                     type="filepath",
-                    interactive=True
+                    interactive=False
                 )
                 video_in = gr.Video(label="Video IN")
                 with gr.Row():
@@ -333,8 +342,11 @@ with gr.Blocks() as demo:
                 submit_btn = gr.Button("Submit")
             with gr.Column():
                 output_result = gr.Image()
-                propagate_btn = gr.Button("Propagate")
-                output_propagated = gr.Gallery()
+                with gr.Row():
+                    vis_frame_type = gr.Radio(choices=["check", "render"], value="render", scale=2)
+                    propagate_btn = gr.Button("Propagate", scale=1)
+                output_propagated = gr.Gallery(visible=False)
+                output_video = gr.Video(visible=False)
                 # output_result_mask = gr.Image()
     
     clear_points_btn.click(
@@ -366,8 +378,8 @@ with gr.Blocks() as demo:
 
     propagate_btn.click(
         fn = propagate_to_all,
-        inputs = [checkpoint, stored_inference_state, stored_frame_names, video_frames_dir],
-        outputs = [output_propagated]
+        inputs = [checkpoint, stored_inference_state, stored_frame_names, video_frames_dir, vis_frame_type],
+        outputs = [output_propagated, output_video]
     )
 
 demo.launch(show_api=False, show_error=True)