Update app.py

app.py

@@ -70,7 +70,7 @@ def preprocess_video_in(video_path):
     cap.release()
 
     # 'image' is the first frame extracted from video_in
-    return first_frame, gr.State([]), gr.State([]), first_frame, first_frame, output_dir, gr.State(), gr.State()
+    return first_frame, gr.State([]), gr.State([]), first_frame, first_frame, output_dir, None, None
 
 def get_point(point_type, tracking_points, trackings_input_label, first_frame_path, evt: gr.SelectData):
     print(f"You selected {evt.value} at {evt.index} from {evt.target}")
@@ -256,12 +256,13 @@ 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):   
+def propagate_to_all(checkpoint, stored_inference_state, stored_frame_names, video_frames_dir):   
     #### 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)
@@ -365,7 +366,7 @@ with gr.Blocks() as demo:
 
     propagate_btn.click(
         fn = propagate_to_all,
-        inputs = [checkpoint, stored_inference_state, stored_frame_names],
+        inputs = [checkpoint, stored_inference_state, stored_frame_names, video_frames_dir],
         outputs = [output_propagated]
     )