Update app.py

app.py

@@ -299,7 +299,7 @@ def get_mask_sam_process(
         available_frames_to_check.append(working_frame)
         print(available_frames_to_check)
     
-    return "output_first_frame.jpg", frame_names, inference_state, gr.update(choices=available_frames_to_check, value=working_frame, visible=True)
+    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)):   
     #### PROPAGATION ####
@@ -392,12 +392,19 @@ def switch_working_frame(working_frame, scanned_frames, video_frames_dir):
             new_working_frame = os.path.join(video_frames_dir, scanned_frames[ann_frame_idx])
             return new_working_frame, gr.State([]), gr.State([]), new_working_frame, new_working_frame, new_working_frame
 
+def reset_propagation(predictor, stored_inference_state):
+    
+    predictor.reset_state(stored_inference_state)
+    print(f"RESET State: {stored_inference_state} ")
+    return stored_inference_state
+    
 with gr.Blocks() as demo:
     first_frame_path = gr.State()
     tracking_points = gr.State([])
     trackings_input_label = gr.State([])
     video_frames_dir = gr.State()
     scanned_frames = gr.State()
+    loaded_predictor = gr.State()
     stored_inference_state = gr.State()
     stored_frame_names = gr.State()
     available_frames_to_check = gr.State([])
@@ -442,6 +449,7 @@ with gr.Blocks() as demo:
                 with gr.Row():
                     vis_frame_type = gr.Radio(label="Propagation level", choices=["check", "render"], value="check", scale=2)
                     propagate_btn = gr.Button("Propagate", scale=1)
+                    reset_prpgt_brn = gr.Button("Reset", scale=0.75)
                 output_propagated = gr.Gallery(label="Propagated Mask samples gallery", visible=False)
                 output_video = gr.Video(visible=False)
                 # output_result_mask = gr.Image()
@@ -524,11 +532,19 @@ with gr.Blocks() as demo:
         outputs = [
             output_result, 
             stored_frame_names, 
+            loaded_predictor,
             stored_inference_state,
             working_frame,
         ]
     )
 
+    reset_prpgt_brn.click(
+        fn = reset_propagation,
+        inputs = [loaded_predictor, stored_inference_state],
+        outputs = [stored_inference_state],
+        queue=False
+    )
+
     propagate_btn.click(
         fn = update_ui,
         inputs = [vis_frame_type],