efficient track anything demo

app.py

@@ -0,0 +1,638 @@
+import subprocess
+import re
+from typing import List, Tuple, Optional
+
+import gradio as gr
+from datetime import datetime
+import os
+os.environ["TORCH_CUDNN_SDPA_ENABLED"] = "0,1,2,3,4,5,6,7"
+import torch
+import numpy as np
+import cv2
+import matplotlib.pyplot as plt
+from PIL import Image, ImageFilter
+from sam2.build_sam import build_sam2_video_predictor
+
+from moviepy.editor import ImageSequenceClip
+
+# Description
+title = "<center><strong><font size='8'>Efficient Track Anything (EfficientTAM)<font></strong></center>"
+
+description_e = """This is a demo of [Efficient Track Anything (EfficientTAM) Model](https://github.com/yformer/EfficientTAM).
+              """
+
+description_p = """# Interactive Video Segmentation
+                - Built our demo based on [SAM2-Video-Predictor](https://huggingface.co/spaces/fffiloni/SAM2-Video-Predictor). Thanks to Sylvain Filoni.
+                - Instruction
+                <ol>
+                <li> Upload one video or click one example video</li>
+                <li> Click 'include' point type, select the object to segment and track</li>
+                <li> Click 'exclude' point type (optional), select the area you want to avoid segmenting and tracking</li>
+                <li> Click the 'Segment' button, obtain the mask of the first frame </li>
+                <li> Click the 'coarse' level and the 'Track' button, segment and track the object every 15 frames </li>
+                <li> Click the corresponding frame to add points on the object for mask refining (optional) </li>
+                <li> Click the 'fine' level and the 'Track' button, obtain masklet and masked video </li>
+                <li> Click the 'Reset' button to restart </li>
+                </ol>
+                - Github [link](https://github.com/yformer/EfficientTAM)
+              """
+
+# examples
+examples = [
+    ["examples/videos/cat.mp4"],
+    ["examples/videos/coffee.mp4"],
+    ["examples/videos/car.mp4"],
+    ["examples/videos/chick.mp4"],
+    ["examples/videos/cups.mp4"],
+    ["examples/videos/dog.mp4"],
+    ["examples/videos/goat.mp4"],
+    ["examples/videos/juggle.mp4"],
+    ["examples/videos/street.mp4"],
+    ["examples/videos/yacht.mp4"],
+]
+
+default_example = examples[0]
+
+def get_video_fps(video_path):
+    # Open the video file
+    cap = cv2.VideoCapture(video_path)
+    
+    if not cap.isOpened():
+        print("Error: Could not open video.")
+        return None
+    
+    # Get the FPS of the video
+    fps = cap.get(cv2.CAP_PROP_FPS)
+
+    return fps
+
+def clear_points(image):
+    # we clean all
+    return [
+        image,   # first_frame_path
+        gr.State([]),      # tracking_points
+        gr.State([]),      # trackings_input_label
+        image,   # points_map
+        #gr.State()     # stored_inference_state
+    ]
+
+def preprocess_video_in(video_path):
+    if video_path is None:
+        return None, gr.State([]), gr.State([]), None, None, None, None, None, None, gr.update(open=True)
+
+    # Generate a unique ID based on the current date and time
+    unique_id = datetime.now().strftime('%Y%m%d%H%M%S')
+    
+    # Set directory with this ID to store video frames 
+    extracted_frames_output_dir = f'frames_{unique_id}'
+    
+    # Create the output directory
+    os.makedirs(extracted_frames_output_dir, exist_ok=True)
+
+    ### Process video frames ###
+    # Open the video file
+    cap = cv2.VideoCapture(video_path)
+    
+    if not cap.isOpened():
+        print("Error: Could not open video.")
+        return None
+
+    # Get the frames per second (FPS) of the video
+    fps = cap.get(cv2.CAP_PROP_FPS)
+    
+    # Calculate the number of frames to process (10 seconds of video)
+    max_frames = int(fps * 10)
+    
+    frame_number = 0
+    first_frame = None
+    
+    while True:
+        ret, frame = cap.read()
+        if not ret or frame_number >= max_frames:
+            break
+        
+        # Format the frame filename as '00000.jpg'
+        frame_filename = os.path.join(extracted_frames_output_dir, f'{frame_number:05d}.jpg')
+        
+        # Save the frame as a JPEG file
+        cv2.imwrite(frame_filename, frame)
+        
+        # Store the first frame
+        if frame_number == 0:
+            first_frame = frame_filename
+        
+        frame_number += 1
+    
+    # Release the video capture object
+    cap.release()
+    
+    # scan all the JPEG frame names in this directory
+    scanned_frames = [
+        p for p in os.listdir(extracted_frames_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]))
+    # print(f"SCANNED_FRAMES: {scanned_frames}")
+    
+    return [
+        first_frame,           # first_frame_path
+        gr.State([]),          # tracking_points
+        gr.State([]),          # trackings_input_label
+        first_frame,           # input_first_frame_image
+        first_frame,           # points_map
+        extracted_frames_output_dir,            # video_frames_dir
+        scanned_frames,        # scanned_frames
+        None,                  # stored_inference_state
+        None,                  # stored_frame_names
+        gr.update(open=False)  # video_in_drawer
+    ]
+
+
+def get_point(point_type, tracking_points, trackings_input_label, input_first_frame_image, evt: gr.SelectData):
+    if input_first_frame_image is None:
+        return gr.State([]), gr.State([]), None
+    print(f"You selected {evt.value} at {evt.index} from {evt.target}")
+
+    tracking_points.value.append(evt.index)
+    print(f"TRACKING POINT: {tracking_points.value}")
+
+    if point_type == "include":
+        trackings_input_label.value.append(1)
+    elif point_type == "exclude":
+        trackings_input_label.value.append(0)
+    print(f"TRACKING INPUT LABEL: {trackings_input_label.value}")
+    
+    # Open the image and get its dimensions
+    transparent_background = Image.open(input_first_frame_image).convert('RGBA')
+    w, h = transparent_background.size
+    
+    # Define the circle radius as a fraction of the smaller dimension
+    fraction = 0.02  # You can adjust this value as needed
+    radius = int(fraction * min(w, h))
+    
+    # Create a transparent layer to draw on
+    transparent_layer = np.zeros((h, w, 4), dtype=np.uint8)
+    
+    for index, track in enumerate(tracking_points.value):
+        if trackings_input_label.value[index] == 1:
+            cv2.circle(transparent_layer, track, radius, (0, 255, 0, 255), -1)
+        else:
+            cv2.circle(transparent_layer, track, radius, (255, 0, 0, 255), -1)
+
+    # Convert the transparent layer back to an image
+    transparent_layer = Image.fromarray(transparent_layer, 'RGBA')
+    selected_point_map = Image.alpha_composite(transparent_background, transparent_layer)
+    
+    return tracking_points, trackings_input_label, selected_point_map
+
+DEVICE = 'cuda'
+# use bfloat16 for the entire notebook
+torch.autocast(device_type="cuda", dtype=torch.bfloat16).__enter__()
+if torch.cuda.get_device_properties(0).major >= 8:
+    # turn on tfloat32 for Ampere GPUs (https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices)
+    torch.backends.cuda.matmul.allow_tf32 = True
+    torch.backends.cudnn.allow_tf32 = True
+    
+def show_mask(mask, ax, obj_id=None, random_color=False):
+    if random_color:
+        color = np.concatenate([np.random.random(3), np.array([0.6])], axis=0)
+    else:
+        cmap = plt.get_cmap("tab10")
+        cmap_idx = 0 if obj_id is None else obj_id
+        color = np.array([*cmap(cmap_idx)[:3], 0.6])
+    h, w = mask.shape[-2:]
+    mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1)
+    ax.axis('off')
+    ax.imshow(mask_image)
+
+
+def show_points(coords, labels, ax, marker_size=200):
+    pos_points = coords[labels==1]
+    neg_points = coords[labels==0]
+    ax.scatter(pos_points[:, 0], pos_points[:, 1], color='green', marker='*', s=marker_size, edgecolor='white', linewidth=1.25)
+    ax.scatter(neg_points[:, 0], neg_points[:, 1], color='red', marker='*', s=marker_size, edgecolor='white', linewidth=1.25)
+
+def show_box(box, ax):
+    x0, y0 = box[0], box[1]
+    w, h = box[2] - box[0], box[3] - box[1]
+    ax.add_patch(plt.Rectangle((x0, y0), w, h, edgecolor='green', facecolor=(0, 0, 0, 0), lw=2))    
+
+
+def load_model(checkpoint):
+    # Load model accordingly to user's choice
+    if checkpoint == "efficienttam-s":
+        efficienttam_checkpoint = "./checkpoints/efficienttam_s.pt"
+        model_cfg = "efficienttam-s.yaml"
+        return [efficienttam_checkpoint, model_cfg]
+    elif checkpoint == "efficienttam-ti":
+        efficienttam_checkpoint = "./checkpoints/efficienttam-ti.pt"
+        model_cfg = "efficienttam-ti.yaml"
+        return [efficienttam_checkpoint, model_cfg]
+    elif checkpoint == "efficienttam-s-512x512":
+        efficienttam_checkpoint = "./checkpoints/efficienttam-s-512x512.pt"
+        model_cfg = "efficienttam_s_512x512.yaml"
+        return [efficienttam_checkpoint, model_cfg]
+    elif checkpoint == "efficienttam-ti-512x512":
+        efficienttam_checkpoint = "./checkpoints/efficienttam-ti-512x512.pt"
+        model_cfg = "efficienttam_ti_512x512.yaml"
+        return [efficienttam_checkpoint, model_cfg]
+    elif checkpoint == "efficienttam-s-1":
+        efficienttam_checkpoint = "./checkpoints/efficienttam-s-1.pt"
+        model_cfg = "efficienttam-s-1.yaml"
+        return [efficienttam_checkpoint, model_cfg]
+    elif checkpoint == "efficienttam-s-2":
+        efficienttam_checkpoint = "./checkpoints/efficienttam-s-2.pt"
+        model_cfg = "efficienttam-s-2.yaml"
+        return [efficienttam_checkpoint, model_cfg]
+    elif checkpoint == "efficienttam-ti-1":
+        efficienttam_checkpoint = "./checkpoints/efficienttam-ti-1.pt"
+        model_cfg = "efficienttam-ti-1.yaml"
+        return [efficienttam_checkpoint, model_cfg]
+    elif checkpoint == "efficienttam-ti-2":
+        efficienttam_checkpoint = "./checkpoints/efficienttam-ti-2.pt"
+        model_cfg = "efficienttam-ti-2.yaml"
+        return [efficienttam_checkpoint, model_cfg]
+    else:
+        efficienttam_checkpoint = "./checkpoints/demo/efficienttam_s.pt"
+        model_cfg = "efficienttam-s.yaml"
+        return [efficienttam_checkpoint, model_cfg]
+    
+def get_mask_sam_process(
+    stored_inference_state,
+    input_first_frame_image, 
+    checkpoint, 
+    tracking_points, 
+    trackings_input_label, 
+    video_frames_dir, # extracted_frames_output_dir defined in 'preprocess_video_in' function
+    scanned_frames, 
+    working_frame: str = None, # current frame being added points
+    available_frames_to_check: List[str] = [],
+):
+
+    if len(tracking_points.value) == 0:
+        return gr.update(visible=False), None, gr.State(), None, stored_inference_state, working_frame
+    # get model and model config paths
+    print(f"USER CHOSEN CHECKPOINT: {checkpoint}")
+    sam2_checkpoint, model_cfg = load_model(checkpoint)
+    print("MODEL LOADED")
+
+    # set predictor 
+    predictor = build_sam2_video_predictor(model_cfg, sam2_checkpoint, device="cuda")
+    print("PREDICTOR READY")
+
+    # `video_dir` a directory of JPEG frames with filenames like `<frame_index>.jpg`
+    # 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 = scanned_frames
+
+    # print(f"STORED INFERENCE STEP: {stored_inference_state}")
+    if stored_inference_state is None:
+        # Init SAM2 inference_state
+        inference_state = predictor.init_state(video_path=video_dir, device="cuda")
+        print("NEW INFERENCE_STATE INITIATED")
+    else:
+        inference_state = stored_inference_state
+
+    # segment and track one object
+    # predictor.reset_state(inference_state) # if any previous tracking, reset
+
+    ### HANDLING WORKING FRAME
+    # new_working_frame = None
+    # Add new point
+    if working_frame is None:
+        ann_frame_idx = 0  # the frame index we interact with, 0 if it is the first frame
+        working_frame = "frame_0.jpg"
+    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
+            
+    print(f"NEW_WORKING_FRAME PATH: {working_frame}")
+    
+    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
+    points = np.array(tracking_points.value, dtype=np.float32)
+    # for labels, `1` means positive click and `0` means negative click
+    labels = np.array(trackings_input_label.value, np.int32)
+    _, out_obj_ids, out_mask_logits = predictor.add_new_points(
+        inference_state=inference_state,
+        frame_idx=ann_frame_idx,
+        obj_id=ann_obj_id,
+        points=points,
+        labels=labels,
+    )
+
+    # Create the plot
+    plt.figure(figsize=(12, 8))
+    plt.title(f"frame {ann_frame_idx}")
+    plt.imshow(Image.open(os.path.join(video_dir, frame_names[ann_frame_idx])))
+    show_points(points, labels, plt.gca())
+    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
+    first_frame_output_filename = "output_first_frame.jpg"
+    plt.savefig(first_frame_output_filename, format='jpg')
+    plt.close()
+    torch.cuda.empty_cache()
+
+    # Assuming available_frames_to_check.value is a list
+    if working_frame not in available_frames_to_check:
+        available_frames_to_check.append(working_frame)
+        print(available_frames_to_check)
+    
+    return gr.update(visible=True), "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(tracking_points, video_in, checkpoint, stored_inference_state, stored_frame_names, video_frames_dir, vis_frame_type, available_frames_to_check, working_frame):
+    if tracking_points is None or video_in is None or checkpoint is None or stored_inference_state is None:
+        return gr.update(value=None), gr.update(value=None), gr.update(value=None), available_frames_to_check, gr.update(visible=False)
+    #### PROPAGATION ####
+    sam2_checkpoint, model_cfg = load_model(checkpoint)
+    predictor = build_sam2_video_predictor(model_cfg, sam2_checkpoint, device="cuda")
+    
+    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)
+    
+    # Initialize a list to store file paths of saved images
+    jpeg_images = []
+
+    # run propagation throughout the video and collect the results in a dict
+    video_segments = {}  # video_segments contains the per-frame segmentation results
+    print("starting propagate_in_video")
+    for out_frame_idx, out_obj_ids, out_mask_logits in predictor.propagate_in_video(inference_state):
+        video_segments[out_frame_idx] = {
+            out_obj_id: (out_mask_logits[i] > 0.0).cpu().numpy()
+            for i, out_obj_id in enumerate(out_obj_ids)
+        }    
+    
+    # obtain the segmentation results every few frames
+    if vis_frame_type == "coarse":
+        vis_frame_stride = 15
+    elif vis_frame_type == "fine":
+        vis_frame_stride = 1
+    
+    plt.close("all")
+    for out_frame_idx in range(0, len(frame_names), vis_frame_stride):
+        plt.figure(figsize=(6, 4))
+        plt.title(f"frame {out_frame_idx}")
+        plt.imshow(Image.open(os.path.join(video_dir, frame_names[out_frame_idx])))
+        for out_obj_id, out_mask in video_segments[out_frame_idx].items():
+            show_mask(out_mask, plt.gca(), obj_id=out_obj_id)
+
+        # Define the output filename and save the figure as a JPEG file
+        output_filename = os.path.join(frames_output_dir, f"frame_{out_frame_idx}.jpg")
+        plt.savefig(output_filename, format='jpg')
+    
+        # Close the plot
+        plt.close()
+
+        # Append the file path to the list
+        jpeg_images.append(output_filename)
+
+        if f"frame_{out_frame_idx}.jpg" not in available_frames_to_check:
+            available_frames_to_check.append(f"frame_{out_frame_idx}.jpg")
+
+    torch.cuda.empty_cache()
+    print(f"JPEG_IMAGES: {jpeg_images}")
+
+    if vis_frame_type == "coarse":
+        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, gr.update(visible=True)
+    elif vis_frame_type == "fine":
+        # Create a video clip from the image sequence
+        original_fps = get_video_fps(video_in)
+        fps = original_fps  # Frames per second
+        total_frames = len(jpeg_images)
+        clip = ImageSequenceClip(jpeg_images, fps=fps)
+        # Write the result to a file
+        final_vid_output_path = "output_video.mp4"
+        
+        # Write the result to a file
+        clip.write_videofile(
+            final_vid_output_path,
+            codec='libx264'
+        )
+        
+        return gr.update(value=None), gr.update(value=final_vid_output_path), working_frame, available_frames_to_check, gr.update(visible=True)
+
+def update_ui(vis_frame_type):
+    if vis_frame_type == "coarse":
+        return gr.update(visible=True), gr.update(visible=False)
+    elif vis_frame_type == "fine":
+        return gr.update(visible=False), gr.update(visible=True)
+
+def switch_working_frame(working_frame, scanned_frames, video_frames_dir):
+    new_working_frame = None
+    if working_frame == None:
+        new_working_frame = os.path.join(video_frames_dir, scanned_frames[0])
+        
+    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
+            new_working_frame = os.path.join(video_frames_dir, scanned_frames[ann_frame_idx])
+    return gr.State([]), gr.State([]), new_working_frame, new_working_frame
+
+def reset_propagation(first_frame_path, predictor, stored_inference_state):
+    predictor.reset_state(stored_inference_state)
+    # print(f"RESET State: {stored_inference_state} ")
+    return first_frame_path, gr.State([]), gr.State([]), gr.update(value=None, visible=False), stored_inference_state, None, ["frame_0.jpg"], first_frame_path, "frame_0.jpg", gr.update(visible=False)
+    
+with gr.Blocks() as demo:
+    first_frame_path = gr.State()
+    tracking_points = gr.State([])
+    trackings_input_label = gr.State([])
+    video_frames_dir = gr.State()
+    scanned_frames = gr.State()
+    loaded_predictor = gr.State()
+    stored_inference_state = gr.State()
+    stored_frame_names = gr.State()
+    available_frames_to_check = gr.State([])
+    with gr.Column():
+        # Title
+        gr.Markdown(title)
+        with gr.Row():
+            
+            with gr.Column():
+                # Instructions
+                gr.Markdown(description_p)
+
+                # video_exp = gr.Video(label="Input Example", format="mp4", visible=False)
+                with gr.Accordion("Input Video", open=True) as video_in_drawer:
+                    video_in = gr.Video(label="Input Video", format="mp4")
+                
+                with gr.Row():
+                    point_type = gr.Radio(label="point type", choices=["include", "exclude"], value="include", scale=2)
+                    clear_points_btn = gr.Button("Clear Points", scale=1)
+                
+                input_first_frame_image = gr.Image(label="input image", interactive=False, type="filepath", visible=False)                 
+                
+                points_map = gr.Image(
+                    label="Frame with Point Prompt", 
+                    type="filepath",
+                    interactive=False
+                )
+
+                with gr.Row():
+                    checkpoint = gr.Dropdown(label="Checkpoint", choices=["efficienttam-s", "efficienttam-ti", "efficienttam-s-512x512", "efficienttam-ti-512x512", "efficienttam-s-1", "efficienttam-s-2", "efficienttam-ti-1", "efficienttam-ti-2"], value="efficienttam-s")
+                    submit_btn = gr.Button("Segment", size="lg")
+                
+            
+            with gr.Column():
+                gr.Markdown("# Try some of the examples below ⬇️")
+                gr.Examples(
+                    examples=examples,
+                    inputs=[video_in,],
+                )
+                gr.Markdown('\n\n\n\n\n\n\n\n\n\n\n')
+                gr.Markdown('\n\n\n\n\n\n\n\n\n\n\n')
+                gr.Markdown('\n\n\n\n\n\n\n\n\n\n\n')
+                with gr.Row():
+                    working_frame = gr.Dropdown(label="Frame ID", choices=[""], value=None, visible=False, allow_custom_value=False, interactive=True)
+                    change_current = gr.Button("change current", visible=False)
+                output_result = gr.Image(label="Reference Mask")
+                with gr.Row():
+                    vis_frame_type = gr.Radio(label="Track level", choices=["coarse", "fine"], value="coarse", scale=2)
+                    propagate_btn = gr.Button("Track", scale=1)
+                reset_prpgt_brn = gr.Button("Reset", visible=False)
+                output_propagated = gr.Gallery(label="Masklets", columns=4, visible=False)
+                output_video = gr.Video(visible=False)
+    
+    
+
+    # When new video is uploaded
+    video_in.upload(
+        fn = preprocess_video_in, 
+        inputs = [video_in], 
+        outputs = [
+            first_frame_path, 
+            tracking_points, # update Tracking Points in the gr.State([]) object
+            trackings_input_label, # update Tracking Labels in the gr.State([]) object
+            input_first_frame_image, # hidden component used as ref when clearing points
+            points_map, # Image component where we add new tracking points
+            video_frames_dir, # Array where frames from video_in are deep stored
+            scanned_frames, # Scanned frames by EfficientTAM
+            stored_inference_state, # EfficientTAM inference state
+            stored_frame_names, # 
+            video_in_drawer, # Accordion to hide uploaded video player
+        ],
+        queue = False
+    )
+
+    video_in.change(
+        fn = preprocess_video_in, 
+        inputs = [video_in], 
+        outputs = [
+            first_frame_path, 
+            tracking_points, # update Tracking Points in the gr.State([]) object
+            trackings_input_label, # update Tracking Labels in the gr.State([]) object
+            input_first_frame_image, # hidden component used as ref when clearing points
+            points_map, # Image component where we add new tracking points
+            video_frames_dir, # Array where frames from video_in are deep stored
+            scanned_frames, # Scanned frames by EfficientTAM
+            stored_inference_state, # EfficientTAM inference state
+            stored_frame_names, # 
+            video_in_drawer, # Accordion to hide uploaded video player
+        ],
+        queue = False
+    )
+
+    
+    # triggered when we click on image to add new points
+    points_map.select(
+        fn = get_point, 
+        inputs = [
+            point_type, # "include" or "exclude"
+            tracking_points, # get tracking_points values
+            trackings_input_label, # get tracking label values
+            input_first_frame_image, # gr.State() first frame path
+        ], 
+        outputs = [
+            tracking_points, # updated with new points
+            trackings_input_label, # updated with corresponding labels
+            points_map, # updated image with points
+        ], 
+        queue = False
+    )
+
+    # Clear every points clicked and added to the map
+    clear_points_btn.click(
+        fn = clear_points,
+        inputs = input_first_frame_image, # we get the untouched hidden image
+        outputs = [
+            first_frame_path, 
+            tracking_points, 
+            trackings_input_label, 
+            points_map, 
+        ],
+        queue=False
+    )
+
+    
+    change_current.click(
+        fn = switch_working_frame,
+        inputs = [working_frame, scanned_frames, video_frames_dir],
+        outputs = [tracking_points, trackings_input_label, input_first_frame_image, points_map],
+        queue=False
+    )
+    
+
+    submit_btn.click(
+        fn = get_mask_sam_process,
+        inputs = [
+            stored_inference_state,
+            input_first_frame_image, 
+            checkpoint, 
+            tracking_points, 
+            trackings_input_label, 
+            video_frames_dir, 
+            scanned_frames, 
+            working_frame,
+            available_frames_to_check,
+        ],
+        outputs = [
+            change_current,
+            output_result, 
+            stored_frame_names, 
+            loaded_predictor,
+            stored_inference_state,
+            working_frame,
+        ],
+        concurrency_limit=10,
+        queue=False
+    )
+
+    reset_prpgt_brn.click(
+        fn = reset_propagation,
+        inputs = [first_frame_path, loaded_predictor, stored_inference_state],
+        outputs = [points_map, tracking_points, trackings_input_label, output_propagated, stored_inference_state, output_result, available_frames_to_check, input_first_frame_image, working_frame, reset_prpgt_brn],
+        queue=False
+    )
+
+    propagate_btn.click(
+        fn = update_ui,
+        inputs = [vis_frame_type],
+        outputs = [output_propagated, output_video],
+        queue=False
+    ).then(
+        fn = propagate_to_all,
+        inputs = [tracking_points, 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, reset_prpgt_brn],
+        concurrency_limit=10,
+        queue=False
+    )
+
+demo.queue()
+demo.launch()

requirements.txt

@@ -0,0 +1,204 @@
+# This file may be used to create an environment using:
+# $ conda create --name <env> --file <this file>
+# platform: linux-64
+_libgcc_mutex=0.1=main
+_openmp_mutex=5.1=1_gnu
+aiofiles=23.2.1=pypi_0
+aiohappyeyeballs=2.4.3=pypi_0
+aiohttp=3.11.7=pypi_0
+aiosignal=1.3.1=pypi_0
+altair=5.5.0=pypi_0
+annotated-types=0.7.0=pypi_0
+antlr4-python3-runtime=4.9.3=pypi_0
+anyio=3.7.1=pypi_0
+aom=3.6.0=h6a678d5_0
+attrs=24.2.0=pypi_0
+bcrypt=4.2.1=pypi_0
+blas=1.0=mkl
+brotli-python=1.0.9=py312h6a678d5_8
+build=1.2.2.post1=pypi_0
+bzip2=1.0.8=h5eee18b_6
+ca-certificates=2024.8.30=hbcca054_0
+cairo=1.16.0=hb05425b_5
+certifi=2024.8.30=pyhd8ed1ab_0
+cffi=1.17.1=pypi_0
+charset-normalizer=3.3.2=pyhd3eb1b0_0
+click=8.1.7=pypi_0
+cmake=3.31.1=pypi_0
+contourpy=1.3.1=pypi_0
+cryptography=43.0.3=pypi_0
+cycler=0.12.1=pypi_0
+dav1d=1.2.1=h5eee18b_0
+decorator=5.1.1=pyhd3eb1b0_0
+defusedxml=0.7.1=pypi_0
+einops=0.8.0=pypi_0
+expat=2.6.2=h6a678d5_0
+fastapi=0.115.5=pypi_0
+ffmpeg=4.2.2=h167e202_0
+ffmpy=0.4.0=pypi_0
+filelock=3.16.1=pypi_0
+flash-attn=2.7.0.post2=pypi_0
+fontconfig=2.14.1=h4c34cd2_2
+fonttools=4.55.0=pypi_0
+freetype=2.12.1=h4a9f257_0
+frozenlist=1.5.0=pypi_0
+fsspec=2024.10.0=pypi_0
+giflib=5.2.1=h5eee18b_3
+glib=2.78.4=h6a678d5_0
+glib-tools=2.78.4=h6a678d5_0
+gmp=6.2.1=h58526e2_0
+gnutls=3.6.13=h85f3911_1
+gradio=4.44.0=pypi_0
+gradio-client=1.3.0=pypi_0
+gradio-image-prompter=0.1.0=pypi_0
+graphite2=1.3.14=h295c915_1
+h11=0.14.0=pypi_0
+harfbuzz=4.3.0=hf52aaf7_2
+httpcore=1.0.7=pypi_0
+httpx=0.27.0=pypi_0
+huggingface-hub=0.26.2=pypi_0
+hydra-core=1.3.2=pypi_0
+icu=73.1=h6a678d5_0
+idna=3.7=py312h06a4308_0
+imageio=2.9.0=pyhd3eb1b0_0
+imageio-ffmpeg=0.5.1=pyhd8ed1ab_0
+importlib-resources=6.4.5=pypi_0
+iniconfig=2.0.0=pypi_0
+intel-openmp=2023.1.0=hdb19cb5_46306
+iopath=0.1.10=pypi_0
+jinja2=3.1.4=pypi_0
+jpeg=9e=h5eee18b_3
+jsonschema=4.23.0=pypi_0
+jsonschema-specifications=2024.10.1=pypi_0
+kiwisolver=1.4.7=pypi_0
+lame=3.100=h7b6447c_0
+lcms2=2.12=h3be6417_0
+ld_impl_linux-64=2.38=h1181459_1
+leptonica=1.82.0=h42c8aad_2
+lerc=3.0=h295c915_0
+libarchive=3.6.2=h6ac8c49_3
+libdeflate=1.17=h5eee18b_1
+libffi=3.4.4=h6a678d5_1
+libgcc-ng=11.2.0=h1234567_1
+libglib=2.78.4=hdc74915_0
+libgomp=11.2.0=h1234567_1
+libiconv=1.16=h5eee18b_3
+libogg=1.3.5=h27cfd23_1
+libopus=1.3.1=h7b6447c_0
+libpng=1.6.39=h5eee18b_0
+libstdcxx-ng=11.2.0=h1234567_1
+libtheora=1.1.1=h7f8727e_3
+libtiff=4.5.1=h6a678d5_0
+libuuid=1.41.5=h5eee18b_0
+libvorbis=1.3.7=h7b6447c_0
+libvpx=1.13.1=h6a678d5_0
+libwebp=1.3.2=h11a3e52_0
+libwebp-base=1.3.2=h5eee18b_0
+libxcb=1.15=h7f8727e_0
+libxml2=2.10.4=hfdd30dd_2
+linkify-it-py=2.0.3=pypi_0
+lz4-c=1.9.4=h6a678d5_1
+markdown-it-py=2.2.0=pypi_0
+markupsafe=2.1.5=pypi_0
+matplotlib=3.9.2=pypi_0
+mdit-py-plugins=0.3.3=pypi_0
+mdurl=0.1.2=pypi_0
+mkl=2023.1.0=h213fc3f_46344
+mkl-service=2.4.0=py312h5eee18b_1
+mkl_fft=1.3.10=py312h5eee18b_0
+mkl_random=1.2.7=py312h526ad5a_0
+moviepy=1.0.3=pyhd8ed1ab_1
+mpmath=1.3.0=pypi_0
+multidict=6.1.0=pypi_0
+narwhals=1.14.2=pypi_0
+ncurses=6.4=h6a678d5_0
+nettle=3.6=he412f7d_0
+networkx=3.4.2=pypi_0
+ninja=1.11.1.2=pypi_0
+numpy=1.26.4=pypi_0
+nvidia-cublas-cu12=12.4.5.8=pypi_0
+nvidia-cuda-cupti-cu12=12.4.127=pypi_0
+nvidia-cuda-nvrtc-cu12=12.4.127=pypi_0
+nvidia-cuda-runtime-cu12=12.4.127=pypi_0
+nvidia-cudnn-cu12=9.1.0.70=pypi_0
+nvidia-cufft-cu12=11.2.1.3=pypi_0
+nvidia-curand-cu12=10.3.5.147=pypi_0
+nvidia-cusolver-cu12=11.6.1.9=pypi_0
+nvidia-cusparse-cu12=12.3.1.170=pypi_0
+nvidia-nccl-cu12=2.21.5=pypi_0
+nvidia-nvjitlink-cu12=12.4.127=pypi_0
+nvidia-nvtx-cu12=12.4.127=pypi_0
+omegaconf=2.3.0=pypi_0
+opencv-python=4.10.0.84=pypi_0
+openh264=2.1.1=h4ff587b_0
+openjpeg=2.5.2=he7f1fd0_0
+openssl=3.0.15=h5eee18b_0
+orjson=3.10.12=pypi_0
+packaging=24.2=pypi_0
+pandas=2.2.3=pypi_0
+paramiko=3.5.0=pypi_0
+pcre2=10.42=hebb0a14_1
+pillow=10.4.0=py312h5eee18b_0
+pip=24.0=py312h06a4308_0
+pixman=0.40.0=h7f8727e_1
+pluggy=1.5.0=pypi_0
+portalocker=3.0.0=pypi_0
+proglog=0.1.9=py_0
+propcache=0.2.0=pypi_0
+psutil=5.9.8=pypi_0
+pycparser=2.22=pypi_0
+pycryptodome=3.21.0=pypi_0
+pydantic=2.10.2=pypi_0
+pydantic-core=2.27.1=pypi_0
+pydub=0.25.1=pypi_0
+pygments=2.18.0=pypi_0
+pynacl=1.5.0=pypi_0
+pyparsing=3.2.0=pypi_0
+pyproject-hooks=1.2.0=pypi_0
+pysocks=1.7.1=py312h06a4308_0
+pytest=8.3.3=pypi_0
+python=3.12.4=h5148396_1
+python-dateutil=2.9.0.post0=pypi_0
+python-multipart=0.0.12=pypi_0
+pytz=2024.2=pypi_0
+pyyaml=6.0.2=pypi_0
+readline=8.2=h5eee18b_0
+referencing=0.35.1=pypi_0
+requests=2.32.3=py312h06a4308_0
+rich=13.9.4=pypi_0
+rpds-py=0.21.0=pypi_0
+ruff=0.8.0=pypi_0
+safehttpx=0.1.1=pypi_0
+scipy=1.14.1=pypi_0
+semantic-version=2.10.0=pypi_0
+setuptools=72.1.0=py312h06a4308_0
+shellingham=1.5.4=pypi_0
+six=1.16.0=pypi_0
+sniffio=1.3.1=pypi_0
+spaces=0.30.4=pypi_0
+sqlite=3.45.3=h5eee18b_0
+starlette=0.41.3=pypi_0
+supervision=0.25.0=pypi_0
+sympy=1.13.1=pypi_0
+tbb=2021.8.0=hdb19cb5_0
+tesseract=5.2.0=h6a678d5_0
+tk=8.6.14=h39e8969_0
+tomlkit=0.12.0=pypi_0
+torch=2.5.1+cu124=pypi_0
+torchaudio=2.5.1+cu124=pypi_0
+torchvision=0.20.1+cu124=pypi_0
+tqdm=4.67.1=pypi_0
+triton=3.1.0=pypi_0
+typer=0.13.1=pypi_0
+typing-extensions=4.12.2=pypi_0
+tzdata=2024.2=pypi_0
+uc-micro-py=1.0.3=pypi_0
+urllib3=2.2.2=py312h06a4308_0
+uvicorn=0.32.1=pypi_0
+websockets=11.0.3=pypi_0
+wheel=0.43.0=py312h06a4308_0
+x264=1!152.20180806=h14c3975_0
+xz=5.4.6=h5eee18b_1
+yarl=1.18.0=pypi_0
+zlib=1.2.13=h5eee18b_1
+zstd=1.5.5=hc292b87_2