Add simple demo

.github/workflows/check_fmt.yml

@@ -1,17 +1,17 @@
-name: SAM2/fmt
-on:
-  pull_request:
-    branches:
-    - main
-jobs:
-  ufmt_check:
-    runs-on: ubuntu-latest
-    steps:
-      - name: Check formatting
-        uses: omnilib/ufmt@action-v1
-        with:
-          path: sam2 tools
-          version: "2.0.0b2"
-          python-version: "3.10"
-          black-version: "24.2.0"
-          usort-version: "1.0.2"
+name: SAM2/fmt
+on:
+  pull_request:
+    branches:
+    - main
+jobs:
+  ufmt_check:
+    runs-on: ubuntu-latest
+    steps:
+      - name: Check formatting
+        uses: omnilib/ufmt@action-v1
+        with:
+          path: sam2 tools
+          version: "2.0.0b2"
+          python-version: "3.10"
+          black-version: "24.2.0"
+          usort-version: "1.0.2"

.gitignore

@@ -3,6 +3,7 @@
 __pycache__/
 *-checkpoint.ipynb
 .venv
+venv3.11
 *.egg*
 build/*
 _C.*

demo/README.md

@@ -0,0 +1,47 @@
+# SAM-2 Plus Interactive Video Object Segmentation
+
+This project provides a lightweight GUI to interactively prompt an object on the first frame using points and automatically propagate the segmentation across all frames using **SAM-2 Plus**.
+
+## Folder Structure
+
+```
+demo
+├── app_core.py
+├── app_gui.py
+├── engine.py
+├── point_manager.py
+├── data
+│   ├── images          # Input frames
+│   └── vos_results     # Output segmentation overlays
+└── requirements.txt
+```
+
+## How to Run
+
+1. Activate your Python environment
+2. `cd demo`
+3. `pip install -r requirements.txt`
+4. `python app_gui.py --image_dir data/images`
+
+## GUI Interaction
+
+1) Mouse Controls
+
+- **Foreground point (green)**: Left Click
+- **Background point (red)**: Right Click (windows) or Ctrl + Left Click (macOS)
+
+2) Buttons
+
+- **Clear**: Clears all placed points so you can re-annotate.
+- **Visualization**: Shows how SAM-2 Plus interprets the current points on the first frame only. You can add more points and visualize again for refinement.
+- **Segment & Run**: Finalizes the object, closes the GUI, and automatically propagates the segmentation to all frames.
+
+## Output
+
+Segmented results are saved to `data/vos_results` as overlay images.
+
+## Notes
+
+- Works on **Windows, macOS, and Linux**
+- CUDA is used automatically if available; otherwise, CPU is used
+- On macOS, use **Ctrl + Left Click** for background points

demo/app_core.py

@@ -0,0 +1,63 @@
+import os
+import cv2
+import torch
+from engine import initialize_object, overlay_mask
+
+
+class AppCore:
+    def __init__(self, predictor, frames, state, point_manager):
+        self.predictor = predictor
+        self.frames = frames
+        self.state = state
+        self.pm = point_manager
+        self.obj_id = 0
+
+    def initialize_object(self):
+        """
+        Initialize SAM-2 object memory using user prompts on frame 0.
+        """
+        print("[INFO] Initializing object with user prompt...")
+
+        initialize_object(
+            self.predictor,
+            self.state,
+            frame_idx=0,
+            points=self.pm.points,
+            labels=self.pm.labels,
+            obj_id=self.obj_id,
+        )
+
+        print("[OK] Object initialized.")
+
+    def run_full_propagation(self):
+        """
+        Run SAM-2 propagation on all remaining frames and save results.
+        Output directory: ../vos_results relative to image directory.
+        """
+        print("[INFO] Running SAM-2 propagation...")
+
+        # ------------------------------------------------
+        # Output directory: data/vos_results
+        # ------------------------------------------------
+        images_dir = os.path.dirname(self.frames[0])      # .../data/images
+        output_dir = os.path.abspath(
+            os.path.join(images_dir, "..", "vos_results")
+        )
+        os.makedirs(output_dir, exist_ok=True)
+
+        # ------------------------------------------------
+        # SAM-2 propagation (generator – must be called ONCE)
+        # ------------------------------------------------
+        for frame_idx, obj_ids, logits, _, _ in self.predictor.propagate_in_video(self.state):
+
+            # logits shape: (num_objects, 1, H, W)
+            logits = logits[0].squeeze().cpu()
+            mask = (torch.sigmoid(logits) > 0.5).numpy().astype("uint8")
+
+            img = cv2.imread(self.frames[frame_idx])
+            vis = overlay_mask(img, mask)
+
+            out_path = os.path.join(output_dir, f"frame_{frame_idx}.png")
+            cv2.imwrite(out_path, vis)
+
+        print(f"[OK] Propagation completed. Results saved to:\n{output_dir}")

demo/app_gui.py

@@ -0,0 +1,173 @@
+import os
+import cv2
+import torch
+import argparse
+import tkinter as tk
+from PIL import Image, ImageTk
+
+from sam2_plus.build_sam import build_sam2_video_predictor_plus
+from point_manager import PointManager
+from app_core import AppCore
+from engine import overlay_mask
+import warnings
+
+warnings.filterwarnings(
+    "ignore",
+    message="cannot import name '_C' from 'sam2'"
+)
+
+# ---------------- ARGUMENTS ----------------
+parser = argparse.ArgumentParser()
+parser.add_argument("--image_dir", required=True, type=str)
+args = parser.parse_args()
+
+os.environ["TQDM_DISABLE"] = "1"
+
+IMAGE_DIR = args.image_dir
+
+MODEL_CFG = "configs/sam2.1/sam2.1_hiera_b+_predmasks_decoupled_MAME.yaml"
+MODEL_CKPT = f"{os.path.dirname(__file__)}/../checkpoints/SAM2-Plus/checkpoint_phase123.pt"
+
+DISPLAY_W, DISPLAY_H = 900, 500
+
+
+# ---------------- LOAD FRAMES ----------------
+frames = sorted(
+    [os.path.join(IMAGE_DIR, f) for f in os.listdir(IMAGE_DIR)
+     if f.lower().endswith((".jpg", ".png"))]
+)
+assert len(frames) > 0, "No images found"
+
+
+# ---------------- LOAD MODEL ----------------
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+predictor = build_sam2_video_predictor_plus(
+    MODEL_CFG,
+    MODEL_CKPT,
+    device=device,
+)
+
+print("[OK] SAM-2 model and dependencies loaded successfully.")
+
+state = predictor.init_state(IMAGE_DIR)
+
+
+# ---------------- CORE ----------------
+pm = PointManager()
+core = AppCore(predictor, frames, state, pm)
+
+preview_mask = None   # for Visualization button
+
+
+# ---------------- GUI ----------------
+root = tk.Tk()
+root.title("SAM-2 Prompt (Once)")
+
+canvas = tk.Canvas(root, width=DISPLAY_W, height=DISPLAY_H, bg="black")
+canvas.pack()
+
+tk_img = None
+
+
+def draw(img):
+    global tk_img
+    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+    img = cv2.resize(img, (DISPLAY_W, DISPLAY_H))
+    tk_img = ImageTk.PhotoImage(Image.fromarray(img))
+    canvas.create_image(0, 0, anchor=tk.NW, image=tk_img)
+
+
+def draw_points(img):
+    for (x, y), lbl in zip(pm.points, pm.labels):
+        color = (0, 255, 0) if lbl == 1 else (0, 0, 255)
+        cv2.circle(img, (x, y), 6, color, -1)
+
+
+def update():
+    img = cv2.imread(frames[0])
+
+    if preview_mask is not None:
+        img = overlay_mask(img, preview_mask)
+
+    draw_points(img)
+    draw(img)
+
+
+# ---------------- MOUSE ----------------
+def on_click(event):
+    img = cv2.imread(frames[0])
+    h, w = img.shape[:2]
+
+    x = int(event.x * w / DISPLAY_W)
+    y = int(event.y * h / DISPLAY_H)
+
+    # if event.num == 1:
+    #     pm.add_point(x, y, 1)   # FG
+    # elif event.num == 3:
+    #     pm.add_point(x, y, 0)   # BG
+
+    if event.num == 1 and not (event.state & 0x0004):  # Left click
+        pm.add_point(x, y, 1)   # Foreground
+    else:
+        pm.add_point(x, y, 0)   # Background (Right click OR Ctrl+Left)
+
+
+    update()
+
+
+canvas.bind("<Button-1>", on_click)
+canvas.bind("<Control-Button-1>", on_click)
+canvas.bind("<Button-3>", on_click)
+
+
+# ---------------- BUTTON ACTIONS ----------------
+def clear_points():
+    global preview_mask
+    pm.clear()
+    preview_mask = None
+    update()
+
+
+def visualize_prompt():
+    """
+    Visualize SAM-2 response on frame 0 ONLY.
+    Does NOT commit object memory.
+    """
+    global preview_mask
+
+    if len(pm.points) == 0:
+        return
+
+    res = predictor.add_new_points_or_box(
+        inference_state=state,
+        frame_idx=0,
+        points=pm.points,
+        labels=pm.labels,
+        obj_id=0,
+    )
+
+    logits = res[2][0].squeeze().cpu()
+    preview_mask = (torch.sigmoid(logits) > 0.5).numpy().astype("uint8")
+
+    update()
+
+
+def segment_and_run():
+    core.initialize_object()
+    root.destroy()
+    core.run_full_propagation()
+
+
+# ---------------- BUTTONS ----------------
+btn_frame = tk.Frame(root)
+btn_frame.pack(pady=10)
+
+tk.Button(btn_frame, text="Clear", width=12, command=clear_points).grid(row=0, column=0, padx=5)
+tk.Button(btn_frame, text="Visualization", width=12, command=visualize_prompt).grid(row=0, column=1, padx=5)
+tk.Button(btn_frame, text="Segment & Run", width=14, command=segment_and_run).grid(row=0, column=2, padx=5)
+
+
+# ---------------- START ----------------
+update()
+root.mainloop()

demo/engine.py

@@ -0,0 +1,65 @@
+import numpy as np
+import torch
+import cv2
+
+
+# ============================================================
+# ONE-TIME INTERACTION (PROMPT ONCE)
+# ============================================================
+def initialize_object(predictor, state, frame_idx, points, labels, obj_id=0):
+    """
+    Initialize object memory using user prompt.
+    This MUST be called exactly once.
+    """
+    if len(points) == 0:
+        raise RuntimeError("No points provided for initialization")
+
+    res = predictor.add_new_points_or_box(
+        inference_state=state,
+        frame_idx=frame_idx,
+        points=np.array(points),
+        labels=np.array(labels),
+        obj_id=obj_id,
+    )
+
+    # SAM-2 output format:
+    # res = (obj_id, obj_ids, mask_logits)
+    logits = res[2]            # tensor (1, 1, H, W)
+
+    logits = logits.squeeze().cpu()
+    mask = (torch.sigmoid(logits) > 0.5).numpy().astype("uint8")
+
+    return mask
+
+
+# ============================================================
+# PURE PROPAGATION (NO PROMPTS)
+# ============================================================
+def propagate_next(predictor, state):
+    """
+    Propagate object to the next frame using SAM-2 VOS generator.
+    """
+    # propagate_in_video returns a generator
+    gen = predictor.propagate_in_video(state)
+
+    # Get next frame result
+    frame_idx, obj_ids, logits = next(gen)
+
+    # logits shape: (num_objects, 1, H, W)
+    logits = logits[0].squeeze().cpu()   # first object only
+
+    mask = (torch.sigmoid(logits) > 0.5).numpy().astype("uint8")
+    return mask
+
+
+# ============================================================
+# VISUALIZATION
+# ============================================================
+def overlay_mask(img, mask, alpha=0.5):
+    """
+    Overlay binary mask on image.
+    """
+    mask = cv2.resize(mask, (img.shape[1], img.shape[0]))
+    overlay = img.copy()
+    overlay[mask == 1] = (0, 255, 0)
+    return cv2.addWeighted(overlay, alpha, img, 1 - alpha, 0)