Buckets:
| from pathlib import Path | |
| import cv2 | |
| import numpy as np | |
| class SAM31Tracker: | |
| def __init__(self, checkpoint, device="cuda"): | |
| import torch | |
| from sam3.model_builder import build_sam3_multiplex_video_predictor | |
| if not device.startswith("cuda") or not torch.cuda.is_available(): | |
| raise RuntimeError("SAM 3.1 video tracking requires CUDA") | |
| torch.cuda.set_device(torch.device(device)) | |
| self.predictor = build_sam3_multiplex_video_predictor( | |
| checkpoint_path=str(checkpoint), | |
| use_fa3=False, | |
| compile=False, | |
| async_loading_frames=False, | |
| ) | |
| self.predictor.bf16_context.__exit__(None, None, None) | |
| self.model = self.predictor.model | |
| def track(self, frames_folder, prompts): | |
| import torch | |
| frames_folder = Path(frames_folder) | |
| files = sorted(frames_folder.glob("*.jpg")) | |
| if not files or not prompts: | |
| raise ValueError("Frames and tracking prompts are required") | |
| height, width = cv2.imread(str(files[0])).shape[:2] | |
| names = list(prompts) | |
| masks = np.zeros((len(names), len(files), height, width), dtype=bool) | |
| confidence = np.zeros((len(names), len(files)), dtype=float) | |
| with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16): | |
| state = self.model.init_state( | |
| resource_path=str(frames_folder), | |
| offload_video_to_cpu=True, | |
| async_loading_frames=False, | |
| ) | |
| try: | |
| for name_index, name in enumerate(names): | |
| prompt = prompts[name] | |
| seed = int(prompt["frame"]) | |
| if not 0 <= seed < len(files): | |
| raise ValueError(f"Invalid seed frame for {name}") | |
| x0, y0, x1, y1 = prompt["box"] | |
| roi = np.zeros((height, width), dtype=bool) | |
| roi[int(y0) : int(y1), int(x0) : int(x1)] = True | |
| box = [ | |
| [x0 / width, y0 / height, (x1 - x0) / width, (y1 - y0) / height] | |
| ] | |
| for reverse in (False, True): | |
| self.model.reset_state(state) | |
| _, output = self.model.add_prompt( | |
| state, | |
| seed, | |
| text_str=prompt["text"], | |
| boxes_xywh=box, | |
| box_labels=[1], | |
| ) | |
| object_ids = output["out_obj_ids"] | |
| if not len(object_ids): | |
| raise RuntimeError( | |
| f"SAM 3.1 did not detect {name} at the seed frame" | |
| ) | |
| overlaps = [ | |
| np.count_nonzero(mask & roi) | |
| / max(np.count_nonzero(mask | roi), 1) | |
| for mask in output["out_binary_masks"] | |
| ] | |
| index = int(np.argmax(overlaps)) | |
| if overlaps[index] <= 0: | |
| raise RuntimeError( | |
| f"SAM 3.1 detections do not overlap the seed for {name}" | |
| ) | |
| selected_id = int(object_ids[index]) | |
| for frame, output in self.model.propagate_in_video( | |
| state, | |
| start_frame_idx=seed, | |
| reverse=reverse, | |
| ): | |
| ids = np.asarray(output["out_obj_ids"]) | |
| match = np.flatnonzero(ids == selected_id) | |
| if len(match) and ( | |
| (reverse and frame < seed) | |
| or (not reverse and frame >= seed) | |
| ): | |
| index = int(match[0]) | |
| masks[name_index, frame] = output["out_binary_masks"][ | |
| index | |
| ] | |
| confidence[name_index, frame] = float( | |
| output["out_probs"][index] | |
| ) | |
| finally: | |
| self.model.reset_state(state) | |
| winner = (masks * confidence[:, :, None, None]).argmax(axis=0) | |
| masks &= winner[None] == np.arange(len(names))[:, None, None, None] | |
| observed = masks.sum(axis=(2, 3)) >= 16 | |
| return names, masks, confidence, observed | |
Xet Storage Details
- Size:
- 4.67 kB
- Xet hash:
- fa3c66294d40da77ce53e5bbbd37fdd0e99484ad81b7617502bfafd5b06d2ecf
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.