import csv import inspect import json import math import os import random import sys from pathlib import Path from types import SimpleNamespace import cv2 import numpy as np import pandas as pd from PIL import Image, ImageDraw, ImageFont import torch import torch.nn as nn import torch.nn.functional as F from torch.utils.data import Dataset import torchvision.transforms.functional as TF SCRIPT_DIR = Path(__file__).resolve().parent BUNDLE_ROOT = SCRIPT_DIR.parent SAM3_REPO = Path(os.environ.get("SAM3_REPO", BUNDLE_ROOT / "runtime" / "sam3_repo")) if str(SAM3_REPO) not in sys.path: sys.path.insert(0, str(SAM3_REPO)) if not torch.cuda.is_available() and not getattr(F.linear, "_sam3_cpu_dtype_patch", False): _orig_linear = F.linear _orig_conv2d = F.conv2d _orig_pin_memory = torch.Tensor.pin_memory def _linear_dtype_safe(input, weight, bias=None): if torch.is_tensor(input) and torch.is_tensor(weight) and input.dtype != weight.dtype: input = input.to(weight.dtype) if bias is not None and bias.dtype != weight.dtype: bias = bias.to(weight.dtype) return _orig_linear(input, weight, bias) def _conv2d_dtype_safe(input, weight, bias=None, stride=1, padding=0, dilation=1, groups=1): if torch.is_tensor(input) and torch.is_tensor(weight) and input.dtype != weight.dtype: input = input.to(weight.dtype) if bias is not None and bias.dtype != weight.dtype: bias = bias.to(weight.dtype) return _orig_conv2d(input, weight, bias, stride, padding, dilation, groups) _linear_dtype_safe._sam3_cpu_dtype_patch = True F.linear = _linear_dtype_safe F.conv2d = _conv2d_dtype_safe torch.Tensor.pin_memory = lambda self, *args, **kwargs: self torch.Tensor.pin_memory._sam3_cpu_dtype_patch = True def seed_everything(seed): random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed_all(seed) def read_image_rgb(path): return Image.open(path).convert("RGB") def read_mask_binary(path): arr = np.asarray(Image.open(path).convert("L")) return (arr > 0).astype(np.uint8) def resize_pad_image_and_mask(img, mask, out_size): w, h = img.size scale = float(out_size) / float(max(h, w)) new_h = max(1, int(round(h * scale))) new_w = max(1, int(round(w * scale))) img_r = img.resize((new_w, new_h), Image.BILINEAR) mask_r = Image.fromarray((mask * 255).astype(np.uint8)).resize((new_w, new_h), Image.NEAREST) canvas = Image.new("RGB", (out_size, out_size), (0, 0, 0)) mask_canvas = Image.new("L", (out_size, out_size), 0) left = (out_size - new_w) // 2 top = (out_size - new_h) // 2 canvas.paste(img_r, (left, top)) mask_canvas.paste(mask_r, (left, top)) return canvas, (np.asarray(mask_canvas) > 0).astype(np.uint8) def bbox_from_mask_tensor(mask): mask_np = mask.detach().cpu().numpy() > 0.5 boxes = [] for m in mask_np: yy, xx = np.where(m[0]) if len(xx) == 0: boxes.append([0.0, 0.0, 1.0, 1.0]) else: boxes.append([float(xx.min()), float(yy.min()), float(xx.max()), float(yy.max())]) return torch.tensor(boxes, dtype=torch.float32, device=mask.device) def xyxy_to_normalized_cxcywh(boxes_xyxy, image_size): boxes = boxes_xyxy.float() x1, y1, x2, y2 = boxes.unbind(-1) w = (x2 - x1 + 1.0).clamp(min=1.0) h = (y2 - y1 + 1.0).clamp(min=1.0) cx = x1 + 0.5 * w cy = y1 + 0.5 * h scale = float(image_size) out = torch.stack([cx / scale, cy / scale, w / scale, h / scale], dim=-1) return out.clamp(0.0, 1.0) def bbox_prompt_from_mask_tensor(mask, image_size): boxes_xyxy = bbox_from_mask_tensor(mask) return xyxy_to_normalized_cxcywh(boxes_xyxy, image_size) class PublicSegmentationDataset(Dataset): def __init__(self, df, image_size=512, augment=False): self.df = df.reset_index(drop=True).copy() self.image_size = int(image_size) self.augment = bool(augment) def __len__(self): return len(self.df) def __getitem__(self, idx): row = self.df.iloc[idx] img = read_image_rgb(row["image_path"]) mask = read_mask_binary(row["mask_path"]) img, mask = resize_pad_image_and_mask(img, mask, self.image_size) img_t = TF.to_tensor(img) mask_t = torch.from_numpy(mask).unsqueeze(0).float() if self.augment: img_t, mask_t = self.apply_aug(img_t, mask_t) return { "image": img_t, "mask": mask_t, "dataset": str(row["dataset"]), "label": str(row["label"]), "label_id": int(row["label_id"]), "case_id": str(row["new_case_id"]), "image_name": str(row["image_name"]), "mask_name": str(row["mask_name"]), "image_path": str(row["image_path"]), "mask_path": str(row["mask_path"]), "mask_bbox_xyxy": str(row.get("mask_bbox_xyxy", "")), "prompt": str(row.get("prompt", "")), } def apply_aug(self, img_t, mask_t): if random.random() < 0.5: img_t = TF.hflip(img_t) mask_t = TF.hflip(mask_t) if random.random() < 0.35: angle = random.uniform(-10.0, 10.0) img_t = TF.rotate(img_t, angle, interpolation=TF.InterpolationMode.BILINEAR, fill=0) mask_t = TF.rotate(mask_t, angle, interpolation=TF.InterpolationMode.NEAREST, fill=0) if random.random() < 0.45: img_t = TF.adjust_brightness(img_t, random.uniform(0.85, 1.15)) img_t = TF.adjust_contrast(img_t, random.uniform(0.85, 1.15)) mask_t = (mask_t > 0.5).float() return img_t.clamp(0, 1), mask_t def subset_limit(df, max_samples, seed): if max_samples is None or int(max_samples) <= 0 or len(df) <= int(max_samples): return df.reset_index(drop=True) return df.sample(n=int(max_samples), random_state=seed).reset_index(drop=True) def case_group(row): return f"{row['dataset']}::{row['new_case_id']}" def parse_dataset_list(value): if value is None: return [] if isinstance(value, (list, tuple)): return [str(v) for v in value if str(v)] return [v.strip() for v in str(value).split(",") if v.strip()] def split_dataset( df, protocol, heldout_dataset=None, train_datasets=None, test_dataset=None, seed=42, max_train_samples=None, max_val_samples=None, max_test_samples=None, ): df = df.copy() df["split"] = df["split"].astype(str).str.lower() if protocol == "all_public_indomain": train = df[df["split"] == "train"] val = df[df["split"] == "val"] test = df[df["split"] == "test"] elif protocol == "leave_one_dataset_out": if not heldout_dataset: raise ValueError("--heldout_dataset is required for leave_one_dataset_out") test = df[df["dataset"].astype(str) == str(heldout_dataset)] pool = df[df["dataset"].astype(str) != str(heldout_dataset)] train = pool[pool["split"] == "train"].copy() val = pool[pool["split"] == "val"].copy() if len(val) < max(8, int(0.05 * len(pool))): train, val = make_group_val_split(pool, seed=seed) elif protocol == "source_to_target": train_dataset_list = parse_dataset_list(train_datasets) target_dataset = test_dataset or heldout_dataset if not train_dataset_list: raise ValueError("--train_datasets is required for source_to_target") if not target_dataset: raise ValueError("--test_dataset is required for source_to_target") pool = df[df["dataset"].astype(str).isin(train_dataset_list)].copy() test = df[df["dataset"].astype(str) == str(target_dataset)].copy() train = pool[pool["split"] == "train"].copy() val = pool[pool["split"] == "val"].copy() if len(val) < max(8, int(0.05 * len(pool))): train, val = make_group_val_split(pool, seed=seed) else: raise ValueError(f"Unknown protocol: {protocol}") train = subset_limit(train, max_train_samples, seed) val = subset_limit(val, max_val_samples, seed + 1) test = subset_limit(test, max_test_samples, seed + 2) if len(train) == 0 or len(val) == 0 or len(test) == 0: raise ValueError(f"Empty split: train={len(train)}, val={len(val)}, test={len(test)}") return train, val, test def make_group_val_split(pool, seed=42, val_frac=0.15): rng = random.Random(seed) groups = [] for group_id, g in pool.groupby(pool.apply(case_group, axis=1)): labels = g["label"].astype(str).value_counts().to_dict() dominant = "malignant" if labels.get("malignant", 0) >= labels.get("benign", 0) else "benign" groups.append((group_id, dominant, len(g))) rng.shuffle(groups) target = max(1, int(round(len(pool) * val_frac))) val_groups = set() counts = {"benign": 0, "malignant": 0} for label in ["benign", "malignant"]: for group_id, dominant, size in groups: if dominant == label and group_id not in val_groups: val_groups.add(group_id) counts[label] += size break current = sum(size for group_id, _, size in groups if group_id in val_groups) for group_id, _, size in groups: if current >= target: break if group_id not in val_groups: val_groups.add(group_id) current += size group_series = pool.apply(case_group, axis=1) val = pool[group_series.isin(val_groups)].copy() train = pool[~group_series.isin(val_groups)].copy() return train, val def call_with_supported_kwargs(func, **kwargs): sig = inspect.signature(func) return func(**{k: v for k, v in kwargs.items() if k in sig.parameters}) class TorchCudaToCpuPatch: def __init__(self): self.active = not torch.cuda.is_available() self.originals = {} def __enter__(self): if not self.active: return self for name in ["zeros", "ones", "empty", "full", "arange", "tensor"]: fn = getattr(torch, name) self.originals[name] = fn def make_wrapper(orig): def wrapper(*args, **kwargs): if str(kwargs.get("device", "")) == "cuda": kwargs["device"] = "cpu" return orig(*args, **kwargs) return wrapper setattr(torch, name, make_wrapper(fn)) self.originals["tensor_cuda"] = torch.Tensor.cuda torch.Tensor.cuda = lambda self, *args, **kwargs: self return self def __exit__(self, exc_type, exc, tb): if not self.active: return False for name, fn in self.originals.items(): if name == "tensor_cuda": torch.Tensor.cuda = fn else: setattr(torch, name, fn) return False def patch_sam3_rope(): try: import sam3.model.vitdet as vitdet except Exception: return if getattr(vitdet.apply_rotary_enc, "_sam3_public_patch", False): return original_apply_rotary_enc = vitdet.apply_rotary_enc def patched_apply_rotary_enc(xq, xk, freqs_cis): target_l = xq.shape[-2] current_l = freqs_cis.shape[0] * freqs_cis.shape[1] if freqs_cis.dim() == 3 else freqs_cis.shape[0] side_old = freqs_cis.shape[0] if freqs_cis.dim() == 3 else int(math.sqrt(current_l)) if current_l != target_l: side_new = int(math.sqrt(target_l)) d = freqs_cis.shape[-1] if freqs_cis.dim() == 2: real = freqs_cis.real.view(1, side_old, side_old, d).permute(0, 3, 1, 2) imag = freqs_cis.imag.view(1, side_old, side_old, d).permute(0, 3, 1, 2) else: real = freqs_cis.real.permute(2, 0, 1).unsqueeze(0) imag = freqs_cis.imag.permute(2, 0, 1).unsqueeze(0) real = F.interpolate(real, size=(side_new, side_new), mode="bicubic", align_corners=False) imag = F.interpolate(imag, size=(side_new, side_new), mode="bicubic", align_corners=False) if freqs_cis.dim() == 2: freqs_cis = torch.complex(real, imag).permute(0, 2, 3, 1).reshape(target_l, d).to(xq.device) else: freqs_cis = torch.complex(real, imag).squeeze(0).permute(1, 2, 0).to(xq.device) return original_apply_rotary_enc(xq, xk, freqs_cis) patched_apply_rotary_enc._sam3_public_patch = True vitdet.apply_rotary_enc = patched_apply_rotary_enc def patch_sam3_fused_for_grad(): """Use a regular linear+activation fallback when SAM3 fused MLP sees grad mode. SAM3's inference fused addmm path intentionally rejects grad-enabled forward. LoRA on the image encoder makes later frozen blocks receive grad-enabled tensors, so this fallback is needed for encoder LoRA smoke/training. """ try: import sam3.perflib.fused as fused import sam3.model.vitdet as vitdet except Exception: return if getattr(fused.addmm_act, "_sam3_grad_fallback_patch", False): return original_addmm_act = fused.addmm_act def patched_addmm_act(activation, linear, mat1): if not torch.is_grad_enabled(): return original_addmm_act(activation, linear, mat1) out = F.linear(mat1, linear.weight, linear.bias) if activation in [torch.nn.functional.relu, torch.nn.ReLU]: return F.relu(out) if activation in [torch.nn.functional.gelu, torch.nn.GELU]: return F.gelu(out) raise ValueError(f"Unexpected activation {activation}") patched_addmm_act._sam3_grad_fallback_patch = True fused.addmm_act = patched_addmm_act vitdet.addmm_act = patched_addmm_act def fix_backbone_rope(backbone, image_size): side_new = max(1, image_size // 14) for module in backbone.modules(): if hasattr(module, "freqs_cis") and module.freqs_cis is not None: old_f = module.freqs_cis d = old_f.shape[-1] side_old = int(np.sqrt(old_f.shape[0])) if old_f.dim() == 2 else old_f.shape[0] if old_f.dim() == 2: view_f = old_f.view(side_old, side_old, d).permute(2, 0, 1).unsqueeze(0) else: view_f = old_f.permute(2, 0, 1).unsqueeze(0) f_r = F.interpolate(view_f.real, size=(side_new, side_new), mode="bicubic", align_corners=False) f_i = F.interpolate(view_f.imag, size=(side_new, side_new), mode="bicubic", align_corners=False) new_f = torch.complex(f_r, f_i).squeeze(0).permute(1, 2, 0) if old_f.dim() == 2: new_f = new_f.reshape(side_new * side_new, d) del module.freqs_cis module.freqs_cis = new_f.to(old_f.device).detach() class LoRALinear(nn.Module): def __init__(self, original_linear, r=8, alpha=16): super().__init__() self.linear = original_linear self.r = int(r) self.alpha = float(alpha) self.scaling = self.alpha / max(1, self.r) for p in self.linear.parameters(): p.requires_grad = False device = original_linear.weight.device dtype = original_linear.weight.dtype self.lora_A = nn.Parameter(torch.zeros(original_linear.in_features, self.r, device=device, dtype=dtype)) self.lora_B = nn.Parameter(torch.zeros(self.r, original_linear.out_features, device=device, dtype=dtype)) nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5)) nn.init.zeros_(self.lora_B) @property def weight(self): return self.linear.weight @property def bias(self): return self.linear.bias def forward(self, x): return self.linear(x) + (x @ self.lora_A @ self.lora_B) * self.scaling def inject_lora_to_module(module, target_keywords, r=8, alpha=16): targets = [] for name, sub_module in module.named_modules(): if isinstance(sub_module, nn.Linear) and any(k in name for k in target_keywords): targets.append((name, sub_module)) for name, sub_module in targets: parent = module parts = name.split(".") for part in parts[:-1]: parent = getattr(parent, part) setattr(parent, parts[-1], LoRALinear(sub_module, r=r, alpha=alpha).to(sub_module.weight.device)) return len(targets) class SAM3FeatureModel(nn.Module): def __init__( self, checkpoint_path, image_size=512, encoder_trainable="frozen", decoder_name="cnn", prompt_type="none", prompt_text="breast tumor", lora_rank=8, lora_alpha=16, train_native_head=True, ): super().__init__() from sam3 import build_sam3_image_model patch_sam3_rope() if encoder_trainable == "lora": patch_sam3_fused_for_grad() self.image_size = int(image_size) self.decoder_name = decoder_name self.prompt_type = "semantic_text" if prompt_type == "text" else prompt_type self.prompt_text = "" if self.prompt_type == "none" else prompt_text self.encoder_trainable = encoder_trainable self.encoder_forward_no_grad = False self.lora_rank = int(lora_rank) self.lora_alpha = float(lora_alpha) self.train_native_head = bool(train_native_head) self.encoder_lora_layers = 0 self.decoder_lora_layers = 0 with TorchCudaToCpuPatch(): self.sam = build_sam3_image_model(checkpoint_path=checkpoint_path, load_from_HF=False) self.backbone = self.sam.backbone fix_backbone_rope(self.backbone, self.image_size) self.feature_meta = self._infer_feature_meta() for p in self.sam.parameters(): p.requires_grad = False if encoder_trainable == "frozen": for p in self.backbone.parameters(): p.requires_grad = False self.backbone.eval() self.encoder_forward_no_grad = True elif encoder_trainable == "lora": if decoder_name != "sam3_native": raise ValueError("encoder_trainable=lora is currently supported only with decoder=sam3_native") for p in self.sam.parameters(): p.requires_grad = False self.encoder_lora_layers = inject_lora_to_module( self.backbone, ["qkv", "q_proj", "v_proj"], r=self.lora_rank, alpha=self.lora_alpha, ) transformer = getattr(self.sam, "transformer", None) if transformer is not None: self.decoder_lora_layers = inject_lora_to_module( transformer, ["cross_attn", "ca_text", "ca_image", "self_attn", "q_proj", "v_proj"], r=self.lora_rank, alpha=self.lora_alpha, ) self.encoder_forward_no_grad = False elif encoder_trainable == "last_block": print("WARNING: SAM3 last-block fine-tuning is currently disabled because SAM3 fused ops require grad disabled.") self.encoder_trainable = "frozen" for p in self.backbone.parameters(): p.requires_grad = False self.backbone.eval() self.encoder_forward_no_grad = True else: raise ValueError("encoder_trainable must be frozen, last_block, or lora") if decoder_name == "sam3_native": self.decoder = SAM3NativeDecoder( self.sam, self.image_size, prompt_type=self.prompt_type, prompt_text=self.prompt_text, encoder_forward_no_grad=self.encoder_forward_no_grad, backbone_autocast=self._backbone_autocast, ) if self.encoder_trainable == "lora": if self.train_native_head: for name, p in self.sam.named_parameters(): if "lora_" not in name and any(k in name for k in ["segmentation_head", "mask_head", "class_embed", "bbox_embed"]): p.requires_grad = True else: for name, p in self.sam.named_parameters(): if not name.startswith("backbone."): p.requires_grad = True else: in_dims = self.feature_meta["channels"] if decoder_name == "cnn": self.decoder = CNNDecoder(in_dims[-1]) elif decoder_name == "unet": self.decoder = UNetStyleDecoder(in_dims) elif decoder_name == "segformer": self.decoder = SegFormerStyleDecoder(in_dims) else: raise ValueError(f"Unknown decoder: {decoder_name}") self._log_trainable_summary() def train(self, mode=True): super().train(mode) if self.encoder_forward_no_grad: self.backbone.eval() return self @staticmethod def _count_params(module, trainable_only=False): params = module.parameters() if trainable_only: params = (p for p in params if p.requires_grad) return sum(p.numel() for p in params) def _log_trainable_summary(self): backbone_trainable = self._count_params(self.backbone, trainable_only=True) if self.decoder_name == "sam3_native": decoder_trainable = sum( p.numel() for name, p in self.sam.named_parameters() if p.requires_grad and not name.startswith("backbone.") ) else: decoder_trainable = self._count_params(self.decoder, trainable_only=True) print(f"encoder_trainable: {self.encoder_trainable}") print(f"backbone trainable params: {backbone_trainable}") print(f"decoder trainable params: {decoder_trainable}") print(f"encoder forward uses no_grad: {self.encoder_forward_no_grad}") if self.encoder_trainable == "lora": lora_params = sum(p.numel() for name, p in self.named_parameters() if "lora_" in name and p.requires_grad) native_head_params = sum(p.numel() for name, p in self.named_parameters() if "lora_" not in name and p.requires_grad) print(f"encoder LoRA layers: {self.encoder_lora_layers}") print(f"native decoder LoRA layers: {self.decoder_lora_layers}") print(f"LoRA trainable params: {lora_params}") print(f"non-LoRA trainable params: {native_head_params}") def _call_backbone(self, x): if hasattr(self.backbone, "forward_image"): return self.backbone.forward_image(x) try: return self.backbone(x, captions=["tumor"] * x.shape[0]) except TypeError: return self.backbone(x) def _backbone_autocast(self, device): return torch.autocast( device_type="cuda", dtype=torch.bfloat16, enabled=(device.type == "cuda"), ) def _infer_feature_meta(self): was_training = self.backbone.training self.backbone.eval() device = next(self.backbone.parameters()).device with torch.no_grad(): with self._backbone_autocast(device): out = self._call_backbone(torch.zeros(1, 3, self.image_size, self.image_size, device=device)) features = normalize_encoder_features(out) if was_training: self.backbone.train() return {"channels": [int(f.shape[1]) for f in features], "strides": [self.image_size // int(f.shape[-1]) for f in features]} def extract_features(self, x): x = F.interpolate(x, size=(self.image_size, self.image_size), mode="bilinear", align_corners=False) device = x.device if not self.encoder_forward_no_grad: with self._backbone_autocast(device): out = self._call_backbone(x) else: with torch.no_grad(): with self._backbone_autocast(device): out = self._call_backbone(x) return [feature.float() for feature in normalize_encoder_features(out)] def forward(self, x, mask=None, return_details=False, bbox_prompt=None, prompt_texts=None): x3 = x.repeat(1, 3, 1, 1) if x.shape[1] == 1 else x if self.decoder_name == "sam3_native": if self.prompt_type == "gt_bbox" and bbox_prompt is None: if mask is None: raise ValueError("prompt_type=gt_bbox requires mask or bbox_prompt") bbox_prompt = bbox_prompt_from_mask_tensor(mask, self.image_size) return self.decoder( x3, output_size=x.shape[-2:], return_details=return_details, bbox_prompt=bbox_prompt, prompt_texts=prompt_texts, ) features = self.extract_features(x3) logits = self.decoder(features) logits = F.interpolate(logits, size=x.shape[-2:], mode="bilinear", align_corners=False) if return_details: return {"mask_logits": logits} return logits def normalize_encoder_features(out): if isinstance(out, dict): for key in ["vision_features", "backbone_fpn", "features", "feature_maps", "multiscale_features"]: if key in out: out = out[key] break else: tensors = [v for v in out.values() if torch.is_tensor(v) and v.dim() == 4] if tensors: out = tensors if torch.is_tensor(out): if out.dim() == 3: b, n, c = out.shape side = int(math.sqrt(n)) out = out.transpose(1, 2).reshape(b, c, side, side) return make_pyramid_from_single(out) if isinstance(out, (list, tuple)): feats = [] for item in out: if torch.is_tensor(item): f = item if f.dim() == 3: b, n, c = f.shape side = int(math.sqrt(n)) f = f.transpose(1, 2).reshape(b, c, side, side) if f.dim() == 4: feats.append(f) if feats: feats = sorted(feats, key=lambda t: t.shape[-1], reverse=True) return feats[:4] if len(feats) >= 4 else make_pyramid_from_single(feats[-1]) raise RuntimeError(f"Could not parse SAM3 backbone features from type {type(out)}") def make_pyramid_from_single(f): return [ F.interpolate(f, scale_factor=4, mode="bilinear", align_corners=False), F.interpolate(f, scale_factor=2, mode="bilinear", align_corners=False), f, F.avg_pool2d(f, kernel_size=2, stride=2), ] class CNNDecoder(nn.Module): def __init__(self, in_ch, mid=256): super().__init__() self.net = nn.Sequential( nn.Conv2d(in_ch, mid, 3, padding=1), nn.BatchNorm2d(mid), nn.ReLU(inplace=True), nn.Conv2d(mid, mid // 2, 3, padding=1), nn.BatchNorm2d(mid // 2), nn.ReLU(inplace=True), nn.Conv2d(mid // 2, 1, 1), ) def forward(self, features): return self.net(features[-1]) class UNetStyleDecoder(nn.Module): def __init__(self, in_dims, width=128): super().__init__() self.proj = nn.ModuleList([nn.Conv2d(c, width, 1) for c in in_dims]) self.fuse = nn.ModuleList([ConvBlock(width * 2, width) for _ in range(len(in_dims) - 1)]) self.head = nn.Conv2d(width, 1, 1) def forward(self, features): feats = [proj(f) for proj, f in zip(self.proj, features)] x = feats[-1] for i in range(len(feats) - 2, -1, -1): x = F.interpolate(x, size=feats[i].shape[-2:], mode="bilinear", align_corners=False) x = self.fuse[i](torch.cat([x, feats[i]], dim=1)) return self.head(x) class ConvBlock(nn.Module): def __init__(self, in_ch, out_ch): super().__init__() self.net = nn.Sequential( nn.Conv2d(in_ch, out_ch, 3, padding=1, bias=False), nn.BatchNorm2d(out_ch), nn.ReLU(inplace=True), nn.Conv2d(out_ch, out_ch, 3, padding=1, bias=False), nn.BatchNorm2d(out_ch), nn.ReLU(inplace=True), ) def forward(self, x): return self.net(x) class SegFormerStyleDecoder(nn.Module): def __init__(self, in_dims, embed_dim=128): super().__init__() self.proj = nn.ModuleList([nn.Conv2d(c, embed_dim, 1) for c in in_dims]) self.fuse = nn.Sequential( nn.Conv2d(embed_dim * len(in_dims), embed_dim, 1, bias=False), nn.BatchNorm2d(embed_dim), nn.ReLU(inplace=True), nn.Dropout2d(0.1), nn.Conv2d(embed_dim, 1, 1), ) def forward(self, features): mapped = [p(f) for p, f in zip(self.proj, features)] target = mapped[0].shape[-2:] mapped = [mapped[0]] + [F.interpolate(f, size=target, mode="bilinear", align_corners=False) for f in mapped[1:]] return self.fuse(torch.cat(mapped, dim=1)) class SAM3NativeDecoder(nn.Module): def __init__( self, sam, image_size, prompt_type="none", prompt_text="breast tumor", encoder_forward_no_grad=True, backbone_autocast=None, ): super().__init__() self.sam = sam self.backbone = sam.backbone self.image_size = image_size self.prompt_type = "semantic_text" if prompt_type == "text" else prompt_type self.prompt_text = "" if self.prompt_type == "none" else prompt_text self.encoder_forward_no_grad = encoder_forward_no_grad self.backbone_autocast = backbone_autocast def _geometry_dtype(self): norm = getattr(getattr(self.sam, "geometry_encoder", None), "img_pre_norm", None) weight = getattr(norm, "weight", None) return weight.dtype if torch.is_tensor(weight) else None def _cast_backbone_out_for_geometry(self, backbone_out): dtype = self._geometry_dtype() if dtype is None: return backbone_out out = dict(backbone_out) for key in ("backbone_fpn", "vision_pos_enc"): value = out.get(key) if isinstance(value, (list, tuple)): out[key] = [x.to(dtype=dtype) if torch.is_tensor(x) and x.dtype != dtype else x for x in value] elif torch.is_tensor(value) and value.dtype != dtype: out[key] = value.to(dtype=dtype) return out def forward(self, x, output_size, return_details=False, bbox_prompt=None, prompt_texts=None): b = x.shape[0] x = F.interpolate(x, size=(self.image_size, self.image_size), mode="bilinear", align_corners=False) device = x.device def autocast_ctx(): return self.backbone_autocast(device) if self.backbone_autocast is not None else torch.autocast( device_type="cuda", dtype=torch.bfloat16, enabled=(device.type == "cuda"), ) if self.encoder_forward_no_grad: with torch.no_grad(): with autocast_ctx(): backbone_out = self.backbone.forward_image(x) if hasattr(self.backbone, "forward_image") else self.backbone(x) else: with autocast_ctx(): backbone_out = self.backbone.forward_image(x) if hasattr(self.backbone, "forward_image") else self.backbone(x) if self.prompt_type == "none": prompt, prompt_mask = None, None elif self.prompt_type == "semantic_text": texts = list(prompt_texts) if prompt_texts is not None else [self.prompt_text] * b if len(texts) != b: raise ValueError(f"Expected {b} prompt texts, got {len(texts)}") with torch.no_grad(): with autocast_ctx(): text_out = ( self.backbone.forward_text(texts, device=x.device) if hasattr(self.backbone, "forward_text") else None ) prompt, prompt_mask = unpack_prompt(text_out, b, x.device) elif self.prompt_type == "gt_bbox": if bbox_prompt is None: raise ValueError("sam3_native prompt_type=gt_bbox requires bbox_prompt") bbox_prompt = bbox_prompt.to(device=x.device, dtype=torch.float32) if bbox_prompt.dim() == 2: bbox_prompt = bbox_prompt.unsqueeze(0) if bbox_prompt.shape != (1, b, 4): raise RuntimeError(f"Expected bbox_prompt shape (1, {b}, 4), got {tuple(bbox_prompt.shape)}") with torch.no_grad(): with autocast_ctx(): text_out = ( self.backbone.forward_text(["visual"] * b, device=x.device) if hasattr(self.backbone, "forward_text") else None ) if isinstance(text_out, dict): backbone_out.update(text_out) backbone_out = self._cast_backbone_out_for_geometry(backbone_out) from sam3.model.geometry_encoders import Prompt find_input = SimpleNamespace( img_ids=torch.arange(b, device=x.device, dtype=torch.long), text_ids=torch.arange(b, device=x.device, dtype=torch.long), ) geometric_prompt = Prompt( box_embeddings=bbox_prompt, box_mask=torch.zeros(b, 1, dtype=torch.bool, device=x.device), box_labels=torch.ones(1, b, dtype=torch.bool, device=x.device), ) prompt, prompt_mask, backbone_out = self.sam._encode_prompt( backbone_out, find_input, geometric_prompt, ) else: raise ValueError(f"Unsupported native prompt_type: {self.prompt_type}") if self.prompt_type != "gt_bbox": find_input = SimpleNamespace( img_ids=torch.arange(b, device=x.device, dtype=torch.long), image_size=(self.image_size, self.image_size), original_size=(self.image_size, self.image_size), ) with autocast_ctx(): enc_out = call_with_supported_kwargs( self.sam._run_encoder, backbone_out=backbone_out, find_input=find_input, prompt=prompt, prompt_mask=prompt_mask, ) encoder_out = enc_out[1] decoder_pkg = enc_out[2] out = decoder_pkg[0] if isinstance(decoder_pkg, (tuple, list)) and decoder_pkg and isinstance(decoder_pkg[0], dict) else decoder_pkg dec_out = call_with_supported_kwargs( self.sam._run_decoder, memory=encoder_out.get("encoder_hidden_states"), pos_embed=encoder_out.get("pos_embed"), level_start_index=encoder_out.get("level_start_index"), spatial_shapes=encoder_out.get("spatial_shapes"), valid_ratios=encoder_out.get("valid_ratios"), src_mask=encoder_out.get("padding_mask"), out=out, encoder_out=encoder_out, prompt=prompt, prompt_mask=prompt_mask, memory_text=prompt, text_attention_mask=prompt_mask, ) hs = dec_out[1] if isinstance(dec_out, (tuple, list)) and len(dec_out) > 1 else None out_for_seg = dec_out[0] if isinstance(dec_out, (tuple, list)) else dec_out seg_ret = call_with_supported_kwargs( self.sam._run_segmentation_heads, backbone_out=backbone_out, img_ids=find_input.img_ids, vis_feat_sizes=encoder_out.get("vis_feat_sizes"), encoder_hidden_states=encoder_out.get("encoder_hidden_states"), prompt=prompt, prompt_mask=prompt_mask, hs=hs, out=out_for_seg, ) score_source = seg_ret if seg_ret is not None else out_for_seg logits = extract_mask_logits(score_source, b) logits = F.interpolate(logits.float(), size=output_size, mode="bilinear", align_corners=False) if return_details: details = extract_native_decoder_scores(score_source, b) details["mask_logits"] = logits return details return logits def unpack_prompt(text_out, batch_size, device): prompt = None prompt_mask = None if isinstance(text_out, dict): for key in ["language_features", "prompt", "memory_text", "text_features", "encoded_text"]: if key in text_out: prompt = text_out[key] break for key in ["language_mask", "prompt_mask", "text_attention_mask", "attention_mask", "mask"]: if key in text_out: prompt_mask = text_out[key] break elif isinstance(text_out, (list, tuple)) and text_out: prompt = text_out[0] if len(text_out) > 1: prompt_mask = text_out[1] elif torch.is_tensor(text_out): prompt = text_out if prompt is None: prompt = torch.zeros(1, batch_size, 256, device=device) prompt = prompt.to(device) if prompt.dim() == 2: prompt = prompt.unsqueeze(0) if prompt.dim() != 3: raise RuntimeError(f"SAM3 text prompt must be 3D, got shape {tuple(prompt.shape)}") if prompt.shape[0] == batch_size and prompt.shape[1] != batch_size: prompt = prompt.transpose(0, 1) if prompt.shape[1] != batch_size: raise RuntimeError(f"SAM3 text prompt batch dimension mismatch: shape={tuple(prompt.shape)}, batch={batch_size}") if prompt_mask is None: seq_len = prompt.shape[0] prompt_mask = torch.zeros(batch_size, seq_len, dtype=torch.bool, device=device) else: prompt_mask = prompt_mask.to(device=device, dtype=torch.bool) if prompt_mask.dim() != 2: raise RuntimeError(f"SAM3 text prompt mask must be 2D, got shape {tuple(prompt_mask.shape)}") if prompt_mask.shape[0] != batch_size and prompt_mask.shape[1] == batch_size: prompt_mask = prompt_mask.transpose(0, 1) if prompt_mask.shape != (batch_size, prompt.shape[0]): raise RuntimeError( f"SAM3 text prompt mask shape mismatch: mask={tuple(prompt_mask.shape)}, " f"expected={(batch_size, prompt.shape[0])}" ) return prompt, prompt_mask def extract_mask_logits(candidate, batch_size): logits = None if isinstance(candidate, dict): for key in ["pred_masks", "masks", "mask_logits", "seg_logits", "pred_mask_logits", "low_res_masks"]: if key in candidate and torch.is_tensor(candidate[key]): logits = candidate[key] break elif isinstance(candidate, (tuple, list)): for item in candidate: if torch.is_tensor(item): logits = item break elif torch.is_tensor(candidate): logits = candidate if logits is None: raise RuntimeError("SAM3 native decoder did not return mask logits") if logits.dim() == 3: logits = logits.unsqueeze(1) if logits.dim() == 4 and logits.shape[0] != batch_size and logits.shape[1] == batch_size: logits = logits.permute(1, 0, 2, 3) if logits.dim() == 4 and logits.shape[1] > 1: logits = logits.max(dim=1, keepdim=True).values if logits.dim() != 4 or logits.shape[1] != 1: raise RuntimeError(f"Unexpected native logits shape: {tuple(logits.shape)}") return logits def _iter_nested_named_tensors(candidate, prefix=""): if isinstance(candidate, dict): for key, value in candidate.items(): name = f"{prefix}.{key}" if prefix else str(key) if torch.is_tensor(value): yield name, value elif isinstance(value, (dict, list, tuple)): yield from _iter_nested_named_tensors(value, name) elif isinstance(candidate, (list, tuple)): for idx, value in enumerate(candidate): name = f"{prefix}.{idx}" if prefix else str(idx) if torch.is_tensor(value): yield name, value elif isinstance(value, (dict, list, tuple)): yield from _iter_nested_named_tensors(value, name) def _score_tensor_to_batch_vector(tensor, batch_size): score = tensor.detach().float() if score.numel() == 0: return None if score.dim() == 0: return score.reshape(1).repeat(batch_size) if score.shape[0] == batch_size: return score.reshape(batch_size, -1).max(dim=1).values if score.dim() >= 2 and score.shape[1] == batch_size: return score.transpose(0, 1).reshape(batch_size, -1).max(dim=1).values if score.numel() == batch_size: return score.reshape(batch_size) return None def extract_native_decoder_scores(candidate, batch_size): """Best-effort extraction of SAM3 native confidence/QC scores. These scores are model-internal confidence estimates, not true mask quality. SAM3 package versions use different names, so this walks nested outputs and accepts common score key variants. """ details = { "predicted_iou": None, "objectness_score": None, "selected_mask_index": None, "multimask_candidate_scores": None, } iou_keys = ("predicted_iou", "predicted_iou_score", "iou_predictions", "iou_scores", "mask_quality_score", "mask_quality_scores") obj_keys = ("objectness_score", "objectness_scores", "object_score_logits", "objectness_logits", "obj_score_logits") multi_keys = ("multimask", "candidate_scores", "mask_scores", "scores") for name, tensor in _iter_nested_named_tensors(candidate): lname = name.lower() if details["predicted_iou"] is None and any(key in lname for key in iou_keys): details["predicted_iou"] = _score_tensor_to_batch_vector(tensor, batch_size) if details["objectness_score"] is None and any(key in lname for key in obj_keys): score = _score_tensor_to_batch_vector(tensor, batch_size) if score is not None and "logit" in lname: score = torch.sigmoid(score) details["objectness_score"] = score if details["multimask_candidate_scores"] is None and any(key in lname for key in multi_keys): score = tensor.detach().float() if score.dim() >= 2 and (score.shape[0] == batch_size or score.shape[1] == batch_size): if score.shape[0] != batch_size and score.shape[1] == batch_size: score = score.transpose(0, 1) score = score.reshape(batch_size, -1) details["multimask_candidate_scores"] = score details["selected_mask_index"] = score.argmax(dim=1) return details def dice_bce_loss(logits, target): bce = F.binary_cross_entropy_with_logits(logits, target) prob = torch.sigmoid(logits) inter = (prob * target).sum(dim=(2, 3)) denom = prob.sum(dim=(2, 3)) + target.sum(dim=(2, 3)) dice = (2 * inter + 1e-6) / (denom + 1e-6) return bce + (1.0 - dice).mean() def binary_metrics(prob, target, threshold=0.5): pred = prob >= threshold gt = target > 0.5 tp = float(np.logical_and(pred, gt).sum()) fp = float(np.logical_and(pred, ~gt).sum()) fn = float(np.logical_and(~pred, gt).sum()) tn = float(np.logical_and(~pred, ~gt).sum()) dice = (2 * tp) / (2 * tp + fp + fn + 1e-8) iou = tp / (tp + fp + fn + 1e-8) precision = tp / (tp + fp + 1e-8) recall = tp / (tp + fn + 1e-8) specificity = tn / (tn + fp + 1e-8) area_error = (float(pred.sum()) - float(gt.sum())) / (float(gt.sum()) + 1e-8) hd95, assd = surface_distances(pred.astype(np.uint8), gt.astype(np.uint8)) return { "dice": dice, "iou": iou, "precision": precision, "recall": recall, "specificity": specificity, "hd95": hd95, "assd": assd, "mask_area_error": area_error, } def surface_distances(pred, gt): if pred.sum() == 0 and gt.sum() == 0: return 0.0, 0.0 if pred.sum() == 0 or gt.sum() == 0: return float("nan"), float("nan") try: from scipy.ndimage import binary_erosion, distance_transform_edt pred_border = pred.astype(bool) ^ binary_erosion(pred.astype(bool)) gt_border = gt.astype(bool) ^ binary_erosion(gt.astype(bool)) dt_pred = distance_transform_edt(~pred_border) dt_gt = distance_transform_edt(~gt_border) d1 = dt_gt[pred_border] d2 = dt_pred[gt_border] d = np.concatenate([d1, d2]).astype(float) return float(np.percentile(d, 95)), float(d.mean()) except Exception: return float("nan"), float("nan") def summarize_metrics(rows, group_cols=None): df = pd.DataFrame(rows) metric_cols = ["dice", "iou", "precision", "recall", "specificity", "hd95", "assd", "mask_area_error"] if group_cols is None: return {f"mean_{m}": float(pd.to_numeric(df[m], errors="coerce").mean()) for m in metric_cols} grouped = [] for keys, g in df.groupby(group_cols, dropna=False): item = {} if not isinstance(keys, tuple): keys = (keys,) for col, val in zip(group_cols, keys): item[col] = val item["n"] = int(len(g)) for m in metric_cols: item[f"mean_{m}"] = float(pd.to_numeric(g[m], errors="coerce").mean()) grouped.append(item) return pd.DataFrame(grouped) def bootstrap_ci(rows, metric, group_key="case_id", n_resamples=2000, seed=42): df = pd.DataFrame(rows) if df.empty or metric not in df: return {} rng = np.random.default_rng(seed) group_means = df.groupby(group_key)[metric].mean().dropna().to_numpy() if group_key in df else df[metric].dropna().to_numpy() if group_means.size == 0: return {} samples = [] for _ in range(int(n_resamples)): idx = rng.integers(0, group_means.size, size=group_means.size) samples.append(float(group_means[idx].mean())) return { f"{metric}_bootstrap_mean": float(group_means.mean()), f"{metric}_ci_low": float(np.percentile(samples, 2.5)), f"{metric}_ci_high": float(np.percentile(samples, 97.5)), f"{metric}_bootstrap_n_groups": int(group_means.size), } def ensure_dir(path): Path(path).mkdir(parents=True, exist_ok=True) def save_prediction(prob, path): arr = ((prob >= 0.5).astype(np.uint8) * 255) Image.fromarray(arr).save(path) def save_overlay(image_path, gt, pred, metrics, path): img = read_image_rgb(image_path).resize((gt.shape[1], gt.shape[0]), Image.BILINEAR) canvas = img.copy() overlay = Image.new("RGBA", canvas.size, (0, 0, 0, 0)) draw = ImageDraw.Draw(overlay) gt_contours = contours_from_mask(gt) pred_contours = contours_from_mask(pred) for pts in gt_contours: if len(pts) > 1: draw.line(pts, fill=(0, 255, 0, 220), width=2) for pts in pred_contours: if len(pts) > 1: draw.line(pts, fill=(255, 0, 0, 220), width=2) canvas = Image.alpha_composite(canvas.convert("RGBA"), overlay).convert("RGB") draw = ImageDraw.Draw(canvas) text = ( f"Dice {metrics['dice']:.3f} IoU {metrics['iou']:.3f} | " f"{metrics.get('dataset','')} {metrics.get('label','')} | {metrics.get('image_name','')}" ) draw.rectangle((0, 0, canvas.size[0], 24), fill=(0, 0, 0)) draw.text((4, 4), text[:160], fill=(255, 255, 255)) canvas.save(path) def contours_from_mask(mask): mask_u8 = (mask > 0).astype(np.uint8) * 255 contours, _ = cv2.findContours(mask_u8, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) out = [] for c in contours: pts = [(int(p[0][0]), int(p[0][1])) for p in c] out.append(pts) return out def write_csv(path, rows): if not rows: return with open(path, "w", newline="") as f: writer = csv.DictWriter(f, fieldnames=list(rows[0].keys())) writer.writeheader() writer.writerows(rows) def write_json(path, obj): Path(path).write_text(json.dumps(obj, indent=2)) def choose_visual_examples(rows, seed=42): rng = random.Random(seed) ordered = sorted(rows, key=lambda r: r["dice"]) worst = ordered[:20] best = ordered[-20:] sample = rows.copy() rng.shuffle(sample) return {"worst": worst, "best": best, "random": sample[:50]} @torch.no_grad() def run_eval(model, loader, device, threshold=0.5, save_dir=None, save_outputs=False, seed=42): model.eval() rows = [] pred_dir = Path(save_dir) / "masks" if save_dir and save_outputs else None vis_dir = Path(save_dir) / "overlays" if save_dir and save_outputs else None if pred_dir: ensure_dir(pred_dir) ensure_dir(vis_dir) cache_for_vis = [] for batch in loader: images = batch["image"].to(device, non_blocking=True) masks = batch["mask"].to(device, non_blocking=True) logits = model(images, masks) probs = torch.sigmoid(logits).detach().cpu().numpy() gt = masks.detach().cpu().numpy() for i in range(probs.shape[0]): prob_i = probs[i, 0] gt_i = gt[i, 0] metrics = binary_metrics(prob_i, gt_i, threshold=threshold) row = { "dataset": batch["dataset"][i], "label": batch["label"][i], "label_id": int(batch["label_id"][i]), "case_id": batch["case_id"][i], "image_name": batch["image_name"][i], "mask_name": batch["mask_name"][i], "image_path": batch["image_path"][i], "threshold": float(threshold), **metrics, } rows.append(row) if save_outputs: safe_name = Path(row["image_name"]).stem save_prediction(prob_i, pred_dir / f"{safe_name}_pred.png") cache_for_vis.append((row, gt_i, prob_i >= threshold)) if save_outputs and cache_for_vis: by_name = {row["image_name"]: (row, gt_i, pred_i) for row, gt_i, pred_i in cache_for_vis} for group, examples in choose_visual_examples(rows, seed=seed).items(): group_dir = vis_dir / group ensure_dir(group_dir) for row in examples: cached = by_name.get(row["image_name"]) if cached is None: continue _, gt_i, pred_i = cached save_overlay(row["image_path"], gt_i, pred_i, row, group_dir / f"{Path(row['image_name']).stem}.png") return rows def save_eval_outputs(rows, output_dir, prefix, bootstrap_resamples, seed): output_dir = Path(output_dir) write_csv(output_dir / f"{prefix}_per_image_metrics.csv", rows) overall = summarize_metrics(rows) overall.update(bootstrap_ci(rows, "dice", n_resamples=bootstrap_resamples, seed=seed)) overall.update(bootstrap_ci(rows, "iou", n_resamples=bootstrap_resamples, seed=seed + 1)) write_json(output_dir / f"{prefix}_overall_metrics.json", overall) by_dataset = summarize_metrics(rows, ["dataset"]) by_label = summarize_metrics(rows, ["label"]) by_dataset.to_csv(output_dir / f"{prefix}_metrics_by_dataset.csv", index=False) by_label.to_csv(output_dir / f"{prefix}_metrics_by_label.csv", index=False) write_json(output_dir / f"{prefix}_metrics_by_dataset.json", by_dataset.to_dict(orient="records")) write_json(output_dir / f"{prefix}_metrics_by_label.json", by_label.to_dict(orient="records")) return overall