| """FlexiCT-2D slice encoder with class-query heads for 18-label CT classification. |
| |
| Model = FlexiCT-2D (LoRA-adapted, patch 16, conv branch on the patch embedding) |
| -> slice position encoding -> class-query pool -> class interaction |
| -> per-class classifier. |
| |
| Module and parameter names match the released checkpoints, so |
| `load_state_dict(strict=True)` works on a model built by `build_from_config`. |
| """ |
| from __future__ import annotations |
|
|
| import math |
| import re |
|
|
| import numpy as np |
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| import torchio as tio |
|
|
| FLEXICT_ID = "ricklisz123/FlexiCT-2D" |
| FLEXICT_REVISION = "fcefac1e223c4a1ed72a6387c8e6f2d1f848fa90" |
|
|
|
|
| |
| |
| |
| class ConvStem(nn.Module): |
| """Small strided CNN that maps an image to the patch grid: (B, C, H, W) -> (B, D, H/p, W/p).""" |
|
|
| def __init__(self, in_ch: int, embed_dim: int, patch, widths=(16, 32, 64, 128), groups: int = 8): |
| super().__init__() |
| patch = tuple(int(p) for p in patch) |
| n_axis = [int(round(math.log2(p))) for p in patch] |
| if any(2 ** n != p for n, p in zip(n_axis, patch)): |
| raise ValueError(f"ConvStem needs power-of-two patches, got {patch}") |
| layers, c_prev = [], in_ch |
| for i in range(max(n_axis)): |
| stride = tuple(2 if i < n else 1 for n in n_axis) |
| c = widths[min(i, len(widths) - 1)] |
| layers += [nn.Conv2d(c_prev, c, 3, stride=stride, padding=1, bias=False), |
| nn.GroupNorm(min(groups, c), c), nn.GELU()] |
| c_prev = c |
| layers += [nn.Conv2d(c_prev, embed_dim, 1, bias=True)] |
| self.net = nn.Sequential(*layers) |
| nn.init.trunc_normal_(self.net[-1].weight, std=0.02) |
| nn.init.zeros_(self.net[-1].bias) |
|
|
| def forward(self, x): |
| return self.net(x) |
|
|
|
|
| class PatchEmbedPlusConv(nn.Module): |
| """Runs the original patch projection and the conv branch in parallel. |
| |
| The branch output is scaled by a learned per-channel gain that starts near |
| zero, so training starts from the pretrained projection. |
| """ |
|
|
| def __init__(self, base_conv: nn.Conv2d, gamma_init: float = 1e-4, checkpoint: bool = False): |
| super().__init__() |
| self.base = base_conv |
| self.conv_branch = ConvStem(base_conv.in_channels, base_conv.out_channels, |
| tuple(base_conv.kernel_size)) |
| self.cb_gamma = nn.Parameter(torch.full((base_conv.out_channels,), float(gamma_init))) |
| self.checkpoint = checkpoint |
|
|
| |
| @property |
| def weight(self): return self.base.weight |
| @property |
| def bias(self): return self.base.bias |
| @property |
| def kernel_size(self): return self.base.kernel_size |
| @property |
| def stride(self): return self.base.stride |
| @property |
| def in_channels(self): return self.base.in_channels |
| @property |
| def out_channels(self): return self.base.out_channels |
|
|
| def forward(self, x): |
| y = self.base(x) |
| if self.checkpoint and self.training: |
| import torch.utils.checkpoint as cp |
| b = cp.checkpoint(self.conv_branch, x, use_reentrant=False) |
| else: |
| b = self.conv_branch(x) |
| return y + self.cb_gamma.view(1, -1, 1, 1) * b |
|
|
|
|
| def _find_patch_embed(model, hint="patch_embed_2D"): |
| for name, module in model.named_modules(): |
| if hint in name and hasattr(module, "proj"): |
| return name, module |
| raise RuntimeError(f"no patch embedding named '{hint}'") |
|
|
|
|
| def add_conv_branch(model, checkpoint: bool = False): |
| _, pe = _find_patch_embed(model) |
| pe.proj = PatchEmbedPlusConv(pe.proj, checkpoint=checkpoint) |
| return pe.proj |
|
|
|
|
| def bake_patch_size(model, new_patch: int): |
| """Replace the 8x8 patch projection by a real 16x16 conv. |
| |
| FlexiCT resamples its kernel on the fly when the patch size differs from |
| the pretrained one, and applies it functionally. That path would bypass |
| LoRA, so the resampled kernel is computed once (FlexiCT's own resampling) |
| and written into a conv module. |
| """ |
| _, pe = _find_patch_embed(model) |
| new_k = (new_patch, new_patch) |
| if tuple(pe.proj.kernel_size) == new_k: |
| return pe |
| with torch.no_grad(): |
| w = pe._resample_conv_weight(pe.proj.weight, new_k) |
| conv = nn.Conv2d(pe.proj.in_channels, pe.proj.out_channels, kernel_size=new_k, |
| stride=new_k, bias=pe.proj.bias is not None) |
| with torch.no_grad(): |
| conv.weight.copy_(w) |
| if pe.proj.bias is not None: |
| conv.bias.copy_(pe.proj.bias) |
| pe.proj = conv.to(w.device, w.dtype) |
| pe.base_patch_size = new_k |
| pe._runtime_patch_size = None |
| pe._pinv_cache = {} |
| return pe |
|
|
|
|
| |
| |
| |
| def body_crop_bbox(vol: np.ndarray, thr: float = -500.0, margin: int = 4): |
| """In-plane bounding box of voxels above `thr` HU over all slices, plus a margin.""" |
| m = vol > thr |
| xs, ys = m.any(axis=(1, 2)), m.any(axis=(0, 2)) |
| if not xs.any() or not ys.any(): |
| return 0, vol.shape[0], 0, vol.shape[1] |
| x0, x1 = int(np.argmax(xs)), int(len(xs) - np.argmax(xs[::-1])) |
| y0, y1 = int(np.argmax(ys)), int(len(ys) - np.argmax(ys[::-1])) |
| return (max(0, x0 - margin), min(vol.shape[0], x1 + margin), |
| max(0, y0 - margin), min(vol.shape[1], y1 + margin)) |
|
|
|
|
| def tissue_start(t: torch.Tensor, n: int, thr: float = -500.0) -> int | None: |
| """Start of an n-slice window centred on the slices that contain tissue; None if no crop is needed.""" |
| s = t.shape[0] |
| if s <= n: |
| return None |
| idx = torch.nonzero((t > thr).any(dim=1).any(dim=1)).flatten() |
| c = int((idx[0] + idx[-1]) // 2) if idx.numel() else s // 2 |
| return max(0, min(c - n // 2, s - n)) |
|
|
|
|
| def fit_xy(sl: torch.Tensor, size: int, fill: torch.Tensor) -> torch.Tensor: |
| """Centre-crop or pad (S, X, Y) to (S, size, size) with a per-slice fill value. No resizing.""" |
| s, x, y = sl.shape |
| out = fill.expand(s, size, size).clone() |
|
|
| def span(n, k): |
| if n >= k: |
| a = (n - k) // 2 |
| return a, a + k, 0, k |
| a = (k - n) // 2 |
| return 0, n, a, a + n |
|
|
| xs0, xs1, xd0, xd1 = span(x, size) |
| ys0, ys1, yd0, yd1 = span(y, size) |
| out[:, xd0:xd1, yd0:yd1] = sl[:, xs0:xs1, ys0:ys1] |
| return out |
|
|
|
|
| def preprocess_ct(volume: np.ndarray, affine: np.ndarray, *, img_size: int, xy_spacing: float, |
| z_spacing: float, max_slices: int, hu_clip=(-1000.0, 1000.0), |
| out_dtype=torch.float16) -> torch.Tensor: |
| """NIfTI array in HU -> (max_slices, 1, img_size, img_size).""" |
| img = tio.ScalarImage(tensor=torch.from_numpy(np.asarray(volume))[None].float(), affine=affine) |
| img = tio.ToCanonical()(img) |
| img = tio.Resample((xy_spacing, xy_spacing, z_spacing))(img) |
| vol = img.data[0].numpy() |
| x0, x1, y0, y1 = body_crop_bbox(vol) |
| vol = vol[x0:x1, y0:y1, :] |
| sl = np.ascontiguousarray(vol.transpose(2, 0, 1)[:, ::-1, ::-1]) |
| t = torch.from_numpy(sl).float().clamp_(*hu_clip) |
| st = tissue_start(t, max_slices) |
| if st is not None: |
| t = t[st:st + max_slices].contiguous() |
| mean = t.mean(dim=(1, 2), keepdim=True) |
| std = t.std(dim=(1, 2), keepdim=True).clamp(min=1e-6) |
| t = (t - mean) / std |
| air = (hu_clip[0] - mean) / std |
| t = fit_xy(t, img_size, air) |
| if t.shape[0] < max_slices: |
| pad = max_slices - t.shape[0] |
| t = F.pad(t, (0, 0, 0, 0, pad // 2, pad - pad // 2), value=float(air.mean())) |
| return t.unsqueeze(1).to(out_dtype).contiguous() |
|
|
|
|
| |
| |
| |
| class FlexiCT2DEncoder(nn.Module): |
| """Encodes every slice with FlexiCT-2D and returns all tokens: (B, S, 1, H, W) -> (B, S*T, D).""" |
|
|
| modality = "2d_slices" |
|
|
| def __init__(self, img_size: int = 480, patch_size: int = 16, max_slices: int = 240, |
| slice_microbatch: int = 240, conv_branch: bool = True, |
| grad_checkpointing: bool = False, model_id: str = FLEXICT_ID, |
| revision: str = FLEXICT_REVISION, load_pretrained: bool = True): |
| super().__init__() |
| from transformers import AutoConfig, AutoModel |
|
|
| cfg = AutoConfig.from_pretrained(model_id, revision=revision, trust_remote_code=True) |
| cfg.image_size = [img_size, img_size] |
| self.model = AutoModel.from_config(cfg, trust_remote_code=True) |
| if load_pretrained: |
| |
| from huggingface_hub import hf_hub_download |
| from safetensors.torch import load_file |
| self.model.load_state_dict( |
| load_file(hf_hub_download(model_id, "model.safetensors", revision=revision)), strict=False) |
| bake_patch_size(self.model, patch_size) |
|
|
| if grad_checkpointing: |
| import torch.utils.checkpoint as cp |
| depth = int(getattr(cfg, "depth", 16)) |
| blocks = next(m for m in self.model.modules() |
| if isinstance(m, nn.ModuleList) and len(m) == depth) |
| for blk in blocks: |
| blk.forward = (lambda f: lambda *a, **k: cp.checkpoint(f, *a, use_reentrant=False, **k))(blk.forward) |
| if conv_branch: |
| add_conv_branch(self.model, checkpoint=grad_checkpointing) |
|
|
| self.size = int(img_size) |
| self.max_slices = int(max_slices) |
| self.slice_microbatch = int(slice_microbatch) |
| self.embed_dim = int(getattr(cfg, "embed_dim", 864)) |
| grid = self.size // patch_size |
| self.tokens_per_slice = 1 + grid * grid |
| |
| |
| self.lora_target_modules = [ |
| n for n, m in self.model.named_modules() |
| if isinstance(m, (nn.Linear, nn.Conv2d, nn.Conv3d)) |
| and "patch_embed_3D" not in n and "text" not in n.lower() and "mlp" not in n |
| and any(k in n for k in ("qkv", "proj", "q_proj", "k_proj", "v_proj")) |
| ] |
|
|
| def _encode_slices(self, flat): |
| outs = [] |
| for i in range(0, flat.shape[0], self.slice_microbatch): |
| o = self.model(pixel_values=flat[i:i + self.slice_microbatch]) |
| cls = o.cls_token |
| if cls.dim() == 2: |
| cls = cls.unsqueeze(1) |
| outs.append(torch.cat([cls, o.patch_tokens], dim=1)) |
| return torch.cat(outs, dim=0) |
|
|
| def forward(self, batch): |
| b, s = batch.shape[:2] |
| f = self._encode_slices(batch.flatten(0, 1)) |
| return f.reshape(b, s * f.shape[1], -1), None |
|
|
|
|
| class SlicePosEnc(nn.Module): |
| """Factorised position code for (B, S*T, D) tokens: slice (z) table + in-plane x/y tables + CLS vector.""" |
|
|
| def __init__(self, dim: int, n_slices: int, tokens_per_slice: int, scale: float = 0.25): |
| super().__init__() |
| self.S, self.T, self.dim = int(n_slices), int(tokens_per_slice), int(dim) |
| self.g = int(round((self.T - 1) ** 0.5)) |
| std = scale * dim ** -0.5 |
| self.z = nn.Parameter(torch.randn(self.S, dim) * std) |
| self.x = nn.Parameter(torch.randn(self.g, dim) * std) |
| self.y = nn.Parameter(torch.randn(self.g, dim) * std) |
| self.cls = nn.Parameter(torch.randn(1, dim) * std) |
|
|
| def forward(self, tokens): |
| b, n, d = tokens.shape |
| xy = (self.x.view(self.g, 1, d) + self.y.view(1, self.g, d)).reshape(self.g * self.g, d) |
| per_slice = torch.cat([self.cls, xy], 0) |
| pe = self.z.unsqueeze(1) + per_slice.unsqueeze(0) |
| return tokens + pe.reshape(1, self.S * self.T, d).to(tokens.dtype) |
|
|
|
|
| |
| |
| |
| class MultiHeadQueryAttentionPool(nn.Module): |
| """One learned query per class attends over all tokens, with several heads.""" |
|
|
| def __init__(self, dim: int, num_queries: int = 18, num_heads: int = 4, |
| d_attn: int | None = None, identity_out: bool = True): |
| super().__init__() |
| d_attn = d_attn or dim |
| self.num_queries, self.num_heads, self.head_dim = num_queries, num_heads, d_attn // num_heads |
| self.query = nn.Parameter(torch.empty(1, num_heads, num_queries, self.head_dim)) |
| self.logit_scale = nn.Parameter(torch.tensor(math.log(4.0))) |
| nn.init.trunc_normal_(self.query, std=0.1) |
| self.k = nn.Linear(dim, d_attn, bias=False) |
| self.v = nn.Linear(dim, dim, bias=False) |
| self.proj_out = nn.Linear(dim, dim) if num_heads > 1 else nn.Identity() |
| if identity_out and isinstance(self.proj_out, nn.Linear): |
| nn.init.eye_(self.proj_out.weight) |
| nn.init.zeros_(self.proj_out.bias) |
|
|
| def forward(self, x, mask=None): |
| b, n, d = x.shape |
| h, q_n = self.num_heads, self.num_queries |
| k = self.k(x).view(b, n, h, self.head_dim).transpose(1, 2) |
| v = self.v(x).view(b, n, h, d // h).transpose(1, 2) |
| gain = self.logit_scale.clamp(max=math.log(100.0)).exp() |
| q = self.query.expand(b, -1, -1, -1) * ((self.head_dim ** -0.5) * gain) |
| out = F.scaled_dot_product_attention(q, k, v, scale=1.0) |
| return self.proj_out(out.transpose(1, 2).contiguous().view(b, q_n, d)) |
|
|
|
|
| class GatedMaxTokenPool(nn.Module): |
| """Attention pool plus the single highest-scoring token per class, scaled by tanh(gate). |
| |
| The gate starts at zero, so the head starts as a plain attention pool. |
| """ |
|
|
| def __init__(self, dim: int, num_queries: int = 18, d_attn: int | None = None): |
| super().__init__() |
| d_attn = d_attn or dim |
| self.query = nn.Parameter(torch.empty(1, num_queries, d_attn)) |
| self.logit_scale = nn.Parameter(torch.tensor(math.log(4.0))) |
| nn.init.trunc_normal_(self.query, std=0.1) |
| self.k = nn.Linear(dim, d_attn, bias=False) |
| self.v = nn.Linear(dim, dim, bias=False) |
| self.scale = d_attn ** -0.5 |
| self.max_gate = nn.Parameter(torch.zeros(1, num_queries, 1)) |
|
|
| def forward(self, x, mask=None): |
| b, n, d = x.shape |
| k, v = self.k(x), self.v(x) |
| gain = self.logit_scale.clamp(max=math.log(100.0)).exp() |
| scores = torch.bmm(self.query.expand(b, -1, -1) * (self.scale * gain), k.transpose(1, 2)) |
| out_mean = torch.bmm(F.softmax(scores, dim=-1), v) |
| idx = scores.argmax(dim=-1).unsqueeze(-1).expand(-1, -1, d) |
| out_max = torch.gather(v, 1, idx) |
| return out_mean + torch.tanh(self.max_gate) * out_max |
|
|
|
|
| class ClassInteractionBlock(nn.Module): |
| """Self-attention over the class dimension; both branches start near identity (layer scale).""" |
|
|
| def __init__(self, dim: int, num_heads: int = 4, dropout: float = 0.1, layerscale: float = 1e-4): |
| super().__init__() |
| self.norm = nn.LayerNorm(dim) |
| self.self_attn = nn.MultiheadAttention(dim, num_heads, dropout=dropout, batch_first=True) |
| self.mlp = nn.Sequential(nn.Linear(dim, dim * 2), nn.GELU(), nn.Dropout(dropout), |
| nn.Linear(dim * 2, dim), nn.Dropout(dropout)) |
| self.norm_mlp = nn.LayerNorm(dim) |
| self.ci_gamma_attn = nn.Parameter(torch.full((dim,), float(layerscale))) |
| self.ci_gamma_mlp = nn.Parameter(torch.full((dim,), float(layerscale))) |
|
|
| def forward(self, x): |
| xn = self.norm(x) |
| x = x + self.ci_gamma_attn * self.self_attn(xn, xn, xn)[0] |
| return x + self.ci_gamma_mlp * self.mlp(self.norm_mlp(x)) |
|
|
|
|
| class ClasswiseHead(nn.Module): |
| """Shared Linear-LayerNorm-GELU-Dropout block, then one weight vector and bias per class.""" |
|
|
| def __init__(self, dim: int, num_classes: int, dropout: float = 0.1): |
| super().__init__() |
| self.shared = nn.Sequential(nn.Linear(dim, dim), nn.LayerNorm(dim), nn.GELU(), nn.Dropout(dropout)) |
| self.cls_w = nn.Parameter(torch.empty(num_classes, dim)) |
| self.cls_b = nn.Parameter(torch.zeros(num_classes)) |
| nn.init.trunc_normal_(self.cls_w, std=0.02) |
|
|
| def forward(self, pooled): |
| return (self.shared(pooled) * self.cls_w.unsqueeze(0)).sum(-1) + self.cls_b |
|
|
|
|
| class RadiologyHead(nn.Module): |
| """tokens (B, N, D) -> logits (B, K): class-query pool -> width reduction -> class interaction -> classifier.""" |
|
|
| def __init__(self, dim: int, num_classes: int = 18, pool: str = "mhead", num_heads: int = 4, |
| d_attn: int | None = None, dropout: float = 0.1, head_dim: int | None = None, |
| layerscale: float = 1e-4, identity_proj: bool = True): |
| super().__init__() |
| if pool == "mhead": |
| self.pool = MultiHeadQueryAttentionPool(dim, num_classes, num_heads, d_attn, identity_proj) |
| elif pool == "hybrid": |
| self.pool = GatedMaxTokenPool(dim, num_classes, d_attn) |
| else: |
| raise ValueError(f"unknown pool {pool!r}") |
| w = int(head_dim) if head_dim else dim |
| self.reduce = nn.Sequential(nn.Linear(dim, w), nn.LayerNorm(w)) if w != dim else None |
| self.class_interaction = ClassInteractionBlock(w, num_heads, dropout, layerscale) |
| self.head = ClasswiseHead(w, num_classes, dropout) |
|
|
| def forward(self, x, mask=None): |
| pooled = self.pool(x, mask) |
| if self.reduce is not None: |
| pooled = self.reduce(pooled) |
| return self.head(self.class_interaction(pooled)) |
|
|
|
|
| class CTClassifier(nn.Module): |
| """encoder (LoRA-wrapped) -> slice position encoding -> radiology head.""" |
|
|
| def __init__(self, encoder, dim: int, num_classes: int, slice_posenc: bool = True, **head_kw): |
| super().__init__() |
| self.encoder = encoder |
| base = encoder.base_model.model if hasattr(encoder, "base_model") else encoder |
| self.pos_enc = SlicePosEnc(base.embed_dim, base.max_slices, base.tokens_per_slice) if slice_posenc else None |
| self.head = RadiologyHead(dim, num_classes, **head_kw) |
|
|
| def forward(self, x): |
| tokens, mask = self.encoder(x) |
| if self.pos_enc is not None: |
| tokens = self.pos_enc(tokens) |
| return self.head(tokens, mask) |
|
|
|
|
| |
| HEAD_PREFIXES = ("head.", "pos_enc.") |
| HEAD_SUBSTRINGS = ("conv_branch", "cb_gamma") |
|
|
|
|
| def head_kwargs(cfg: dict) -> dict: |
| return dict(pool=cfg.get("pool", "mhead"), num_heads=int(cfg.get("head_heads") or 4), |
| d_attn=int(cfg["pool_d_attn"]) if cfg.get("pool_d_attn") else None, |
| dropout=float(cfg.get("head_dropout") or 0.1), |
| head_dim=int(cfg["head_dim"]) if cfg.get("head_dim") else None, |
| layerscale=float(cfg.get("head_layerscale") or 0.0), |
| identity_proj=bool(cfg.get("head_identity_proj", True))) |
|
|
|
|
| _LORA_KEY = re.compile(r"^(?P<mod>.+)\.lora_(?:A|B)\.(?:[^.]+\.)?weight$") |
|
|
|
|
| def lora_targets_from_state_dict(sd) -> list[str]: |
| """Module paths that carry LoRA weights in a checkpoint.""" |
| out = set() |
| for k in sd: |
| m = _LORA_KEY.match(k) |
| if m: |
| mod = m.group("mod") |
| i = mod.rfind("base_model.model.") |
| out.add(mod[i + len("base_model.model."):] if i >= 0 else mod) |
| return sorted(out) |
|
|
|
|
| def build_from_config(cfg: dict, state_dict: dict, num_classes: int = 18) -> CTClassifier: |
| """Rebuild a released model from its config (streams.json) and weights.""" |
| from peft import LoraConfig, get_peft_model |
|
|
| encoder = FlexiCT2DEncoder(img_size=int(cfg["img_size"]), patch_size=int(cfg["patch_size"]), |
| max_slices=int(cfg["max_slices"]), |
| slice_microbatch=int(cfg.get("slice_microbatch") or cfg["max_slices"]), |
| conv_branch=bool(cfg.get("conv_branch")), |
| load_pretrained=False) |
| dim = encoder.embed_dim |
| if cfg.get("use_lora"): |
| encoder = get_peft_model(encoder, LoraConfig( |
| r=int(cfg["lora_r"]), lora_alpha=int(cfg["lora_alpha"]), lora_dropout=0.0, |
| bias="none", target_modules=lora_targets_from_state_dict(state_dict))) |
| model = CTClassifier(encoder, dim, num_classes, slice_posenc=bool(cfg.get("slice_posenc")), |
| **head_kwargs(cfg)) |
| model.load_state_dict(state_dict, strict=True) |
| return model.eval() |
|
|
