facebookresearch / facebookresearch/sam2
Question about implementing multiple object video tracking
- Dominant language
- Jupyter Notebook
- Stars
- 19.9k
- Forks
- 2.5k
- PR merge metrics
- No merged PRs in 30d
Description
I want to segment multiple objects based on bounding boxes. I didn't find code for batch processing video tracking, so I chose to process them one by one and gradually add IDs, and merge all the segmented masks into a single mask. I wrote this code, but it doesn't work. Can you tell me where the problem is? Thank you very much for your help.
```
import torch
import numpy as np
import cv2
# Define the device (prefer CUDA:1 if available, otherwise use CPU)
device = torch.device("cuda:1" if torch.cuda.is_available() else "cpu")
# Path to the model checkpoint
checkpoint = "./sam2_hiera_large.pt"
# Path to the model configuration file
model_cfg = "sam2_hiera_l.yaml"
# Build the SAM2 video predictor
predictor = build_sam2_video_predictor(model_cfg, checkpoint, device=device)
# Directory path containing video frames
video_dir = "./video_frame"
# Initialize the inference state
inference_state = predictor.init_state(video_path=video_dir)
# Define bounding boxes
boxes = np.array([
[616.26, 717.36, 699.88, 760.17],
[64.953, 683.81, 109.53, 710.87],
[20.893, 712.16, 72.295, 738.23],
[1744.1, 629.83, 1800.7, 656.11],
[265.73, 709.02, 334.88, 729.54],
[496.77, 697.25, 566.11, 722.93],
[54.411, 740.81, 136.61, 774.92],
[565.42, 666.2, 618.92, 696.09],
[907.14, 692.92, 988.52, 724.91],
[1813.7, 631.45, 1858.8, 657.15],
[1839.7, 633.64, 1884, 652.13],
[1123.2, 670.81, 1186.8, 697.15],
[1769.9, 628.47, 1839, 654.66]
], dtype=np.float32)
# Frame index for annotation
ann_frame_idx = 0
# Initialize a boolean mask for merging masks
merge_mask = np.full((1080, 1920), False, dtype=bool)
# Labels for positive clicks (1 for positive, 0 for negative)
labels = np.array([1], dtype=np.int32)
# Iterate over each bounding box
for i, box in enumerate(boxes):
# Object ID for this bounding box
ann_obj_id = i + 1
print(box)
# Calculate the center point of the bounding box
points = np.array([[(box[0] + box[2]) / 2, (box[1] + box[3]) / 2]], dtype=np.float32)
print(points)
# Add new points or box to the inference state
_, out_obj_ids, out_mask_logits = predictor.add_new_points_or_box(
inference_state=inference_state,
frame_idx=ann_frame_idx,
obj_id=ann_obj_id,
points=points,
labels=labels,
box=box,
)
# Convert the mask logits to a binary mask
mask = (out_mask_logits[0] > 0.0).cpu().numpy().squeeze()
# Load the first frame of the video
image = cv2.imread("/home/li/zw/segment-anything-2/video_frame/00000.jpg")
# Apply the mask to the image
mask_image = apply_mask_to_image(image, mask)
# Merge the mask into the overall mask
merge_mask |= mask
```
Contributor guide
Assessment
This issue has not been assessed yet.