Spaces:
Running on Zero
Running on Zero
File size: 3,430 Bytes
168e955 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 | """SAM3 mask tracking for the demo, through `transformers`.
Points: clicks on the first frame -> Sam3TrackerVideoModel propagates one object.
Text: a short noun phrase -> Sam3VideoModel finds every matching instance in
every frame; their union is the mask.
"""
import numpy as np
import torch
from PIL import Image
from transformers import (
Sam3TrackerVideoModel,
Sam3TrackerVideoProcessor,
Sam3VideoModel,
Sam3VideoProcessor,
)
SAM3_ID = "facebook/sam3"
DTYPE = torch.bfloat16
class Sam3Tracker:
def __init__(self, device="cuda", token=None):
self.device = device
self.tracker = Sam3TrackerVideoModel.from_pretrained(SAM3_ID, token=token).to(device, dtype=DTYPE).eval()
self.tracker_proc = Sam3TrackerVideoProcessor.from_pretrained(SAM3_ID, token=token)
self.video_model = Sam3VideoModel.from_pretrained(SAM3_ID, token=token).to(device, dtype=DTYPE).eval()
self.video_proc = Sam3VideoProcessor.from_pretrained(SAM3_ID, token=token)
# ---------------------------------------------------------------- points
def _point_session(self, frames, points, labels):
sess = self.tracker_proc.init_video_session(
video=[Image.fromarray(f) for f in frames], inference_device=self.device, dtype=DTYPE)
self.tracker_proc.add_inputs_to_inference_session(
inference_session=sess, frame_idx=0, obj_ids=1,
input_points=[[[[float(x), float(y)] for x, y in points]]],
input_labels=[[[int(l) for l in labels]]],
)
return sess
def _to_mask(self, pred_masks, sess):
m = self.tracker_proc.post_process_masks(
[pred_masks], original_sizes=[[sess.video_height, sess.video_width]], binarize=True)[0]
return (m[0, 0].cpu().numpy() > 0).astype(np.uint8) * 255
@torch.inference_mode()
def first_frame_mask(self, frame, points, labels):
sess = self._point_session([frame], points, labels)
out = self.tracker(inference_session=sess, frame_idx=0)
return self._to_mask(out.pred_masks, sess)
@torch.inference_mode()
def track_points(self, frames, points, labels):
sess = self._point_session(frames, points, labels)
self.tracker(inference_session=sess, frame_idx=0)
masks = [None] * len(frames)
for out in self.tracker.propagate_in_video_iterator(sess):
masks[out.frame_idx] = self._to_mask(out.pred_masks, sess)
empty = np.zeros(frames[0].shape[:2], np.uint8)
return [m if m is not None else empty for m in masks]
# ------------------------------------------------------------------ text
@torch.inference_mode()
def track_text(self, frames, text):
sess = self.video_proc.init_video_session(
video=[Image.fromarray(f) for f in frames], inference_device=self.device, dtype=DTYPE)
sess = self.video_proc.add_text_prompt(inference_session=sess, text=text)
h, w = frames[0].shape[:2]
masks = [np.zeros((h, w), np.uint8) for _ in frames]
for out in self.video_model.propagate_in_video_iterator(inference_session=sess, max_frame_num_to_track=len(frames)):
res = self.video_proc.postprocess_outputs(sess, out)
m = res["masks"]
if m is not None and len(m) > 0:
masks[out.frame_idx] = (m.any(0).cpu().numpy() > 0).astype(np.uint8) * 255
return masks
|