Views
No views yet
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()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()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()import torch
from sam2_plus.sam2_video_predictor import SAM2VideoPredictor_Plus
predictor = SAM2VideoPredictor_Plus.from_pretrained("MCG-NJU/SAM2-Plus")