Run Python SAM2 inference for one variant.
(variant_name, config, checkpoint, image_path, point_x, point_y, dump_dir)
| 39 | |
| 40 | |
| 41 | def run_python_variant(variant_name, config, checkpoint, image_path, point_x, point_y, dump_dir): |
| 42 | """Run Python SAM2 inference for one variant.""" |
| 43 | from sam2.build_sam import build_sam2 |
| 44 | from sam2.sam2_image_predictor import SAM2ImagePredictor |
| 45 | |
| 46 | os.makedirs(dump_dir, exist_ok=True) |
| 47 | |
| 48 | print(f"\n{'='*60}") |
| 49 | print(f"Python: {variant_name}") |
| 50 | print(f"{'='*60}") |
| 51 | |
| 52 | t0 = time.time() |
| 53 | model = build_sam2(config, checkpoint, device="cpu") |
| 54 | model = model.float().eval() |
| 55 | predictor = SAM2ImagePredictor(model) |
| 56 | t_load = time.time() - t0 |
| 57 | |
| 58 | image = np.array(Image.open(image_path).convert("RGB")) |
| 59 | img_h, img_w = image.shape[:2] |
| 60 | |
| 61 | t0 = time.time() |
| 62 | with torch.no_grad(): |
| 63 | predictor.set_image(image) |
| 64 | t_encode = time.time() - t0 |
| 65 | |
| 66 | point_coords = np.array([[point_x, point_y]], dtype=np.float32) |
| 67 | point_labels = np.array([1], dtype=np.int32) |
| 68 | |
| 69 | t0 = time.time() |
| 70 | with torch.no_grad(): |
| 71 | masks, scores, logits = predictor.predict( |
| 72 | point_coords=point_coords, |
| 73 | point_labels=point_labels, |
| 74 | multimask_output=True, |
| 75 | normalize_coords=True, |
| 76 | ) |
| 77 | t_predict = time.time() - t0 |
| 78 | |
| 79 | # Save preprocessed image for C++ consumption |
| 80 | transforms = predictor._transforms |
| 81 | img_tensor = torch.from_numpy(image).permute(2, 0, 1).float() / 255.0 |
| 82 | from torchvision.transforms.functional import resize, normalize |
| 83 | img_resized = resize(img_tensor, [model.image_size, model.image_size], antialias=True) |
| 84 | mean = torch.tensor([0.485, 0.456, 0.406]).view(3, 1, 1) |
| 85 | std = torch.tensor([0.229, 0.224, 0.225]).view(3, 1, 1) |
| 86 | img_normalized = (img_resized - mean) / std |
| 87 | preproc_path = os.path.join(dump_dir, "preprocessed.bin") |
| 88 | img_normalized.numpy().astype(np.float32).tofile(preproc_path) |
| 89 | |
| 90 | # Save masks and scores |
| 91 | logits.astype(np.float32).tofile(os.path.join(dump_dir, "logits.bin")) |
| 92 | scores.astype(np.float32).tofile(os.path.join(dump_dir, "scores.bin")) |
| 93 | masks.astype(np.uint8).tofile(os.path.join(dump_dir, "masks.bin")) |
| 94 | |
| 95 | # Save metadata |
| 96 | meta = { |
| 97 | "variant": variant_name, "img_w": img_w, "img_h": img_h, |
| 98 | "point_x": point_x, "point_y": point_y, |