File size: 4,614 Bytes
590304f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 | #!/usr/bin/env python3
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 # noqa: E402
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
|