Update src\models\mricore_stem.py

src/models/mricore_stem.py

@@ -0,0 +1,108 @@
+"""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)  # [S, 196, 768]
+        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)