Update README.md

README.md

@@ -1,3 +1,317 @@
----
-license: cc-by-nc-4.0
----
+---
+license: cc-by-nc-4.0
+datasets:
+- MCG-NJU/Tracking-Any-Granularity
+library_name: transformers
+---
+
+🏠 [Homepage](https://tracking-any-granularity.github.io/) | 📄 [Paper](https://arxiv.org/abs/2510.18822) | 🔗 [GitHub](https://github.com/MCG-NJU/SAM2-Plus)
+
+Model repository for SAM 2++: Tracking Anything at Any Granularity, a unified video tracking framework that extends the SAM 2 model to track any targets in videos at any granularity, including masks, bounding boxes, and points.
+See the [SAM 2++ paper](https://arxiv.org/abs/2510.18822) for more information.
+
+## Usage
+
+**[Video Object Segmentation (Mask Granularity)](https://github.com/MCG-NJU/SAM2-Plus/blob/main/tools/vos_inference_plus.sh)**
+```
+import os
+import torch
+import matplotlib.pyplot as plt
+from PIL import Image
+
+from sam2_plus.build_sam import build_sam2_video_predictor_plus
+
+from tools.visualization import show_mask, show_box, show_points
+from tools.vos_inference import load_ann_png, get_per_obj_mask, DAVIS_PALETTE, save_masks_to_dir
+
+checkpoint = "./checkpoints/SAM2-Plus/checkpoint_phase123.pt"
+model_cfg = "configs/sam2.1/sam2.1_hiera_b+_predmasks_decoupled_MAME.yaml"
+predictor = build_sam2_video_predictor_plus(model_cfg, checkpoint, task="mask")
+
+input_video_dir = "./examples/JPEGImages/00001"
+input_mask_path = "./examples/Annotations/00001/00000.png"
+output_mask_dir = "./output/Annotations/"
+
+score_thresh = 0
+
+with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16):
+    inference_state = predictor.init_state(video_path=input_video_dir)
+
+    video_name = os.path.basename(input_video_dir)
+    frame_names = [
+        os.path.splitext(p)[0]
+        for p in os.listdir(input_video_dir)
+        if os.path.splitext(p)[-1] in [".jpg", ".jpeg", ".JPG", ".JPEG"]
+    ]
+    frame_names.sort()
+    height = inference_state["video_height"]
+    width = inference_state["video_width"]
+
+    input_frame_idx = 0     # the frame index we interact with
+    object_id = 1           # give a unique id to each object we interact with (it can be any integers)
+
+    input_palette = None
+    input_mask, input_palette = load_ann_png(input_mask_path)
+    per_obj_input_mask = get_per_obj_mask(input_mask)
+    object_mask = per_obj_input_mask[object_id]
+
+    predictor.add_new_mask(
+        inference_state=inference_state,
+        frame_idx=input_frame_idx,
+        obj_id=object_id,
+        mask=object_mask,
+    )
+
+    # run propagation throughout the video and collect the results in a dict
+    os.makedirs(os.path.join(output_mask_dir, video_name), exist_ok=True)
+    output_palette = input_palette or DAVIS_PALETTE
+    video_segments = {}  # video_segments contains the per-frame segmentation results
+    for out_frame_idx, out_obj_ids, out_mask_logits, _, _ in predictor.propagate_in_video(
+        inference_state
+    ):
+        per_obj_output_mask = {
+            out_obj_id: (out_mask_logits[i] > score_thresh).cpu().numpy()
+            for i, out_obj_id in enumerate(out_obj_ids)
+        }
+        video_segments[out_frame_idx] = per_obj_output_mask
+    
+    # write the output masks as palette PNG files to output_mask_dir
+    for out_frame_idx, per_obj_output_mask in video_segments.items():
+        save_masks_to_dir(
+            output_mask_dir=output_mask_dir,
+            video_name=video_name,
+            frame_name=frame_names[out_frame_idx],
+            per_obj_output_mask=per_obj_output_mask,
+            height=height,
+            width=width,
+            per_obj_png_file=False,
+            output_palette=output_palette,
+        )
+    
+    # visualize the tracking results
+    for out_frame_idx in range(0, len(frame_names)):
+        plt.clf()
+        plt.figure()
+        # plt.title(f"frame {out_frame_idx}")
+        plt.imshow(Image.open(os.path.join(input_video_dir, frame_names[out_frame_idx] + ".jpg")))
+        for out_obj_id, out_mask in video_segments[out_frame_idx].items():
+            show_mask(out_mask, plt.gca(), obj_id=out_obj_id)
+        plt.axis('off')
+        plt.subplots_adjust(left=0, right=1, top=1, bottom=0)
+        plt.savefig(f"{output_mask_dir}/{video_name}/{out_frame_idx:05d}_withMask.png", dpi=300, bbox_inches='tight', pad_inches=0)
+```
+
+**[Video Object Tracking (Box Granularity)](https://github.com/MCG-NJU/SAM2-Plus/blob/main/tools/sot_inference_plus.sh)**
+```
+import os
+import torch
+import matplotlib.pyplot as plt
+from PIL import Image
+import numpy as np
+import logging
+
+from sam2_plus.build_sam import build_sam2_video_predictor_plus
+
+from tools.visualization import show_mask, show_box, show_points
+from tools.vos_inference import load_ann_png, get_per_obj_mask, DAVIS_PALETTE, save_masks_to_dir
+from tools.sot_inference import save_boxes_to_dir, save_masks_and_boxes_to_dir
+from training.dataset_plus.box.utils import np_box_xywh_to_xyxy, np_box_xyxy_to_xywh, np_masks_to_boxes, np_box_clamp_xywh
+from benchmarks.sot_benchmark.datasets.utils import load_text
+
+checkpoint = "./checkpoints/SAM2-Plus/checkpoint_phase123.pt"
+model_cfg = "configs/sam2.1/sam2.1_hiera_b+_predmasks_decoupled_MAME.yaml"
+predictor = build_sam2_video_predictor_plus(model_cfg, checkpoint, task="box")
+
+input_video_dir = "./examples/JPEGImages/00001"
+input_box_path = "./examples/Boxes/00001.txt"
+output_box_dir = "./output/Boxes/"
+
+score_thresh = 0
+
+with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16):
+    inference_state = predictor.init_state(video_path=input_video_dir)
+
+    video_name = os.path.basename(input_video_dir)
+    frame_names = [
+        os.path.splitext(p)[0]
+        for p in os.listdir(input_video_dir)
+        if os.path.splitext(p)[-1] in [".jpg", ".jpeg", ".JPG", ".JPEG"]
+    ]
+    frame_names.sort()
+    height = inference_state["video_height"]
+    width = inference_state["video_width"]
+
+    input_frame_idx = 0     # the frame index we interact with
+    object_id = 1           # give a unique id to each object we interact with (it can be any integers)
+
+    input_palette = None
+    if os.path.isfile(input_box_path):
+        input_box_xywh = load_text(str(input_box_path), delimiter=',', dtype=np.float64, backend='numpy').reshape(-1, 4)[0]
+    else:
+        input_box_xywh = [1026,361,222,169]
+    per_obj_input_box_xyxy = {1: np_box_xywh_to_xyxy(np.array(input_box_xywh))}
+    object_box_xyxy = per_obj_input_box_xyxy[object_id]
+
+    frame_idx, obj_ids, masks, _ = predictor.add_new_points_or_box(
+        inference_state=inference_state,
+        frame_idx=input_frame_idx,
+        obj_id=object_id,
+        box=object_box_xyxy,
+    )
+
+    # run propagation throughout the video and collect the results in a dict
+    output_palette = input_palette or DAVIS_PALETTE
+    video_segments = {}  # video_segments contains the per-frame segmentation results
+    video_boxes_xywh = {}  # video_boxes_xyxy contains the per-frame bounding box results
+    for out_frame_idx, out_obj_ids, out_mask_logits, output_box_xyxy, out_obj_score_logits in predictor.propagate_in_video(
+        inference_state=inference_state,
+    ):
+        if torch.any(output_box_xyxy[:,:,0] >= output_box_xyxy[:,:,2]) or torch.any(output_box_xyxy[:,:,1] >= output_box_xyxy[:,:,3]):
+            logging.warning(f"Invalid box prediction: {output_box_xyxy}")
+    
+        per_obj_output_mask = {
+            out_obj_id: (out_mask_logits[i] > score_thresh).cpu().numpy()
+            for i, out_obj_id in enumerate(out_obj_ids)
+        }
+        video_segments[out_frame_idx] = per_obj_output_mask
+        per_obj_output_box_xywh = {
+            out_obj_id: np_box_clamp_xywh(np_box_xyxy_to_xywh(output_box_xyxy[i].cpu().numpy()))
+            for i, out_obj_id in enumerate(out_obj_ids)
+        }
+        video_boxes_xywh[out_frame_idx] = per_obj_output_box_xywh
+    
+    # save the tracking results
+    save_boxes_to_dir(
+        output_bbox_dir=output_box_dir,
+        video_name=video_name,
+        video_boxes_xywh=video_boxes_xywh,
+    )
+    
+    # visualize the tracking results
+    os.makedirs(os.path.join(output_box_dir, video_name), exist_ok=True)
+    for out_frame_idx in range(0, len(frame_names)):
+        plt.clf()
+        plt.figure()
+        # plt.title(f"frame {out_frame_idx}")
+        plt.imshow(Image.open(os.path.join(input_video_dir, frame_names[out_frame_idx] + ".jpg")))
+        for out_obj_id, out_box in video_boxes_xywh[out_frame_idx].items():
+            box_xywh = out_box[0]
+            box_xyxy = np_box_xywh_to_xyxy(np.array(box_xywh))
+            show_box(box_xyxy, plt.gca())
+        plt.axis('off')
+        plt.subplots_adjust(left=0, right=1, top=1, bottom=0)
+        plt.savefig(os.path.join(output_box_dir, video_name, f"{out_frame_idx:05d}_withbox.png"), dpi=300, bbox_inches='tight', pad_inches=0)
+```
+
+**[Point Tracking (Point Granularity)](https://github.com/MCG-NJU/SAM2-Plus/blob/main/tools/pt_inference_plus.sh)**
+```
+import os
+import torch
+import matplotlib.pyplot as plt
+from PIL import Image
+import numpy as np
+import logging
+
+from sam2_plus.build_sam import build_sam2_video_predictor_plus
+
+from tools.visualization import show_mask, show_box, show_points
+from tools.vos_inference import load_ann_png, get_per_obj_mask, DAVIS_PALETTE, save_masks_to_dir
+from tools.sot_inference import save_boxes_to_dir, save_masks_and_boxes_to_dir
+from tools.pt_inference_plus import load_visible_points_from_npz
+from training.dataset_plus.box.utils import np_box_xywh_to_xyxy, np_box_xyxy_to_xywh, np_masks_to_boxes, np_box_clamp_xywh
+from benchmarks.sot_benchmark.datasets.utils import load_text
+
+checkpoint = "./checkpoints/SAM2-Plus/checkpoint_phase123.pt"
+model_cfg = "configs/sam2.1/sam2.1_hiera_b+_predmasks_decoupled_MAME.yaml"
+predictor = build_sam2_video_predictor_plus(model_cfg, checkpoint, task="point")
+
+input_video_dir = "./examples/JPEGImages/00001"
+input_point_path = "./examples/Points/00001.npz"
+output_point_dir = "./output/Points/"
+
+radius, sigma = 5, 2
+
+score_thresh = 0
+
+with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16):
+    inference_state = predictor.init_state(video_path=input_video_dir)
+
+    video_name = os.path.basename(input_video_dir)
+    frame_names = [
+        os.path.splitext(p)[0]
+        for p in os.listdir(input_video_dir)
+        if os.path.splitext(p)[-1] in [".jpg", ".jpeg", ".JPG", ".JPEG"]
+    ]
+    frame_names.sort()
+    height = inference_state["video_height"]
+    width = inference_state["video_width"]
+
+    input_frame_idx = 0     # the frame index we interact with
+    object_id = 1           # give a unique id to each object we interact with (it can be any integers)
+
+    input_palette = None
+    gt_data = np.load(input_point_path, allow_pickle=True)
+    trajs = gt_data['trajs_2d'].astype(np.float32)    # ndarray [N_frames, N_points, 2], xyxy
+    visible = gt_data['visibs'].astype(bool)    # ndarray [N_frames, N_points], bool
+    input_point, input_visible = torch.tensor(trajs), torch.tensor(visible)
+    per_obj_input_point = load_visible_points_from_npz(
+        input_points=input_point,
+        input_visibles=input_visible,
+        frame_idx=input_frame_idx,
+    )
+    object_point = per_obj_input_point[object_id]
+
+    predictor.add_new_points_and_generate_gaussian_mask(
+        inference_state=inference_state,
+        frame_idx=input_frame_idx,
+        obj_id=object_id,
+        points=object_point.unsqueeze(0).numpy(),
+        labels=np.array([1]),
+        radius=radius,
+        sigma=sigma,
+    )
+
+    # run propagation throughout the video and collect the results in a dict
+    num_frames, num_points = len(frame_names), trajs.shape[1]
+    point_array = -np.ones((num_frames, num_points, 2), dtype=np.float32)
+    visible_array = np.zeros((num_frames, num_points), dtype=bool)
+    for out_frame_idx, out_obj_ids, out_mask_logits, out_box_xyxys, out_obj_score_logits in predictor.propagate_in_video(
+        inference_state
+    ):
+        for out_obj_id, out_mask_logit, out_obj_score_logit in zip(out_obj_ids, out_mask_logits, out_obj_score_logits):
+            out_mask_logit, out_obj_score_logit = out_mask_logit.squeeze(0), out_obj_score_logit.squeeze(0)
+            max_index = torch.argmax(out_mask_logit)
+            max_score_y, max_score_x = torch.unravel_index(max_index, out_mask_logit.shape)
+            point_array[out_frame_idx, out_obj_id] = np.array([max_score_x.cpu(), max_score_y.cpu()])
+            visible_array[out_frame_idx, out_obj_id] = (out_obj_score_logit > score_thresh).cpu().numpy()
+    
+    # write the output masks as palette PNG files to output_mask_dir
+    os.makedirs(output_point_dir, exist_ok=True)
+    np.savez(os.path.join(output_point_dir, f"{video_name}.npz"), trajs_2d=point_array, visibs=visible_array, size=(width, height))
+    
+    # visualize the tracking results
+    os.makedirs(os.path.join(output_point_dir, video_name), exist_ok=True)
+    for out_frame_idx in range(0, len(frame_names)):
+        plt.clf()
+        plt.figure()
+        # plt.title(f"frame {out_frame_idx}")
+        plt.imshow(Image.open(os.path.join(input_video_dir, frame_names[out_frame_idx] + ".jpg")))
+        points = point_array[out_frame_idx, object_id].reshape(1, 2)
+        labels = np.array([-1], np.int32)
+        show_points(points, labels, plt.gca(), marker_size=20, edgecolor=None)
+        plt.axis('off')
+        plt.subplots_adjust(left=0, right=1, top=1, bottom=0)
+        plt.savefig(os.path.join(output_point_dir, video_name, f"{out_frame_idx:05d}_withPoint.png"), dpi=300, bbox_inches='tight', pad_inches=0)
+```
+
+### Load from 🤗 Hugging Face
+
+Models can alternatively be loaded from [Hugging Face](https://huggingface.co/MCG-NJU/SAM2-Plus)
+
+```
+import torch
+from sam2_plus.sam2_video_predictor import SAM2VideoPredictor_Plus
+
+predictor = SAM2VideoPredictor_Plus.from_pretrained("MCG-NJU/SAM2-Plus")
+```
+