scripts/train_pet_foundation.py

from __future__ import annotations

import argparse
import inspect
import importlib.util
from pathlib import Path
import sys
import types

import torch
from torch import nn
from torch.utils.data import DataLoader

from pet_vlm_dataset import PETSUVRDataset, collate_pet_suvr


class MedicalNetBottleneck(nn.Module):
    expansion = 4

    def __init__(self, inplanes: int, planes: int, stride: int = 1, downsample: nn.Module | None = None) -> None:
        super().__init__()
        self.conv1 = nn.Conv3d(inplanes, planes, kernel_size=1, bias=False)
        self.bn1 = nn.BatchNorm3d(planes)
        self.conv2 = nn.Conv3d(planes, planes, kernel_size=3, stride=stride, padding=1, bias=False)
        self.bn2 = nn.BatchNorm3d(planes)
        self.conv3 = nn.Conv3d(planes, planes * self.expansion, kernel_size=1, bias=False)
        self.bn3 = nn.BatchNorm3d(planes * self.expansion)
        self.relu = nn.ReLU(inplace=True)
        self.downsample = downsample

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        residual = x
        out = self.relu(self.bn1(self.conv1(x)))
        out = self.relu(self.bn2(self.conv2(out)))
        out = self.bn3(self.conv3(out))
        if self.downsample is not None:
            residual = self.downsample(x)
        out = self.relu(out + residual)
        return out


class MedicalNetResNet50(nn.Module):
    out_dim = 2048

    def __init__(self) -> None:
        super().__init__()
        self.inplanes = 64
        self.conv1 = nn.Conv3d(1, 64, kernel_size=7, stride=(2, 2, 2), padding=(3, 3, 3), bias=False)
        self.bn1 = nn.BatchNorm3d(64)
        self.relu = nn.ReLU(inplace=True)
        self.maxpool = nn.MaxPool3d(kernel_size=3, stride=2, padding=1)
        self.layer1 = self._make_layer(64, 3)
        self.layer2 = self._make_layer(128, 4, stride=2)
        self.layer3 = self._make_layer(256, 6, stride=2)
        self.layer4 = self._make_layer(512, 3, stride=2)
        self.pool = nn.AdaptiveAvgPool3d(1)

    def _make_layer(self, planes: int, blocks: int, stride: int = 1) -> nn.Sequential:
        downsample = None
        if stride != 1 or self.inplanes != planes * MedicalNetBottleneck.expansion:
            downsample = nn.Sequential(
                nn.Conv3d(self.inplanes, planes * MedicalNetBottleneck.expansion, kernel_size=1, stride=stride, bias=False),
                nn.BatchNorm3d(planes * MedicalNetBottleneck.expansion),
            )
        layers = [MedicalNetBottleneck(self.inplanes, planes, stride, downsample)]
        self.inplanes = planes * MedicalNetBottleneck.expansion
        for _ in range(1, blocks):
            layers.append(MedicalNetBottleneck(self.inplanes, planes))
        return nn.Sequential(*layers)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.maxpool(self.relu(self.bn1(self.conv1(x))))
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)
        return self.pool(x).flatten(1)


def load_medicalnet(path: Path) -> MedicalNetResNet50:
    model = MedicalNetResNet50()
    obj = torch.load(path, map_location="cpu")
    state = obj.get("state_dict", obj)
    state = {k.removeprefix("module."): v for k, v in state.items()}
    missing, unexpected = model.load_state_dict(state, strict=False)
    if unexpected:
        print(f"unexpected_keys={unexpected[:8]}", flush=True)
    if missing:
        print(f"missing_keys={missing[:8]}", flush=True)
    return model


class BrainIACEncoder(nn.Module):
    out_dim = 768

    def __init__(self, weights_path: Path) -> None:
        super().__init__()
        from monai.networks.nets import ViT
        from safetensors.torch import load_file

        self.model = ViT(
            in_channels=1,
            img_size=(96, 96, 96),
            patch_size=(16, 16, 16),
            hidden_size=768,
            mlp_dim=3072,
            num_layers=12,
            num_heads=12,
        )
        weights = load_file(str(weights_path))
        missing, unexpected = self.model.load_state_dict(weights, strict=False)
        if unexpected:
            print(f"brainiac_unexpected_keys={unexpected[:8]}", flush=True)
        if missing:
            print(f"brainiac_missing_keys={missing[:8]}", flush=True)

    def forward(self, image: torch.Tensor) -> torch.Tensor:
        output = self.model(image)
        tokens = output[0] if isinstance(output, tuple) else output
        return tokens[:, 0]


class SwinUNETREncoder(nn.Module):
    out_dim = 768

    def __init__(self, weights_path: Path, img_size: tuple[int, int, int]) -> None:
        super().__init__()
        from monai.networks.nets import SwinUNETR

        kwargs = {
            "in_channels": 1,
            "out_channels": 2,
            "feature_size": 48,
            "use_checkpoint": False,
            "spatial_dims": 3,
        }
        if "img_size" in inspect.signature(SwinUNETR).parameters:
            kwargs["img_size"] = img_size
        self.model = SwinUNETR(**kwargs)
        weights = torch.load(weights_path, map_location="cpu", weights_only=False)
        if hasattr(self.model, "load_from"):
            self.model.load_from(weights)
        else:
            state = weights.get("state_dict", weights.get("model", weights))
            remapped = {}
            for key, value in state.items():
                key = key.removeprefix("module.")
                if key.startswith("encoder."):
                    key = "swinViT." + key[len("encoder.") :]
                remapped[key] = value
            missing, unexpected = self.model.load_state_dict(remapped, strict=False)
            if unexpected:
                print(f"swinunetr_unexpected_keys={unexpected[:8]}", flush=True)
            if missing:
                print(f"swinunetr_missing_keys={missing[:8]}", flush=True)
        self.pool = nn.AdaptiveAvgPool3d(1)

    def forward(self, image: torch.Tensor) -> torch.Tensor:
        hidden = self.model.swinViT(image, self.model.normalize)
        feat = hidden[-1]
        return self.pool(feat).flatten(1)


class SAMMed3DEncoder(nn.Module):
    out_dim = 384

    def __init__(self, weights_path: Path) -> None:
        super().__init__()
        try:
            import medim
        except ImportError as exc:
            raise ImportError("SAM-Med3D requires `medim`. Install it before using --backbone sam_med3d.") from exc

        self.model = medim.create_model("SAM-Med3D", pretrained=True, checkpoint_path=str(weights_path))
        self.image_encoder = getattr(self.model, "image_encoder", self.model)
        self.pool = nn.AdaptiveAvgPool3d(1)

    def forward(self, image: torch.Tensor) -> torch.Tensor:
        output = self.image_encoder(image)
        if isinstance(output, dict):
            output = output.get("image_embeddings", output.get("embeddings", next(iter(output.values()))))
        if isinstance(output, (list, tuple)):
            output = output[0]
        if output.ndim == 2:
            return output
        if output.ndim == 3:
            return output.mean(dim=1)
        return self.pool(output).flatten(1)


class BrainFMEncoder(nn.Module):
    out_dim = 2048

    def __init__(self, weights_path: Path, code_root: Path) -> None:
        super().__init__()
        unet_root = (code_root / "Trainer" / "models" / "unet3d").resolve()
        package = types.ModuleType("brainfm_unet")
        package.__path__ = [str(unet_root)]
        sys.modules.setdefault("brainfm_unet", package)
        spec = importlib.util.spec_from_file_location("brainfm_unet.model", unet_root / "model.py")
        if spec is None or spec.loader is None:
            raise RuntimeError(f"Could not load BrainFM model code from {unet_root}")
        module = importlib.util.module_from_spec(spec)
        sys.modules["brainfm_unet.model"] = module
        spec.loader.exec_module(module)

        ckpt = torch.load(weights_path, map_location="cpu", weights_only=False)
        train_args = ckpt["train_args"]
        self.model = module.UNet3D(
            train_args.in_channels,
            train_args.f_maps,
            train_args.layer_order,
            train_args.num_groups,
            train_args.num_levels,
            train_args.unit_feat,
        )
        state = {k.removeprefix("backbone."): v for k, v in ckpt["model"].items() if k.startswith("backbone.")}
        missing, unexpected = self.model.load_state_dict(state, strict=False)
        if unexpected:
            print(f"brainfm_unexpected_keys={unexpected[:8]}", flush=True)
        if missing:
            print(f"brainfm_missing_keys={missing[:8]}", flush=True)
        self.pool = nn.AdaptiveAvgPool3d(1)

    def forward(self, image: torch.Tensor) -> torch.Tensor:
        features = self.model.get_feature(image)
        bottleneck = features[0]
        return self.pool(bottleneck).flatten(1)


class Small3DPETEncoder(nn.Module):
    out_dim = 256

    def __init__(self) -> None:
        super().__init__()
        self.net = nn.Sequential(
            nn.Conv3d(1, 16, 3, stride=2, padding=1),
            nn.BatchNorm3d(16),
            nn.GELU(),
            nn.Conv3d(16, 32, 3, stride=2, padding=1),
            nn.BatchNorm3d(32),
            nn.GELU(),
            nn.Conv3d(32, 64, 3, stride=2, padding=1),
            nn.BatchNorm3d(64),
            nn.GELU(),
            nn.Conv3d(64, 128, 3, stride=2, padding=1),
            nn.BatchNorm3d(128),
            nn.GELU(),
            nn.AdaptiveAvgPool3d(1),
        )
        self.proj = nn.Linear(128, self.out_dim)

    def forward(self, image: torch.Tensor) -> torch.Tensor:
        return self.proj(self.net(image).flatten(1))


class RegionSUVREncoder(nn.Module):
    def __init__(self, n_regions: int, embed_dim: int) -> None:
        super().__init__()
        self.net = nn.Sequential(
            nn.LayerNorm(n_regions),
            nn.Linear(n_regions, embed_dim),
            nn.GELU(),
            nn.Linear(embed_dim, embed_dim),
        )

    def forward(self, suvr: torch.Tensor) -> torch.Tensor:
        return self.net(suvr)


class PETSUVRFoundationModel(nn.Module):
    def __init__(self, pet_encoder: nn.Module, n_regions: int, embed_dim: int = 256, freeze_encoder: bool = True) -> None:
        super().__init__()
        self.pet_encoder = pet_encoder
        self.freeze_encoder = freeze_encoder
        if freeze_encoder:
            for p in self.pet_encoder.parameters():
                p.requires_grad = False
            self.pet_encoder.eval()
        self.pet_projector = nn.Sequential(nn.LayerNorm(pet_encoder.out_dim), nn.Linear(pet_encoder.out_dim, embed_dim))
        self.suvr_encoder = RegionSUVREncoder(n_regions, embed_dim)
        self.suvr_head = nn.Sequential(nn.LayerNorm(embed_dim), nn.Linear(embed_dim, n_regions))
        self.temperature = nn.Parameter(torch.tensor(0.07))

    def forward(self, image: torch.Tensor, suvr: torch.Tensor) -> dict[str, torch.Tensor]:
        if self.freeze_encoder:
            with torch.no_grad():
                pet_feat = self.pet_encoder(image)
        else:
            pet_feat = self.pet_encoder(image)
        pet_z = nn.functional.normalize(self.pet_projector(pet_feat), dim=-1)
        suvr_z = nn.functional.normalize(self.suvr_encoder(suvr), dim=-1)
        pred_suvr = self.suvr_head(pet_z)
        logits = pet_z @ suvr_z.T / self.temperature.clamp_min(0.01)
        return {"logits": logits, "pred_suvr": pred_suvr}


def alignment_loss(
    outputs: dict[str, torch.Tensor],
    suvr: torch.Tensor,
    contrastive_weight: float = 1.0,
    regression_weight: float = 1.0,
) -> tuple[torch.Tensor, dict[str, float]]:
    labels = torch.arange(suvr.shape[0], device=suvr.device)
    loss_i = nn.functional.cross_entropy(outputs["logits"], labels)
    loss_t = nn.functional.cross_entropy(outputs["logits"].T, labels)
    loss_contrastive = 0.5 * (loss_i + loss_t)
    loss_reg = nn.functional.mse_loss(outputs["pred_suvr"], suvr)
    loss = contrastive_weight * loss_contrastive + regression_weight * loss_reg
    return loss, {
        "contrastive": float(loss_contrastive.detach()),
        "regression": float(loss_reg.detach()),
    }


def build_encoder(args: argparse.Namespace) -> nn.Module:
    if args.backbone == "small_cnn":
        return Small3DPETEncoder()
    if args.backbone == "medicalnet":
        return load_medicalnet(args.medicalnet_weights)
    if args.backbone == "brainiac":
        return BrainIACEncoder(args.brainiac_weights)
    if args.backbone == "swinunetr":
        return SwinUNETREncoder(args.swinunetr_weights, tuple(args.output_size))
    if args.backbone == "sam_med3d":
        return SAMMed3DEncoder(args.sam_med3d_weights)
    if args.backbone == "brainfm":
        return BrainFMEncoder(args.brainfm_weights, args.brainfm_code_root)
    raise ValueError(f"Unsupported backbone: {args.backbone}")


def _set_trainable(module: nn.Module | None, trainable: bool) -> None:
    if module is None:
        return
    for p in module.parameters():
        p.requires_grad = trainable


def _last_vit_block(module: nn.Module) -> nn.Module | None:
    blocks = getattr(module, "blocks", None)
    if isinstance(blocks, (nn.ModuleList, list, tuple)) and len(blocks) > 0:
        return blocks[-1]
    return None


def _last_swin_stage(module: nn.Module) -> nn.Module | None:
    swin = getattr(module, "swinViT", None)
    if swin is None:
        return None
    for name in ("layers4", "layers3", "layers2", "layers1"):
        layer = getattr(swin, name, None)
        if layer is not None:
            return layer
    return None


def _brainfm_last_stage(module: nn.Module) -> nn.Module | None:
    for name in ("encoders", "encoder", "down_path"):
        stage = getattr(module, name, None)
        if isinstance(stage, (nn.ModuleList, nn.Sequential)) and len(stage) > 0:
            return stage[-1]
    return None


def _sam_last_stage(module: nn.Module) -> nn.Module | None:
    for name in ("blocks", "layers", "neck"):
        stage = getattr(module, name, None)
        if isinstance(stage, (nn.ModuleList, list, tuple)) and len(stage) > 0:
            return stage[-1]
        if isinstance(stage, nn.Module):
            return stage
    return None


def unfreeze_last_block(encoder: nn.Module) -> None:
    _set_trainable(encoder, False)
    target = None
    if isinstance(encoder, MedicalNetResNet50):
        target = encoder.layer4
    elif isinstance(encoder, BrainIACEncoder):
        target = _last_vit_block(encoder.model)
    elif isinstance(encoder, SwinUNETREncoder):
        target = _last_swin_stage(encoder.model)
    elif isinstance(encoder, BrainFMEncoder):
        target = _brainfm_last_stage(encoder.model)
    elif isinstance(encoder, SAMMed3DEncoder):
        target = _sam_last_stage(encoder.image_encoder)
    if target is None:
        raise ValueError(f"Could not identify a last block for {encoder.__class__.__name__}.")
    _set_trainable(target, True)


def configure_encoder_training(model: PETSUVRFoundationModel, scope: str) -> None:
    if scope == "none":
        for p in model.pet_encoder.parameters():
            p.requires_grad = False
        return
    if scope == "all":
        for p in model.pet_encoder.parameters():
            p.requires_grad = True
        return
    if scope == "layer4":
        for p in model.pet_encoder.parameters():
            p.requires_grad = False
        if not hasattr(model.pet_encoder, "layer4"):
            raise ValueError("encoder_train_scope=layer4 is only supported for MedicalNet-style encoders.")
        for p in model.pet_encoder.layer4.parameters():
            p.requires_grad = True
        return
    if scope == "last_block":
        unfreeze_last_block(model.pet_encoder)
        return
    raise ValueError(f"Unsupported encoder training scope: {scope}")


def set_encoder_mode_for_scope(model: PETSUVRFoundationModel, scope: str) -> None:
    if scope == "none":
        model.pet_encoder.eval()
    elif scope in {"layer4", "last_block"}:
        model.pet_encoder.eval()
        for module in model.pet_encoder.modules():
            if any(p.requires_grad for p in module.parameters(recurse=False)):
                module.train()


def main() -> None:
    parser = argparse.ArgumentParser(description="Train PET-SUVR alignment with pretrained 3D backbones.")
    parser.add_argument("--backbone", choices=["small_cnn", "medicalnet", "brainiac", "brainfm", "swinunetr", "sam_med3d"], default="medicalnet")
    parser.add_argument("--medicalnet-weights", type=Path, default=Path("pretrained/medicalnet/resnet_50_23dataset.pth"))
    parser.add_argument("--brainiac-weights", type=Path, default=Path("pretrained/brainiac/backbone.safetensors"))
    parser.add_argument("--brainfm-weights", type=Path, default=Path("pretrained/brainfm/assets/brainfm_pretrained.pth"))
    parser.add_argument("--brainfm-code-root", type=Path, default=Path("pretrained/brainfm"))
    parser.add_argument("--swinunetr-weights", type=Path, default=Path("pretrained/swinunetr/model_swinvit.pt"))
    parser.add_argument("--sam-med3d-weights", type=Path, default=Path("pretrained/sam-med3d/sam_med3d_turbo.pth"))
    parser.add_argument("--manifest", type=Path, default=Path("metadata/splits/train.csv"))
    parser.add_argument("--val-manifest", type=Path, default=Path("metadata/splits/val.csv"))
    parser.add_argument("--epochs", type=int, default=10)
    parser.add_argument("--batch-size", type=int, default=2)
    parser.add_argument("--lr", type=float, default=1e-4)
    parser.add_argument("--num-workers", type=int, default=2)
    parser.add_argument("--output-size", type=int, nargs=3, default=(96, 96, 96))
    parser.add_argument("--embed-dim", type=int, default=256)
    parser.add_argument("--freeze-encoder", action=argparse.BooleanOptionalAction, default=True)
    parser.add_argument("--encoder-train-scope", choices=["none", "layer4", "last_block", "all"], default=None)
    parser.add_argument("--contrastive-weight", type=float, default=1.0)
    parser.add_argument("--regression-weight", type=float, default=1.0)
    parser.add_argument("--log-every", type=int, default=10)
    parser.add_argument("--out", type=Path, default=Path("runs/foundation.pt"))
    parser.add_argument("--best-out", type=Path, default=None)
    args = parser.parse_args()
    if args.encoder_train_scope is None:
        args.encoder_train_scope = "none" if args.freeze_encoder else "all"
    args.freeze_encoder = args.encoder_train_scope == "none"

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    train_dataset = PETSUVRDataset(args.manifest, output_size=tuple(args.output_size))
    val_dataset = PETSUVRDataset(args.val_manifest, output_size=tuple(args.output_size))
    train_loader = DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True, num_workers=args.num_workers, collate_fn=collate_pet_suvr)
    val_loader = DataLoader(val_dataset, batch_size=args.batch_size, shuffle=False, num_workers=args.num_workers, collate_fn=collate_pet_suvr)

    sample = train_dataset[0]
    encoder = build_encoder(args)
    model = PETSUVRFoundationModel(encoder, int(sample["suvr"].numel()), args.embed_dim, args.freeze_encoder).to(device)
    configure_encoder_training(model, args.encoder_train_scope)
    optimizer = torch.optim.AdamW((p for p in model.parameters() if p.requires_grad), lr=args.lr, weight_decay=1e-4)

    best_val_loss = float("inf")
    best_out = args.best_out or args.out.with_name(args.out.stem + "_best" + args.out.suffix)
    print(
        f"device={device} backbone={args.backbone} encoder_scope={args.encoder_train_scope} "
        f"contrastive_weight={args.contrastive_weight} regression_weight={args.regression_weight} "
        f"train={len(train_dataset)} val={len(val_dataset)}",
        flush=True,
    )
    for epoch in range(1, args.epochs + 1):
        model.train()
        set_encoder_mode_for_scope(model, args.encoder_train_scope)
        train_loss = 0.0
        train_contrastive = 0.0
        train_regression = 0.0
        for step, batch in enumerate(train_loader, start=1):
            image = batch["image"].to(device, non_blocking=True)
            suvr = batch["suvr"].to(device, non_blocking=True)
            outputs = model(image, suvr)
            loss, parts = alignment_loss(outputs, suvr, args.contrastive_weight, args.regression_weight)
            optimizer.zero_grad(set_to_none=True)
            loss.backward()
            optimizer.step()
            train_loss += float(loss.detach()) * image.shape[0]
            train_contrastive += parts["contrastive"] * image.shape[0]
            train_regression += parts["regression"] * image.shape[0]
            if args.log_every and step % args.log_every == 0:
                print(f"epoch={epoch} step={step}/{len(train_loader)} loss={float(loss.detach()):.4f}", flush=True)
        train_loss /= len(train_dataset)
        train_contrastive /= len(train_dataset)
        train_regression /= len(train_dataset)

        model.eval()
        val_loss = 0.0
        val_contrastive = 0.0
        val_regression = 0.0
        with torch.no_grad():
            for batch in val_loader:
                image = batch["image"].to(device, non_blocking=True)
                suvr = batch["suvr"].to(device, non_blocking=True)
                loss, parts = alignment_loss(model(image, suvr), suvr, args.contrastive_weight, args.regression_weight)
                val_loss += float(loss) * image.shape[0]
                val_contrastive += parts["contrastive"] * image.shape[0]
                val_regression += parts["regression"] * image.shape[0]
        val_loss /= len(val_dataset)
        val_contrastive /= len(val_dataset)
        val_regression /= len(val_dataset)
        print(
            f"epoch={epoch} train_loss={train_loss:.4f} train_contrastive={train_contrastive:.4f} "
            f"train_regression={train_regression:.4f} val_loss={val_loss:.4f} "
            f"val_contrastive={val_contrastive:.4f} val_regression={val_regression:.4f}",
            flush=True,
        )
        if val_loss < best_val_loss:
            best_val_loss = val_loss
            best_out.parent.mkdir(parents=True, exist_ok=True)
            torch.save({"model": model.state_dict(), "args": vars(args), "best_val_loss": best_val_loss, "epoch": epoch}, best_out)
            print(f"saved_best {best_out} val_loss={best_val_loss:.4f} epoch={epoch}", flush=True)

    args.out.parent.mkdir(parents=True, exist_ok=True)
    torch.save({"model": model.state_dict(), "args": vars(args), "best_val_loss": best_val_loss}, args.out)
    print(f"saved {args.out}", flush=True)


if __name__ == "__main__":
    main()