Spaces:
Running on Zero
Running on Zero
| """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 | |
| 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) | |
| 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 | |
| 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 | |