diff --git a/.github/workflows/check_fmt.yml b/.github/workflows/check_fmt.yml index 0a29b88..6fc9f4e 100644 --- a/.github/workflows/check_fmt.yml +++ b/.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" diff --git a/.gitignore b/.gitignore index 7df3f9c..8c7c951 100644 --- a/.gitignore +++ b/.gitignore @@ -3,6 +3,7 @@ __pycache__/ *-checkpoint.ipynb .venv +venv3.11 *.egg* build/* _C.* diff --git a/demo/README.md b/demo/README.md new file mode 100644 index 0000000..15404e0 --- /dev/null +++ b/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 diff --git a/demo/app_core.py b/demo/app_core.py new file mode 100644 index 0000000..ca61822 --- /dev/null +++ b/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}") diff --git a/demo/app_gui.py b/demo/app_gui.py new file mode 100644 index 0000000..bf8d137 --- /dev/null +++ b/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("", on_click) +canvas.bind("", on_click) +canvas.bind("", 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() diff --git a/demo/data/images/001.jpg b/demo/data/images/001.jpg new file mode 100644 index 0000000..81a02ab Binary files /dev/null and b/demo/data/images/001.jpg differ diff --git a/demo/data/images/002.jpg b/demo/data/images/002.jpg new file mode 100644 index 0000000..05a7be7 Binary files /dev/null and b/demo/data/images/002.jpg differ diff --git a/demo/data/images/003.jpg b/demo/data/images/003.jpg new file mode 100644 index 0000000..273cad7 Binary files /dev/null and b/demo/data/images/003.jpg differ diff --git a/demo/data/images/004.jpg b/demo/data/images/004.jpg new file mode 100644 index 0000000..74bee68 Binary files /dev/null and b/demo/data/images/004.jpg differ diff --git a/demo/data/images/005.jpg b/demo/data/images/005.jpg new file mode 100644 index 0000000..d79c3f1 Binary files /dev/null and b/demo/data/images/005.jpg differ diff --git a/demo/data/images/006.jpg b/demo/data/images/006.jpg new file mode 100644 index 0000000..fbcd03c Binary files /dev/null and b/demo/data/images/006.jpg differ diff --git a/demo/data/images/007.jpg b/demo/data/images/007.jpg new file mode 100644 index 0000000..5442944 Binary files /dev/null and b/demo/data/images/007.jpg differ diff --git a/demo/data/images/008.jpg b/demo/data/images/008.jpg new file mode 100644 index 0000000..8649860 Binary files /dev/null and b/demo/data/images/008.jpg differ diff --git a/demo/data/images/009.jpg b/demo/data/images/009.jpg new file mode 100644 index 0000000..df2f682 Binary files /dev/null and b/demo/data/images/009.jpg differ diff --git a/demo/data/images/010.jpg b/demo/data/images/010.jpg new file mode 100644 index 0000000..c970503 Binary files /dev/null and b/demo/data/images/010.jpg differ diff --git a/demo/engine.py b/demo/engine.py new file mode 100644 index 0000000..f77a422 --- /dev/null +++ b/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) diff --git a/demo/point_manager.py b/demo/point_manager.py new file mode 100644 index 0000000..92d4af2 --- /dev/null +++ b/demo/point_manager.py @@ -0,0 +1,19 @@ +class PointManager: + """ + Stores user interaction points for SAM-2. + labels: + 1 -> positive (foreground) + 0 -> negative (background) + """ + + def __init__(self): + self.points = [] + self.labels = [] + + def add_point(self, x, y, label): + self.points.append([x, y]) + self.labels.append(label) + + def clear(self): + self.points = [] + self.labels = [] diff --git a/demo/requirements.txt b/demo/requirements.txt new file mode 100644 index 0000000..1bf95c5 --- /dev/null +++ b/demo/requirements.txt @@ -0,0 +1,7 @@ +torch +torchvision +numpy +opencv-python +Pillow +tqdm +pyyaml \ No newline at end of file diff --git a/sam2_plus/modeling/sam/box_head.py b/sam2_plus/modeling/sam/box_head.py index 28fa662..73f8bc4 100755 --- a/sam2_plus/modeling/sam/box_head.py +++ b/sam2_plus/modeling/sam/box_head.py @@ -88,9 +88,12 @@ def __init__(self, inplanes=64, channel=256, feat_sz=20, stride=16, freeze_bn=Fa self.indice = torch.arange(0, self.feat_sz).view(-1, 1) * self.stride # generate mesh-grid self.coord_x = self.indice.repeat((self.feat_sz, 1)) \ - .view((self.feat_sz * self.feat_sz,)).float().cuda() + .view((self.feat_sz * self.feat_sz,)).float() self.coord_y = self.indice.repeat((1, self.feat_sz)) \ - .view((self.feat_sz * self.feat_sz,)).float().cuda() + .view((self.feat_sz * self.feat_sz,)).float() + if torch.cuda.is_available(): + self.coord_x = self.coord_x.cuda() + self.coord_y = self.coord_y.cuda() def forward(self, x, return_dist=False, softmax=True): """ Forward pass with input x. """ @@ -181,9 +184,12 @@ def __init__(self, inplanes=64, channel=256, feat_sz=20, stride=16, freeze_bn=Fa self.indice = torch.arange(0, self.feat_sz).view(-1, 1) * self.stride # generate mesh-grid self.coord_x = self.indice.repeat((self.feat_sz, 1)) \ - .view((self.feat_sz * self.feat_sz,)).float().cuda() + .view((self.feat_sz * self.feat_sz,)).float() self.coord_y = self.indice.repeat((1, self.feat_sz)) \ - .view((self.feat_sz * self.feat_sz,)).float().cuda() + .view((self.feat_sz * self.feat_sz,)).float() + if torch.cuda.is_available(): + self.coord_x = self.coord_x.cuda() + self.coord_y = self.coord_y.cuda() def forward(self, x, return_dist=False, softmax=True): """ Forward pass with input x. """