Update app.py

app.py

@@ -28,7 +28,7 @@ from PIL import Image, ImageFilter
 from sam2.build_sam import build_sam2_video_predictor
 
 def preprocess_image(image):
-    return image, gr.State([]), gr.State([]), image
+    return image, gr.State([]), gr.State([]), image, gr.State([])
 
 def preprocess_video_in(video_path):
 
@@ -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
+    return first_frame, gr.State([]), gr.State([]), first_frame, first_frame, output_dir, gr.State([]), gr.State([])
 
 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}")
@@ -184,12 +184,7 @@ def show_masks(image, masks, scores, point_coords=None, box_coords=None, input_l
 
     return combined_images, mask_images
 
-
-def sam_process(input_first_frame_image, checkpoint, tracking_points, trackings_input_label, frames_output_dir):
-    # 1. We need to preprocess the video and store frames in the right directory
-    # — Penser à utiliser un ID unique pour le dossier
-
-
+def load_model(checkpoint):
     # Load model accordingly to user's choice
     if checkpoint == "tiny":
         sam2_checkpoint = "./checkpoints/sam2_hiera_tiny.pt"
@@ -203,13 +198,20 @@ def sam_process(input_first_frame_image, checkpoint, tracking_points, trackings_
     elif checkpoint == "large":
         sam2_checkpoint = "./checkpoints/sam2_hiera_large.pt"
         model_cfg = "sam2_hiera_l.yaml"
+
+    return sam2_checkpoint, model_cfg
     
+def sam_process(input_first_frame_image, checkpoint, tracking_points, trackings_input_label, video_frames_dir):
+    # 1. We need to preprocess the video and store frames in the right directory
+    # — Penser à utiliser un ID unique pour le dossier
+ 
+    sam2_checkpoint, model_cfg = load_model(checkpoint)
     predictor = build_sam2_video_predictor(model_cfg, sam2_checkpoint)
 
     
     # `video_dir` a directory of JPEG frames with filenames like `<frame_index>.jpg`
-    print(f"STATE FRAME OUTPUT DIRECTORY: {frames_output_dir}")
-    video_dir = frames_output_dir
+    print(f"STATE FRAME OUTPUT DIRECTORY: {video_frames_dir}")
+    video_dir = video_frames_dir
     
     # scan all the JPEG frame names in this directory
     frame_names = [
@@ -248,13 +250,18 @@ def sam_process(input_first_frame_image, checkpoint, tracking_points, trackings_
     show_mask((out_mask_logits[0] > 0.0).cpu().numpy(), plt.gca(), obj_id=out_obj_ids[0])
     
     # Save the plot as a JPG file
-    output_filename = "output_frame.jpg"
-    plt.savefig(output_filename, format='jpg')
+    first_frame_output_filename = "output_first_frame.jpg"
+    plt.savefig(first_frame_output_filename, format='jpg')
     plt.close()
 
-    
-    #### PROPAGATION ####
+    return "output_first_frame.jpg", frame_names, inference_state
 
+def propagate_to_all(checkpoint, stored_inference_state, stored_frame_names):   
+    #### PROPAGATION ####
+    sam2_checkpoint, model_cfg = load_model(checkpoint)
+    predictor = build_sam2_video_predictor(model_cfg, sam2_checkpoint)
+    inference_state = stored_inference_state.value
+    frame_names = stored_frame_names.value
     # Define a directory to save the JPEG images
     frames_output_dir = "frames_output_images"
     os.makedirs(frames_output_dir, exist_ok=True)
@@ -289,16 +296,16 @@ def sam_process(input_first_frame_image, checkpoint, tracking_points, trackings_
     
         # Close the plot
         plt.close()
-    
-    # OLD
 
-    return output_filename, jpeg_images
+    return jpeg_images
 
 with gr.Blocks() as demo:
     first_frame_path = gr.State()
     tracking_points = gr.State([])
     trackings_input_label = gr.State([])
-    frames_output_dir = gr.State()
+    video_frames_dir = gr.State()
+    stored_inference_state = gr.State([])
+    stored_frame_names = gr.State([])
     with gr.Column():
         gr.Markdown("# SAM2 Video Predictor")
         gr.Markdown("This is a simple demo for video segmentation with SAM2.")
@@ -325,20 +332,21 @@ with gr.Blocks() as demo:
                 submit_btn = gr.Button("Submit")
             with gr.Column():
                 output_result = gr.Image()
+                propagate_btn = gr.Button("Propagate")
                 output_propagated = gr.Gallery()
                 # output_result_mask = gr.Image()
     
     clear_points_btn.click(
         fn = preprocess_image,
         inputs = input_first_frame_image, 
-        outputs = [first_frame_path, tracking_points, trackings_input_label, points_map],
+        outputs = [first_frame_path, tracking_points, trackings_input_label, points_map, stored_inference_state],
         queue=False
     )
     
     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, frames_output_dir],
+        outputs = [first_frame_path, tracking_points, trackings_input_label, input_first_frame_image, points_map, video_frames_dir, stored_inference_state, stored_frame_names],
         queue = False
     )
 
@@ -351,8 +359,14 @@ with gr.Blocks() as demo:
 
     submit_btn.click(
         fn = sam_process,
-        inputs = [input_first_frame_image, checkpoint, tracking_points, trackings_input_label, frames_output_dir],
-        outputs = [output_result, output_propagated]
+        inputs = [input_first_frame_image, checkpoint, tracking_points, trackings_input_label, video_frames_dir],
+        outputs = [output_result, stored_frame_names, stored_inference_state]
+    )
+
+    propagate_btn.click(
+        fn = propagate_to_all,
+        inputs = [checkpoint, stored_inference_state, stored_frame_names],
+        outputs = [output_propagated]
     )
 
 demo.launch(show_api=False, show_error=True)