|
|
--- |
|
|
license: apache-2.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 tqdm import tqdm |
|
|
from natsort import natsorted |
|
|
|
|
|
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 |
|
|
|
|
|
predictor = build_sam2_video_predictor_plus( |
|
|
config_file="configs/sam2.1/sam2.1_hiera_b+_predmasks_decoupled_MAME.yaml", |
|
|
ckpt_path="./checkpoints/SAM2-Plus/checkpoint_phase123.pt", |
|
|
apply_postprocessing=False, |
|
|
hydra_overrides_extra=[ |
|
|
"++model.non_overlap_masks=" + ("false") |
|
|
], |
|
|
vos_optimized=False, |
|
|
task='mask' |
|
|
) |
|
|
|
|
|
input_video_dir = "./examples/JPEGImages/horsejump-low" |
|
|
input_mask_path = "./examples/Annotations/horsejump-low/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", ".png", ".PNG"] |
|
|
] |
|
|
frame_names = natsorted(frame_names) |
|
|
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 tqdm(range(0, len(frame_names)), desc="Visualization Results"): |
|
|
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) |
|
|
plt.close() |
|
|
``` |
|
|
|
|
|
**[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 |
|
|
from tqdm import tqdm |
|
|
from natsort import natsorted |
|
|
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 |
|
|
|
|
|
predictor = build_sam2_video_predictor_plus( |
|
|
config_file="configs/sam2.1/sam2.1_hiera_b+_predmasks_decoupled_MAME.yaml", |
|
|
ckpt_path="./checkpoints/SAM2-Plus/checkpoint_phase123.pt", |
|
|
apply_postprocessing=False, |
|
|
hydra_overrides_extra=[ |
|
|
"++model.non_overlap_masks=" + ("false") |
|
|
], |
|
|
vos_optimized=False, |
|
|
task='box' |
|
|
) |
|
|
|
|
|
input_video_dir = "./examples/JPEGImages/horsejump-low" |
|
|
input_box_path = "./examples/Boxes/horsejump-low.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", ".png", ".PNG"] |
|
|
] |
|
|
frame_names = natsorted(frame_names) |
|
|
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: |
|
|
print(f"Box file {input_box_path} not found. Using default box.") |
|
|
input_box_xywh = [316,385,742,488] |
|
|
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 tqdm(range(0, len(frame_names)), desc="Visualization Results"): |
|
|
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) |
|
|
plt.close() |
|
|
``` |
|
|
|
|
|
**[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 |
|
|
from tqdm import tqdm |
|
|
import numpy as np |
|
|
from natsort import natsorted |
|
|
|
|
|
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.pt_inference_plus import load_visible_points_from_npz |
|
|
|
|
|
predictor = build_sam2_video_predictor_plus( |
|
|
config_file="configs/sam2.1/sam2.1_hiera_b+_predmasks_decoupled_MAME.yaml", |
|
|
ckpt_path="./checkpoints/SAM2-Plus/checkpoint_phase123.pt", |
|
|
apply_postprocessing=False, |
|
|
hydra_overrides_extra=[ |
|
|
"++model.non_overlap_masks=" + ("false") |
|
|
], |
|
|
vos_optimized=False, |
|
|
task='point' |
|
|
) |
|
|
|
|
|
input_video_dir = "./examples/JPEGImages/horsejump-low" |
|
|
input_point_path = "./examples/Points/horsejump-low.npz" |
|
|
output_point_dir = "./output/Points/" |
|
|
|
|
|
radius, sigma = 5, 2 |
|
|
score_thresh = 0 |
|
|
|
|
|
with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16): |
|
|
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", ".png", ".PNG"] |
|
|
] |
|
|
frame_names = natsorted(frame_names) |
|
|
|
|
|
inference_state = predictor.init_state(video_path=input_video_dir) |
|
|
height = inference_state["video_height"] |
|
|
width = inference_state["video_width"] |
|
|
|
|
|
input_frame_idx = 0 # the frame index we interact with |
|
|
object_id = 0 # give a unique id to each object we interact with (it can be any integers) |
|
|
num_frames, num_points = len(frame_names), 1 |
|
|
|
|
|
input_data = np.load(input_point_path, allow_pickle=True) |
|
|
input_point, input_visible = torch.tensor(input_data['trajs_2d'].astype(np.float32)), torch.tensor(input_data['visibs'].astype(bool)) |
|
|
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 |
|
|
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 tqdm(range(0, len(frame_names)), desc="Visualization Results"): |
|
|
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) |
|
|
plt.close() |
|
|
``` |
|
|
|
|
|
### 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") |
|
|
``` |
|
|
|
|
|
|
