Work in process

Dockerfile

@@ -18,11 +18,11 @@ RUN add-apt-repository -y -r ppa:jonathonf/ffmpeg-4 \
     && rm -rf /var/lib/apt/lists/*
 WORKDIR /app
 RUN pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu124
-RUN pip3 install numpy matplotlib pillow gradio==3.38.0 opencv-python ffmpeg-python
+RUN pip3 install numpy matplotlib pillow gradio==3.38.0 opencv-python ffmpeg-python moviepy
 RUN git clone https://github.com/facebookresearch/segment-anything-2.git
 WORKDIR /app/segment-anything-2
-RUN pip3 install -e .
-# RUN pip3 install -e ".[demo]"
+# RUN pip3 install -e .
+RUN pip3 install -e ".[demo]"
 WORKDIR /app/segment-anything-2/checkpoints
 RUN ./download_ckpts.sh
 WORKDIR /app

app.py

@@ -11,245 +11,17 @@
 #     print("Command failed with return code:", result.returncode)
 import gc
 import math
+# import multiprocessing as mp
+import torch.multiprocessing as mp
 import os
+from process_wrappers import clear_folder, draw_markers, sam_click_wrapper1, sam_stroke_process, tracking_objects_process
 os.environ["TORCH_CUDNN_SDPA_ENABLED"] = "1"
-import shutil
 import ffmpeg
-import zipfile
-import gradio as gr
-import torch
-import numpy as np
-import matplotlib.pyplot as plt
-from PIL import Image
-from sam2.build_sam import build_sam2
-from sam2.sam2_image_predictor import SAM2ImagePredictor
-from sam2.build_sam import build_sam2_video_predictor
 import cv2
 
-def clean(Seg_Tracker):
-    if Seg_Tracker is not None:
-        predictor, inference_state, image_predictor = Seg_Tracker
-        predictor.reset_state(inference_state)
-        del predictor
-        del inference_state
-        del image_predictor
-        del Seg_Tracker
-        gc.collect()
-        torch.cuda.empty_cache()
-    return None, ({}, {}), None, None, 0, None, None, None, 0
 
-def change_video(input_video):
-    if input_video is None:
-        return 0, 0
-    cap = cv2.VideoCapture(input_video)
-    fps = cap.get(cv2.CAP_PROP_FPS)
-    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
-    cap.release()
-    scale_slider = gr.Slider.update(minimum=1.0,
-                                maximum=fps,
-                                step=1.0,
-                                value=fps,)
-    frame_per = gr.Slider.update(minimum= 0.0,
-                            maximum= total_frames / fps,
-                            step=1.0/fps,
-                            value=0.0,)
-    return scale_slider, frame_per
-
-def get_meta_from_video(Seg_Tracker, input_video, scale_slider, checkpoint):
-
-    output_dir = '/tmp/output_frames'
-    output_masks_dir = '/tmp/output_masks'
-    output_combined_dir = '/tmp/output_combined'
-    clear_folder(output_dir)
-    clear_folder(output_masks_dir)
-    clear_folder(output_combined_dir)
-    if input_video is None:
-        return None, ({}, {}), None, None, 0, None, None, None, 0
-    cap = cv2.VideoCapture(input_video)
-    fps = cap.get(cv2.CAP_PROP_FPS)
-    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
-    cap.release()
-    frame_interval = max(1, int(fps // scale_slider))
-    print(f"frame_interval: {frame_interval}")
-    try:
-        ffmpeg.input(input_video, hwaccel='cuda').output(
-            os.path.join(output_dir, '%07d.jpg'), q=2, start_number=0, 
-            vf=rf'select=not(mod(n\,{frame_interval}))', vsync='vfr'
-        ).run()
-    except:
-        print(f"ffmpeg cuda err")
-        ffmpeg.input(input_video).output(
-            os.path.join(output_dir, '%07d.jpg'), q=2, start_number=0, 
-            vf=rf'select=not(mod(n\,{frame_interval}))', vsync='vfr'
-        ).run()
-
-    first_frame_path = os.path.join(output_dir, '0000000.jpg')
-    first_frame = cv2.imread(first_frame_path)
-    first_frame_rgb = cv2.cvtColor(first_frame, cv2.COLOR_BGR2RGB)
-
-    if Seg_Tracker is not None:
-        del Seg_Tracker
-        Seg_Tracker = None
-        gc.collect()
-        torch.cuda.empty_cache()
-    torch.autocast(device_type="cuda", dtype=torch.bfloat16).__enter__()
-    if torch.cuda.get_device_properties(0).major >= 8:
-        torch.backends.cuda.matmul.allow_tf32 = True
-        torch.backends.cudnn.allow_tf32 = True
-
-    if checkpoint == "tiny":
-        sam2_checkpoint = "segment-anything-2/checkpoints/sam2_hiera_tiny.pt"
-        model_cfg = "sam2_hiera_t.yaml"
-    elif checkpoint == "samll":
-        sam2_checkpoint = "segment-anything-2/checkpoints/sam2_hiera_small.pt"
-        model_cfg = "sam2_hiera_s.yaml"
-    elif checkpoint == "base-plus":
-        sam2_checkpoint = "segment-anything-2/checkpoints/sam2_hiera_base_plus.pt"
-        model_cfg = "sam2_hiera_b+.yaml"
-    elif checkpoint == "large":
-        sam2_checkpoint = "segment-anything-2/checkpoints/sam2_hiera_large.pt"
-        model_cfg = "sam2_hiera_l.yaml"
-    
-    predictor = build_sam2_video_predictor(model_cfg, sam2_checkpoint, device="cuda")
-    sam2_model = build_sam2(model_cfg, sam2_checkpoint, device="cuda")
-
-    image_predictor = SAM2ImagePredictor(sam2_model)
-    inference_state = predictor.init_state(video_path=output_dir)
-    predictor.reset_state(inference_state)
-    frame_per = gr.Slider.update(minimum= 0.0,
-                            maximum= total_frames / fps,
-                            step=frame_interval / fps,
-                            value=0.0,)
-    return (predictor, inference_state, image_predictor), ({}, {}), first_frame_rgb, first_frame_rgb, frame_per, None, None, None, 0
-
-def mask2bbox(mask):
-    if len(np.where(mask > 0)[0]) == 0:
-        print(f'not mask')
-        return np.array([0, 0, 0, 0]).astype(np.int64), False
-    x_ = np.sum(mask, axis=0)
-    y_ = np.sum(mask, axis=1)
-    x0 = np.min(np.nonzero(x_)[0])
-    x1 = np.max(np.nonzero(x_)[0])
-    y0 = np.min(np.nonzero(y_)[0])
-    y1 = np.max(np.nonzero(y_)[0])
-    return np.array([x0, y0, x1, y1]).astype(np.int64), True
-
-def sam_stroke(Seg_Tracker, drawing_board, last_draw, frame_num, ann_obj_id):
-    predictor, inference_state, image_predictor = Seg_Tracker
-    image_path = f'/tmp/output_frames/{frame_num:07d}.jpg'
-    image = cv2.imread(image_path)
-    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
-    display_image = drawing_board["image"]
-    image_predictor.set_image(image)
-    input_mask = drawing_board["mask"]
-    input_mask[input_mask != 0] = 255
-    if last_draw is not None:
-        diff_mask = cv2.absdiff(input_mask, last_draw)
-        input_mask = diff_mask
-    bbox, hasMask = mask2bbox(input_mask[:, :, 0]) 
-    if not hasMask :
-        return Seg_Tracker, display_image, display_image
-    masks, scores, logits = image_predictor.predict( point_coords=None, point_labels=None, box=bbox[None, :], multimask_output=False,)
-    mask = masks > 0.0
-    masked_frame = show_mask(mask, display_image, ann_obj_id)
-    masked_with_rect = draw_rect(masked_frame, bbox, ann_obj_id)
-    frame_idx, object_ids, masks = predictor.add_new_mask(inference_state, frame_idx=frame_num, obj_id=ann_obj_id, mask=mask[0])
-    last_draw = drawing_board["mask"]
-    return Seg_Tracker, masked_with_rect, masked_with_rect, last_draw
-
-def draw_rect(image, bbox, obj_id):
-    cmap = plt.get_cmap("tab10")
-    color = np.array(cmap(obj_id)[:3])
-    rgb_color = tuple(map(int, (color[:3] * 255).astype(np.uint8)))
-    inv_color = tuple(map(int, (255 - color[:3] * 255).astype(np.uint8)))
-    x0, y0, x1, y1 = bbox
-    image_with_rect = cv2.rectangle(image.copy(), (x0, y0), (x1, y1), rgb_color, thickness=2)
-    return image_with_rect
-
-def sam_click(Seg_Tracker, frame_num, point_mode, click_stack, ann_obj_id, evt: gr.SelectData):
-    points_dict, labels_dict = click_stack
-    predictor, inference_state, image_predictor = Seg_Tracker
-    ann_frame_idx = frame_num  # the frame index we interact with
-    print(f'ann_frame_idx: {ann_frame_idx}')
-    point = np.array([[evt.index[0], evt.index[1]]], dtype=np.float32)
-    if point_mode == "Positive":
-        label = np.array([1], np.int32)
-    else:
-        label = np.array([0], np.int32)
-
-    if ann_frame_idx not in points_dict:
-        points_dict[ann_frame_idx] = {}
-    if ann_frame_idx not in labels_dict:
-        labels_dict[ann_frame_idx] = {}
-
-    if ann_obj_id not in points_dict[ann_frame_idx]:
-        points_dict[ann_frame_idx][ann_obj_id] = np.empty((0, 2), dtype=np.float32)
-    if ann_obj_id not in labels_dict[ann_frame_idx]:
-        labels_dict[ann_frame_idx][ann_obj_id] = np.empty((0,), dtype=np.int32)
-
-    points_dict[ann_frame_idx][ann_obj_id] = np.append(points_dict[ann_frame_idx][ann_obj_id], point, axis=0)
-    labels_dict[ann_frame_idx][ann_obj_id] = np.append(labels_dict[ann_frame_idx][ann_obj_id], label, axis=0)
-
-    click_stack = (points_dict, labels_dict)
-
-    frame_idx, 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_dict[ann_frame_idx][ann_obj_id],
-        labels=labels_dict[ann_frame_idx][ann_obj_id],
-    )
-
-    image_path = f'/tmp/output_frames/{ann_frame_idx:07d}.jpg'
-    image = cv2.imread(image_path)
-    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
-
-    masked_frame = image.copy()
-    for i, obj_id in enumerate(out_obj_ids):
-        mask = (out_mask_logits[i] > 0.0).cpu().numpy()
-        masked_frame = show_mask(mask, image=masked_frame, obj_id=obj_id)
-    masked_frame_with_markers = draw_markers(masked_frame, points_dict[ann_frame_idx], labels_dict[ann_frame_idx])
-
-    return Seg_Tracker, masked_frame_with_markers, masked_frame_with_markers, click_stack
-
-def draw_markers(image, points_dict, labels_dict):
-    cmap = plt.get_cmap("tab10")
-    image_h, image_w = image.shape[:2]
-    marker_size = max(1, int(min(image_h, image_w) * 0.05))
-
-    for obj_id in points_dict:
-        color = np.array(cmap(obj_id)[:3])
-        rgb_color = tuple(map(int, (color[:3] * 255).astype(np.uint8)))
-        inv_color = tuple(map(int, (255 - color[:3] * 255).astype(np.uint8)))
-        for point, label in zip(points_dict[obj_id], labels_dict[obj_id]):
-            x, y = int(point[0]), int(point[1])
-            if label == 1:
-                cv2.drawMarker(image, (x, y), inv_color, markerType=cv2.MARKER_CROSS, markerSize=marker_size, thickness=2)
-            else:
-                cv2.drawMarker(image, (x, y), inv_color, markerType=cv2.MARKER_TILTED_CROSS, markerSize=int(marker_size / np.sqrt(2)), thickness=2)
-    
-    return image
-
-def show_mask(mask, image=None, obj_id=None):
-    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)
-    mask_image = (mask_image * 255).astype(np.uint8)
-    if image is not None:
-        image_h, image_w = image.shape[:2]
-        if (image_h, image_w) != (h, w):
-            raise ValueError(f"Image dimensions ({image_h}, {image_w}) and mask dimensions ({h}, {w}) do not match")
-        colored_mask = np.zeros_like(image, dtype=np.uint8)
-        for c in range(3):
-            colored_mask[..., c] = mask_image[..., c]
-        alpha_mask = mask_image[..., 3] / 255.0
-        for c in range(3):
-            image[..., c] = np.where(alpha_mask > 0, (1 - alpha_mask) * image[..., c] + alpha_mask * colored_mask[..., c], image[..., c])
-        return image
-    return mask_image
+def clean():
+    return ({}, {}, {}), None, None, 0, None, None, None, 0
 
 def show_res_by_slider(frame_per, click_stack):
     image_path = '/tmp/output_frames'
@@ -274,85 +46,11 @@ def show_res_by_slider(frame_per, click_stack):
         print(f"{chosen_frame_path}")
         chosen_frame_show = cv2.imread(chosen_frame_path)
         chosen_frame_show = cv2.cvtColor(chosen_frame_show, cv2.COLOR_BGR2RGB)
-        points_dict, labels_dict = click_stack
+        points_dict, labels_dict, masks_dict = click_stack
         if frame_num in points_dict and frame_num in labels_dict:
             chosen_frame_show = draw_markers(chosen_frame_show, points_dict[frame_num], labels_dict[frame_num])
         return chosen_frame_show, chosen_frame_show, frame_num
 
-def clear_folder(folder_path):
-    if os.path.exists(folder_path):
-        shutil.rmtree(folder_path)
-    os.makedirs(folder_path)
-
-def zip_folder(folder_path, output_zip_path):
-    with zipfile.ZipFile(output_zip_path, 'w', zipfile.ZIP_STORED) as zipf:
-        for root, _, files in os.walk(folder_path):
-            for file in files:
-                file_path = os.path.join(root, file)
-                zipf.write(file_path, os.path.relpath(file_path, folder_path))
-
-def tracking_objects(Seg_Tracker, frame_num, input_video):
-    output_dir = '/tmp/output_frames'
-    output_masks_dir = '/tmp/output_masks'
-    output_combined_dir = '/tmp/output_combined'
-    output_video_path = '/tmp/output_video.mp4'
-    output_zip_path = '/tmp/output_masks.zip'
-    clear_folder(output_masks_dir)
-    clear_folder(output_combined_dir)
-    if os.path.exists(output_video_path):
-        os.remove(output_video_path)
-    if os.path.exists(output_zip_path):
-        os.remove(output_zip_path)
-    video_segments = {}
-    predictor, inference_state, image_predictor = Seg_Tracker
-    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)
-        }
-    frame_files = sorted([f for f in os.listdir(output_dir) if f.endswith('.jpg')])
-    # for frame_idx in sorted(video_segments.keys()):
-    for frame_file in frame_files:
-        frame_idx = int(os.path.splitext(frame_file)[0])
-        frame_path = os.path.join(output_dir, frame_file)
-        image = cv2.imread(frame_path)
-        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
-        masked_frame = image.copy()
-        if frame_idx in video_segments:
-            for obj_id, mask in video_segments[frame_idx].items():
-                masked_frame = show_mask(mask, image=masked_frame, obj_id=obj_id)
-                mask_output_path = os.path.join(output_masks_dir, f'{obj_id}_{frame_idx:07d}.png')
-                cv2.imwrite(mask_output_path, show_mask(mask))
-        combined_output_path = os.path.join(output_combined_dir, f'{frame_idx:07d}.png')
-        combined_image_bgr = cv2.cvtColor(masked_frame, cv2.COLOR_RGB2BGR)
-        cv2.imwrite(combined_output_path, combined_image_bgr)
-        if frame_idx == frame_num:
-            final_masked_frame = masked_frame
-
-    cap = cv2.VideoCapture(input_video)
-    fps = cap.get(cv2.CAP_PROP_FPS)
-    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
-    frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
-    frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
-    cap.release()
-    # output_frames = int(total_frames * scale_slider)
-    output_frames = len([name for name in os.listdir(output_combined_dir) if os.path.isfile(os.path.join(output_combined_dir, name)) and name.endswith('.png')])
-    out_fps = fps * output_frames / total_frames
-    # ffmpeg.input(os.path.join(output_combined_dir, '%07d.png'), framerate=out_fps).output(output_video_path, vcodec='h264_nvenc', pix_fmt='yuv420p').run()
-    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
-    out = cv2.VideoWriter(output_video_path, fourcc, out_fps, (frame_width, frame_height))
-
-    for i in range(output_frames):
-        frame_path = os.path.join(output_combined_dir, f'{i:07d}.png')
-        frame = cv2.imread(frame_path)
-        out.write(frame)
-
-    out.release()
-
-    zip_folder(output_masks_dir, output_zip_path)
-    print("done")
-    return final_masked_frame, final_masked_frame, output_video_path, output_video_path, output_zip_path
-
 def increment_ann_obj_id(ann_obj_id):
     ann_obj_id += 1
     return ann_obj_id
@@ -360,7 +58,87 @@ def increment_ann_obj_id(ann_obj_id):
 def drawing_board_get_input_first_frame(input_first_frame):
     return input_first_frame
 
+def sam_stroke_wrapper(click_stack, checkpoint, drawing_board, last_draw, frame_num, ann_obj_id):
+    queue = mp.Queue()
+    p = mp.Process(target=sam_stroke_process, args=(queue, click_stack, checkpoint, drawing_board, last_draw, frame_num, ann_obj_id))
+    p.start()
+    error, result = queue.get()
+    p.join()
+    if error:
+        raise Exception(f"Error in sam_stroke_process: {error}")
+    return result
+
+def tracking_objects_wrapper(click_stack, checkpoint, frame_num, input_video):
+    queue = mp.Queue()
+    p = mp.Process(target=tracking_objects_process, args=(queue, click_stack, checkpoint, frame_num, input_video))
+    p.start()
+    error, result = queue.get()
+    p.join()
+    if error:
+        raise Exception(f"Error in sam_stroke_process: {error}")
+    return result
+
 def seg_track_app():
+    import gradio as gr
+
+    def sam_click_wrapper(checkpoint, frame_num, point_mode, click_stack, ann_obj_id, evt: gr.SelectData):
+        return sam_click_wrapper1(checkpoint, frame_num, point_mode, click_stack, ann_obj_id, [evt.index[0], evt.index[1]])
+    
+    def change_video(input_video):
+        import gradio as gr
+        if input_video is None:
+            return 0, 0
+        cap = cv2.VideoCapture(input_video)
+        fps = cap.get(cv2.CAP_PROP_FPS)
+        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+        cap.release()
+        scale_slider = gr.Slider.update(minimum=1.0,
+                                    maximum=fps,
+                                    step=1.0,
+                                    value=fps,)
+        frame_per = gr.Slider.update(minimum= 0.0,
+                                maximum= total_frames / fps,
+                                step=1.0/fps,
+                                value=0.0,)
+        return scale_slider, frame_per
+
+    def get_meta_from_video(input_video, scale_slider):
+        import gradio as gr
+        output_dir = '/tmp/output_frames'
+        output_masks_dir = '/tmp/output_masks'
+        output_combined_dir = '/tmp/`output_combined`'
+        clear_folder(output_dir)
+        clear_folder(output_masks_dir)
+        clear_folder(output_combined_dir)
+        if input_video is None:
+            return ({}, {}, {}), None, None, 0, None, None, None, 0
+        cap = cv2.VideoCapture(input_video)
+        fps = cap.get(cv2.CAP_PROP_FPS)
+        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+        cap.release()
+        frame_interval = max(1, int(fps // scale_slider))
+        print(f"frame_interval: {frame_interval}")
+        try:
+            ffmpeg.input(input_video, hwaccel='cuda').output(
+                os.path.join(output_dir, '%07d.jpg'), q=2, start_number=0, 
+                vf=rf'select=not(mod(n\,{frame_interval}))', vsync='vfr'
+            ).run()
+        except:
+            print(f"ffmpeg cuda err")
+            ffmpeg.input(input_video).output(
+                os.path.join(output_dir, '%07d.jpg'), q=2, start_number=0, 
+                vf=rf'select=not(mod(n\,{frame_interval}))', vsync='vfr'
+            ).run()
+
+        first_frame_path = os.path.join(output_dir, '0000000.jpg')
+        first_frame = cv2.imread(first_frame_path)
+        first_frame_rgb = cv2.cvtColor(first_frame, cv2.COLOR_BGR2RGB)
+
+        frame_per = gr.Slider.update(minimum= 0.0,
+                                maximum= total_frames / fps,
+                                step=frame_interval / fps,
+                                value=0.0,)
+        return ({}, {}, {}), first_frame_rgb, first_frame_rgb, frame_per, None, None, None, 0
 
     ##########################################################
     ######################  Front-end ########################
@@ -438,8 +216,7 @@ def seg_track_app():
             '''
         )
 
-        click_stack = gr.State(({}, {}))
-        Seg_Tracker = gr.State(None)
+        click_stack = gr.State(({}, {}, {}))
         frame_num = gr.State(value=(int(0)))
         ann_obj_id = gr.State(value=(int(0)))
         last_draw = gr.State(None)
@@ -474,7 +251,7 @@ def seg_track_app():
                             
                     tab_click = gr.Tab(label="Point Prompt")
                     with tab_click:
-                        input_first_frame = gr.Image(label='Segment result of first frame',interactive=True).style(height=550)
+                        input_first_frame = gr.Image(label='Segment result of first frame',interactive=True, height=550)
                         with gr.Row():
                             point_mode = gr.Radio(
                                         choices=["Positive",  "Negative"],
@@ -549,18 +326,16 @@ def seg_track_app():
         preprocess_button.click(
             fn=get_meta_from_video,
             inputs=[
-                Seg_Tracker,
                 input_video,
                 scale_slider,
-                checkpoint
             ],
             outputs=[
-                Seg_Tracker, click_stack, input_first_frame, drawing_board, frame_per, output_video, output_mp4, output_mask, ann_obj_id
+                click_stack, input_first_frame, drawing_board, frame_per, output_video, output_mp4, output_mask, ann_obj_id
             ]
         )
 
         frame_per.release(
-            fn=show_res_by_slider, 
+            fn= show_res_by_slider, 
             inputs=[
                 frame_per, click_stack
                 ], 
@@ -571,20 +346,21 @@ def seg_track_app():
 
         # Interactively modify the mask acc click
         input_first_frame.select(
-            fn=sam_click,
+            fn=sam_click_wrapper,
             inputs=[
-                Seg_Tracker, frame_num, point_mode, click_stack, ann_obj_id
+                checkpoint, frame_num, point_mode, click_stack, ann_obj_id
             ],
             outputs=[
-                Seg_Tracker, input_first_frame, drawing_board, click_stack
+                input_first_frame, drawing_board, click_stack
             ]
         )
 
         # Track object in video
         track_for_video.click(
-            fn=tracking_objects,
+            fn=tracking_objects_wrapper,
             inputs=[
-                Seg_Tracker,
+                click_stack,
+                checkpoint,
                 frame_num,
                 input_video,
             ],
@@ -599,11 +375,9 @@ def seg_track_app():
 
         reset_button.click(
             fn=clean,
-            inputs=[
-                Seg_Tracker
-            ],
+            inputs=[],
             outputs=[
-                Seg_Tracker, click_stack, input_first_frame, drawing_board, frame_per, output_video, output_mp4, output_mask, ann_obj_id
+                click_stack, input_first_frame, drawing_board, frame_per, output_video, output_mp4, output_mask, ann_obj_id
             ]
         )
 
@@ -624,12 +398,12 @@ def seg_track_app():
         )
 
         seg_acc_stroke.click(
-            fn=sam_stroke,
+            fn=sam_stroke_wrapper,
             inputs=[
-                Seg_Tracker, drawing_board, last_draw, frame_num, ann_obj_id
+                click_stack, checkpoint, drawing_board, last_draw, frame_num, ann_obj_id
             ],
             outputs=[
-                Seg_Tracker, input_first_frame, drawing_board, last_draw
+                click_stack, input_first_frame, drawing_board, last_draw
             ]
         )
 
@@ -640,7 +414,8 @@ def seg_track_app():
         )
         
     app.queue(concurrency_count=1)
-    app.launch(debug=True, enable_queue=True, share=False)
+    app.launch(debug=True, share=False)
 
 if __name__ == "__main__":
-    seg_track_app()
+    mp.set_start_method('spawn', force=True)
+    seg_track_app()

docker-compose.yml

@@ -0,0 +1,16 @@
+
+services:
+  webuimedsam2:
+    init: true
+    restart: "always"
+    image: webuimedsam2
+    deploy:
+      resources:
+        reservations:
+          devices:
+            - driver: nvidia
+              capabilities: [gpu]
+    ports:
+      - "7860:7860"
+
+