| """MRI-CORE stem: MRI-pretrained ViT-B feature extractor. |
| |
| Loads the MRI-CORE checkpoint (DINOv2-style ViT-B, SAM-init, pretrained on |
| 6.1M MRI slices) and exposes a timm-compatible surface so the existing |
| MoR/MoE cores can use it as their frozen stem: |
| |
| stem.embed_dim -> 768 |
| stem.num_prefix_tokens -> 1 (CLS only, no registers) |
| stem.pos_embed -> [1, 197, 768] |
| stem.forward_features(x)-> [S, 197, 768] (frozen, eval mode) |
| |
| Key layout in the checkpoint: teacher.backbone.{cls_token, pos_embed, |
| patch_embed.proj, blocks.<g>.<i>, norm} with blocks nested 4 x 3 = 12. |
| """ |
| from __future__ import annotations |
|
|
| import os |
|
|
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
| _CKPT = "MRI_CORE_vitb.pth" |
|
|
|
|
| class Block(nn.Module): |
| def __init__(self, dim: int = 768, n_heads: int = 12, mlp_ratio: float = 4.0): |
| super().__init__() |
| self.norm1 = nn.LayerNorm(dim) |
| self.attn = nn.MultiheadAttention(dim, n_heads, batch_first=True) |
| self.norm2 = nn.LayerNorm(dim) |
| self.mlp = nn.Sequential( |
| nn.Linear(dim, int(dim * mlp_ratio)), |
| nn.GELU(), |
| nn.Linear(int(dim * mlp_ratio), dim), |
| ) |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| x = x + self.attn(self.norm1(x), self.norm1(x), self.norm1(x))[0] |
| x = x + self.mlp(self.norm2(x)) |
| return x |
|
|
|
|
| class MRICoreStem(nn.Module): |
| """Frozen MRI-CORE ViT-B feature extractor (embed_dim=768, 197 tokens).""" |
|
|
| def __init__(self, ckpt_path: str | None = None, freeze: bool = True): |
| super().__init__() |
| if ckpt_path is None: |
| for base in ( |
| r"G:\RSNA-Knee\cache", |
| os.path.join(os.path.dirname(__file__), "..", "..", "weights"), |
| ): |
| cand = os.path.join(base, _CKPT) |
| if os.path.isfile(cand): |
| ckpt_path = cand |
| break |
| if ckpt_path is None or not os.path.isfile(ckpt_path): |
| raise FileNotFoundError(f"MRI-CORE checkpoint not found ({_CKPT})") |
| self.embed_dim = 768 |
| self.num_prefix_tokens = 1 |
| self.patch_embed = nn.Conv2d(3, 768, kernel_size=16, stride=16, bias=True) |
| self.cls_token = nn.Parameter(torch.zeros(1, 1, 768)) |
| self.pos_embed = nn.Parameter(torch.zeros(1, 197, 768)) |
| self.mask_token = nn.Parameter(torch.zeros(1, 768)) |
| self.blocks = nn.ModuleList( |
| [nn.ModuleList([Block() for _ in range(3)]) for _ in range(4)] |
| ) |
| self.norm = nn.LayerNorm(768) |
| st = torch.load(ckpt_path, map_location="cpu", weights_only=True) |
| inner = st["teacher"] |
| if isinstance(inner, dict) and "state_dict" in inner: |
| inner = inner["state_dict"] |
| own = self.state_dict() |
| prefix = "backbone." |
| missing = [] |
| for k in list(inner.keys()): |
| if not k.startswith(prefix): |
| inner.pop(k) |
| continue |
| name = k[len(prefix):] |
| if name not in own: |
| continue |
| if tuple(inner[k].shape) == tuple(own[name].shape): |
| own[name] = inner[k] |
| else: |
| missing.append((name, tuple(inner[k].shape), tuple(own[name].shape))) |
| self.load_state_dict(own) |
| if missing: |
| print(f"mricore: {len(missing)} shape mismatches skipped " |
| f"(e.g. {missing[0]})", flush=True) |
| if freeze: |
| for p in self.parameters(): |
| p.requires_grad_(False) |
| self.eval() |
|
|
| def forward_features(self, x: torch.Tensor) -> torch.Tensor: |
| """[S, 3, H, W] -> [S, 197, 768] (CLS + 196 patch tokens).""" |
| x = self.patch_embed(x).flatten(2).transpose(1, 2) |
| x = torch.cat([self.cls_token.expand(x.size(0), -1, -1), x], dim=1) |
| x = x + self.pos_embed |
| for group in self.blocks: |
| for blk in group: |
| x = blk(x) |
| return self.norm(x) |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| return self.forward_features(x) |
|
|
