Homepage|Paper|GitHub
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 for more information.
Usage
Video Object Segmentation (Mask Granularity)
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)
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)
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
import torch
from sam2_plus.sam2_video_predictor import SAM2VideoPredictor_Plus
predictor = SAM2VideoPredictor_Plus.from_pretrained("MCG-NJU/SAM2-Plus")
