Buckets:
| import argparse | |
| import json | |
| from pathlib import Path | |
| import cv2 | |
| import numpy as np | |
| from real2sim import video | |
| from real2sim.tracking.masks import SAM31Tracker | |
| def main(): | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument("folder", type=Path) | |
| parser.add_argument("--prompts", type=Path, required=True) | |
| parser.add_argument( | |
| "--checkpoint", type=Path, default=Path("data/sam3/sam3.1_multiplex.pt") | |
| ) | |
| parser.add_argument("--device", default="cuda") | |
| args = parser.parse_args() | |
| prompts = json.loads(args.prompts.read_text()) | |
| with np.load(args.folder / "sim_trajectory.npz") as data: | |
| indices, timestamps = data["frame_indices"], data["timestamps"] | |
| replay = video.read(args.folder / "mujoco_replay.mp4") | |
| width = replay.shape[2] // 3 | |
| frames = replay[:, :, :width] | |
| folder = args.folder / "tracking_frames" | |
| folder.mkdir(exist_ok=True) | |
| expected = {f"{i:06d}.jpg" for i in range(len(frames))} | |
| if any(p.name not in expected for p in folder.glob("*.jpg")): | |
| raise ValueError("Tracking frames directory contains unexpected frames") | |
| for i, frame in enumerate(frames): | |
| if not cv2.imwrite( | |
| str(folder / f"{i:06d}.jpg"), cv2.cvtColor(frame, cv2.COLOR_RGB2BGR) | |
| ): | |
| raise OSError("Cannot write tracking frame") | |
| tracker = SAM31Tracker(args.checkpoint, args.device) | |
| names, masks, confidence, observed = tracker.track(folder, prompts) | |
| np.savez_compressed( | |
| args.folder / "real_masks.npz", | |
| masks=masks, | |
| names=names, | |
| confidence=confidence, | |
| observed=observed, | |
| model="sam3.1_multiplex", | |
| frame_indices=indices, | |
| timestamps=timestamps, | |
| ) | |
| palette = np.asarray( | |
| [[65, 210, 120], [70, 145, 255], [250, 80, 165]], dtype=np.uint8 | |
| ) | |
| overlays = frames.astype(np.float32) | |
| for index, mask in enumerate(masks): | |
| overlays[mask] = 0.5 * overlays[mask] + 0.5 * palette[index % len(palette)] | |
| step = float(np.median(np.diff(timestamps))) | |
| video.write(args.folder / "mask_overlay.mp4", overlays.astype(np.uint8), 1 / step) | |
| cv2.imwrite( | |
| str(args.folder / "mask_preview.png"), | |
| cv2.cvtColor(overlays[len(frames) // 2].astype(np.uint8), cv2.COLOR_RGB2BGR), | |
| ) | |
| metrics = { | |
| name: { | |
| "observed_frames": int(observed[i].sum()), | |
| "median_mask_pixels": float(np.median(masks[i].sum(axis=(1, 2)))), | |
| "mean_tracking_score": float(confidence[i][observed[i]].mean()), | |
| } | |
| for i, name in enumerate(names) | |
| } | |
| (args.folder / "tracking_metrics.json").write_text( | |
| json.dumps(metrics, indent=2) + "\n" | |
| ) | |
| print(json.dumps(metrics)) | |
| if __name__ == "__main__": | |
| main() | |
Xet Storage Details
- Size:
- 2.76 kB
- Xet hash:
- cb929c470905628b6b857542f247c6f3b83d6b75e29e6da471526946f661fb18
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.