Update app.py

app.py

@@ -301,7 +301,7 @@ def get_mask_sam_process(
     
     return "output_first_frame.jpg", frame_names, predictor, inference_state, gr.update(choices=available_frames_to_check, value=working_frame, visible=True)
 
-def propagate_to_all(video_in, checkpoint, stored_inference_state, stored_frame_names, video_frames_dir, vis_frame_type, progress=gr.Progress(track_tqdm=True)):   
+def propagate_to_all(video_in, checkpoint, stored_inference_state, stored_frame_names, video_frames_dir, vis_frame_type, available_frames_to_check, progress=gr.Progress(track_tqdm=True)):   
     #### PROPAGATION ####
     sam2_checkpoint, model_cfg = load_model(checkpoint)
     predictor = build_sam2_video_predictor(model_cfg, sam2_checkpoint)
@@ -349,11 +349,14 @@ def propagate_to_all(video_in, checkpoint, stored_inference_state, stored_frame_
         # Append the file path to the list
         jpeg_images.append(output_filename)
 
+        if f"frame_{out_frame_idx}.jpg" not in available_frames_to_check:
+            available_frames_to_check.append(f"frame_{out_frame_idx}.jpg")
+
     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=jpeg_images, value=None, visible=True)
+        return gr.update(value=jpeg_images), gr.update(value=None), gr.update(choices=available_frames_to_check, value=None, visible=True), available_frames_to_check
     elif vis_frame_type == "render":
         # Create a video clip from the image sequence
         original_fps = get_video_fps(video_in)
@@ -369,7 +372,7 @@ def propagate_to_all(video_in, checkpoint, stored_inference_state, stored_frame_
             codec='libx264'
         )
         
-        return gr.update(value=None), gr.update(value=final_vid_output_path), None
+        return gr.update(value=None), gr.update(value=final_vid_output_path), None, available_frames_to_check
 
 def update_ui(vis_frame_type):
     if vis_frame_type == "check":
@@ -396,7 +399,7 @@ def reset_propagation(predictor, stored_inference_state):
     
     predictor.reset_state(stored_inference_state)
     print(f"RESET State: {stored_inference_state} ")
-    return gr.update(value=None, visible=False), stored_inference_state, None
+    return gr.update(value=None, visible=False), stored_inference_state, None, gr.State([])
     
 with gr.Blocks() as demo:
     first_frame_path = gr.State()
@@ -437,7 +440,7 @@ with gr.Blocks() as demo:
                 
                 with gr.Row():
                     checkpoint = gr.Dropdown(label="Checkpoint", choices=["tiny", "small", "base-plus", "large"], value="tiny")
-                    submit_btn = gr.Button("Submit", size="lg")
+                    submit_btn = gr.Button("Get Mask", size="lg")
 
                 with gr.Accordion("Your video IN", open=True) as video_in_drawer:
                     video_in = gr.Video(label="Video IN")
@@ -541,7 +544,7 @@ with gr.Blocks() as demo:
     reset_prpgt_brn.click(
         fn = reset_propagation,
         inputs = [loaded_predictor, stored_inference_state],
-        outputs = [output_propagated, stored_inference_state, output_result],
+        outputs = [output_propagated, stored_inference_state, output_result, available_frames_to_check],
         queue=False
     )
 
@@ -552,8 +555,8 @@ with gr.Blocks() as demo:
         queue=False
     ).then(
         fn = propagate_to_all,
-        inputs = [video_in, checkpoint, stored_inference_state, stored_frame_names, video_frames_dir, vis_frame_type],
-        outputs = [output_propagated, output_video, working_frame]
+        inputs = [video_in, checkpoint, stored_inference_state, stored_frame_names, video_frames_dir, vis_frame_type, available_frames_to_check],
+        outputs = [output_propagated, output_video, working_frame, available_frames_to_check]
     )
 
 demo.launch(show_api=False, show_error=True)