Update app.py

app.py

@@ -83,9 +83,16 @@ def preprocess_video_in(video_path):
     
     # Release the video capture object
     cap.release()
+    
+    # scan all the JPEG frame names in this directory
+    scanned_frames = [
+        p for p in os.listdir(output_dir)
+        if os.path.splitext(p)[-1] in [".jpg", ".jpeg", ".JPG", ".JPEG"]
+    ]
+    scanned_frames.sort(key=lambda p: int(os.path.splitext(p)[0]))
 
     # 'image' is the first frame extracted from video_in
-    return first_frame, gr.State([]), gr.State([]), first_frame, first_frame, output_dir, None, None, gr.update(open=False)
+    return first_frame, gr.State([]), gr.State([]), first_frame, first_frame, output_dir, scanned_frames, None, None, gr.update(open=False)
 
 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}")
@@ -220,7 +227,7 @@ def load_model(checkpoint):
 
     
     
-def sam_process(input_first_frame_image, checkpoint, tracking_points, trackings_input_label, video_frames_dir, progress=gr.Progress(track_tqdm=True)):
+def sam_process(input_first_frame_image, checkpoint, tracking_points, trackings_input_label, video_frames_dir, scanned_frames, 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
  
@@ -233,12 +240,7 @@ def sam_process(input_first_frame_image, checkpoint, tracking_points, trackings_
     video_dir = video_frames_dir
     
     # scan all the JPEG frame names in this directory
-    frame_names = [
-        p for p in os.listdir(video_dir)
-        if os.path.splitext(p)[-1] in [".jpg", ".jpeg", ".JPG", ".JPEG"]
-    ]
-    frame_names.sort(key=lambda p: int(os.path.splitext(p)[0]))
-    
+    frame_names = scanned_frames
     
     inference_state = predictor.init_state(video_path=video_dir)
 
@@ -357,6 +359,7 @@ with gr.Blocks() as demo:
     tracking_points = gr.State([])
     trackings_input_label = gr.State([])
     video_frames_dir = gr.State()
+    scanned_frames = gr.State()
     stored_inference_state = gr.State()
     stored_frame_names = gr.State()
     with gr.Column():
@@ -412,7 +415,7 @@ with gr.Blocks() as demo:
     video_in.upload(
         fn = preprocess_video_in, 
         inputs = [video_in], 
-        outputs = [first_frame_path, tracking_points, trackings_input_label, input_first_frame_image, points_map, video_frames_dir, stored_inference_state, stored_frame_names, video_in_drawer],
+        outputs = [first_frame_path, tracking_points, trackings_input_label, input_first_frame_image, points_map, video_frames_dir, scanned_frames, stored_inference_state, stored_frame_names, video_in_drawer],
         queue = False
     )
 
@@ -425,7 +428,7 @@ with gr.Blocks() as demo:
 
     submit_btn.click(
         fn = sam_process,
-        inputs = [input_first_frame_image, checkpoint, tracking_points, trackings_input_label, video_frames_dir],
+        inputs = [input_first_frame_image, checkpoint, tracking_points, trackings_input_label, video_frames_dir, scanned_frames],
         outputs = [output_result, stored_frame_names, stored_inference_state]
     )