| |
| from __future__ import annotations |
|
|
| import json |
| import math |
| import os |
| import sys |
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
| import torchvision.transforms.functional as TF |
| from PIL import Image |
|
|
| SCRIPT_DIR = Path(__file__).resolve().parent |
| BUNDLE_ROOT = SCRIPT_DIR.parent |
| SAM3_REPO = Path(os.environ.get("SAM3_REPO", BUNDLE_ROOT / "runtime" / "sam3_repo")) |
| for path in (SCRIPT_DIR, SAM3_REPO): |
| if str(path) not in sys.path: |
| sys.path.insert(0, str(path)) |
|
|
| from sam3_decoder_experiment_lib import SAM3FeatureModel, resize_pad_image_and_mask |
|
|
|
|
| def bbox_xyxy_to_padded_cxcywh(box, orig_w: int, orig_h: int, image_size: int): |
| x1, y1, x2, y2 = [float(v) for v in box] |
| scale = float(image_size) / float(max(orig_h, orig_w)) |
| new_w = orig_w * scale |
| new_h = orig_h * scale |
| left = (image_size - new_w) / 2.0 |
| top = (image_size - new_h) / 2.0 |
| x1p, x2p = x1 * scale + left, x2 * scale + left |
| y1p, y2p = y1 * scale + top, y2 * scale + top |
| cx = (x1p + x2p) / 2.0 / image_size |
| cy = (y1p + y2p) / 2.0 / image_size |
| w = max(1.0, x2p - x1p) / image_size |
| h = max(1.0, y2p - y1p) / image_size |
| return torch.tensor([[[cx, cy, w, h]]], dtype=torch.float32) |
|
|
|
|
| def unpad_resize_prediction(prob_512: np.ndarray, orig_w: int, orig_h: int, image_size: int): |
| scale = float(image_size) / float(max(orig_h, orig_w)) |
| new_w = max(1, int(round(orig_w * scale))) |
| new_h = max(1, int(round(orig_h * scale))) |
| left = (image_size - new_w) // 2 |
| top = (image_size - new_h) // 2 |
| crop = prob_512[top : top + new_h, left : left + new_w] |
| return np.asarray(Image.fromarray(crop).resize((orig_w, orig_h), Image.BILINEAR)) |
|
|
|
|
| class SAM3BuscotPredictor: |
| def __init__( |
| self, |
| sam3_checkpoint: str, |
| checkpoint_path: str | None = None, |
| prompt_type: str = "semantic_text", |
| prompt_text: str = "breast tumor", |
| image_size: int = 512, |
| device: str = "cuda", |
| encoder_trainable: str = "frozen", |
| lora_rank: int = 8, |
| lora_alpha: float = 16, |
| threshold: float = 0.5, |
| ): |
| self.image_size = int(image_size) |
| self.threshold = float(threshold) |
| self.device = torch.device(device if device == "cuda" and torch.cuda.is_available() else "cpu") |
| self.model = SAM3FeatureModel( |
| sam3_checkpoint, |
| image_size=self.image_size, |
| encoder_trainable=encoder_trainable, |
| decoder_name="sam3_native", |
| prompt_type=prompt_type, |
| prompt_text=prompt_text, |
| lora_rank=lora_rank, |
| lora_alpha=lora_alpha, |
| ).to(self.device) |
| if checkpoint_path: |
| ckpt = torch.load(checkpoint_path, map_location=self.device) |
| state = ckpt["model"] if isinstance(ckpt, dict) and "model" in ckpt else ckpt |
| cleaned = {} |
| for k, v in state.items(): |
| for prefix in ("module.", "model."): |
| if k.startswith(prefix): |
| k = k[len(prefix) :] |
| cleaned[k] = v |
| self.model.load_state_dict(cleaned, strict=False) |
| self.model.eval() |
|
|
| @torch.no_grad() |
| def predict(self, image_path: str, bbox_xyxy=None): |
| image = Image.open(image_path).convert("RGB") |
| orig_w, orig_h = image.size |
| dummy = np.zeros((orig_h, orig_w), dtype=np.uint8) |
| padded, _ = resize_pad_image_and_mask(image, dummy, self.image_size) |
| x = TF.to_tensor(padded).unsqueeze(0).to(self.device) |
| bbox_prompt = None |
| if bbox_xyxy is not None: |
| bbox_prompt = bbox_xyxy_to_padded_cxcywh(bbox_xyxy, orig_w, orig_h, self.image_size).to(self.device) |
| with torch.amp.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=self.device.type == "cuda"): |
| raw = self.model(x, return_details=True, bbox_prompt=bbox_prompt) |
| logits = raw["mask_logits"] if isinstance(raw, dict) else raw |
| prob_512 = torch.sigmoid(logits)[0, 0].detach().float().cpu().numpy() |
| prob = unpad_resize_prediction(prob_512, orig_w, orig_h, self.image_size) |
| pred = prob >= self.threshold |
| details = {} |
| if isinstance(raw, dict): |
| for key, value in raw.items(): |
| if key == "mask_logits": |
| continue |
| try: |
| details[key] = float(torch.as_tensor(value).detach().flatten()[0].cpu()) |
| except Exception: |
| details[key] = str(value) |
| return pred.astype(np.uint8), details |
|
|
