Update app.py

app.py

@@ -1,4 +1,5 @@
 import subprocess
+import re
 
 # Define the command to be executed
 command = ["python", "setup.py", "build_ext", "--inplace"]
@@ -233,7 +234,7 @@ def load_model(checkpoint):
 
     
     
-def sam_process(input_first_frame_image, checkpoint, tracking_points, trackings_input_label, video_frames_dir, scanned_frames, progress=gr.Progress(track_tqdm=True)):
+def sam_process(input_first_frame_image, checkpoint, tracking_points, trackings_input_label, video_frames_dir, scanned_frames, working_frame, progress=gr.Progress(track_tqdm=True)):
     # 1. We need to preprocess the video and store frames in the right directory
     # — Penser à utiliser un ID unique pour le dossier
  
@@ -254,7 +255,16 @@ def sam_process(input_first_frame_image, checkpoint, tracking_points, trackings_
     # predictor.reset_state(inference_state) # if any previous tracking, reset
 
     # Add new point
-    ann_frame_idx = 0  # the frame index we interact with
+    if working_frame == None:
+        ann_frame_idx = 0  # the frame index we interact with
+    else:
+        # Use a regular expression to find the integer
+        match = re.search(r'frame_(\d+)', working_frame)
+        if match:
+            # Extract the integer from the match
+            frame_number = int(match.group(1))
+            ann_frame_idx = frame_number
+    
     ann_obj_id = 1  # give a unique id to each object we interact with (it can be any integers)
     
     # Let's add a positive click at (x, y) = (210, 350) to get started
@@ -360,6 +370,15 @@ def update_ui(vis_frame_type):
     elif vis_frame_type == "render":
         return gr.update(visible=False), gr.update(visible=True)
 
+def switch_working_frame(working_frame, scanned_frames, video_frames_dir):
+    # Use a regular expression to find the integer
+        match = re.search(r'frame_(\d+)', working_frame)
+        if match:
+            # Extract the integer from the match
+            frame_number = int(match.group(1))
+            ann_frame_idx = frame_number
+            new_working_frame = os.path.join(video_frames_dir, scanned_frames[ann_frame_idx_frame_idx])
+    return new_working_frame, gr.State([]), gr.State([]), new_working_frame, new_working_frame
 
 with gr.Blocks() as demo:
     first_frame_path = gr.State()
@@ -404,7 +423,7 @@ with gr.Blocks() as demo:
                     video_in = gr.Video(label="Video IN")
             
             with gr.Column():
-                working_frame = gr.Dropdown(label="working frame ID", choices=[""], value=None, visible=True, allow_custom_value=True, interactive=True)
+                working_frame = gr.Dropdown(label="working frame ID", choices=[""], value=None, visible=False, allow_custom_value=False, interactive=True)
                 output_result = gr.Image(label="current working mask ref")
                 with gr.Row():
                     vis_frame_type = gr.Radio(label="Propagation level", choices=["check", "render"], value="check", scale=2)
@@ -434,6 +453,13 @@ with gr.Blocks() as demo:
         queue = False
     )
 
+    working_frame.change(
+        fn = switch_working_frame,
+        inputs = [working_frame, scanned_frames, video_frames_dir],
+        outputs = [first_frame_path, tracking_points, trackings_input_label, input_first_frame_image, points_map],
+        queue=False
+    )
+
     submit_btn.click(
         fn = sam_process,
         inputs = [input_first_frame_image, checkpoint, tracking_points, trackings_input_label, video_frames_dir, scanned_frames],