Update app.py

app.py

@@ -123,7 +123,7 @@ def preprocess_video_in(video_path):
         gr.update(open=False)  # video_in_drawer
     ]
 
-def get_point(point_type, tracking_points, trackings_input_label, first_frame_path, evt: gr.SelectData):
+def get_point(point_type, tracking_points, trackings_input_label, input_first_frame_path, evt: gr.SelectData):
     print(f"You selected {evt.value} at {evt.index} from {evt.target}")
 
     tracking_points.value.append(evt.index)
@@ -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, available_frames_to_check, 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, working_frame, progress=gr.Progress(track_tqdm=True)):   
     #### PROPAGATION ####
     sam2_checkpoint, model_cfg = load_model(checkpoint)
     predictor = build_sam2_video_predictor(model_cfg, sam2_checkpoint)
@@ -356,7 +356,7 @@ def propagate_to_all(video_in, checkpoint, stored_inference_state, stored_frame_
     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=None, visible=True), available_frames_to_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
     elif vis_frame_type == "render":
         # Create a video clip from the image sequence
         original_fps = get_video_fps(video_in)
@@ -372,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, available_frames_to_check
+        return gr.update(value=None), gr.update(value=final_vid_output_path), working_frame, available_frames_to_check
 
 def update_ui(vis_frame_type):
     if vis_frame_type == "check":
@@ -556,7 +556,7 @@ 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, available_frames_to_check],
+        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]
     )