scripts/sam3_decoder_experiment_lib.py

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