Merge pull request #17 from lukasugar/asparagus_evals/segmentation

Asparagus evals - segmentation

Author: clane9

pyproject.toml

@@ -60,3 +60,6 @@ line-length = 100
 
 [tool.ruff.lint]
 ignore = ["F722"]
+
+[tool.pytest.ini_options]
+norecursedirs = ["third_party", ".scratch"]

src/asparagus_bridge/README.md

@@ -143,7 +143,72 @@ can be used as a reference for Task 3 regression runs.
 
 #### Segmentation
 
-TBD
+##### Task 2
+Task 2 is `SEG009_FOMO26_Meningioma`; for sMRI MAE use the FLAIR-only custom variant `SEG009_FOMO26_Meningioma_FLAIR`.
+
+Prepare the raw FOMO task folders and convert them to asparagus tensors:
+
+```sh
+cd "$ASPARAGUS_SOURCE"
+unzip -n Task_2.zip -d Task_2
+
+cd /Users/lukasecerovic/Documents/repos/sMRI/smri-fm
+uv run asp_process \
+  --dataset SEG009_FOMO26_Meningioma_CUSTOM \
+  --task_name SEG009_FOMO26_Meningioma_FLAIR \
+  --modalities flair \
+  --save_as_tensor \
+  --num_workers 4
+```
+
+The segmentation processors write `split_80_10_10.json` and
+`TEST_80_10_10.json`. Override the asparagus segmentation defaults to use
+those splits when finetuning.
+
+Convert the sMRI MAE checkpoint:
+
+```sh
+uv run python -c 'from asparagus_bridge.checkpoint import convert_checkpoint; convert_checkpoint("smri_mae", "runs/mae/checkpoint-last.pth", "runs/mae/asparagus.ckpt")'
+```
+
+Task 2 smoke test:
+
+```sh
+uv run asp_finetune_seg --config-name finetuning/smoke_test_seg_task_2.yaml
+```
+
+##### Task 4
+Task 4 is `SEG010_FOMO26_TrigeminalNeuralgia` and is already single-channel T2w.
+
+```sh
+cd "$ASPARAGUS_SOURCE"
+unzip -n Task_4.zip -d Task_4
+
+cd <repo_root>
+
+uv run asp_process --dataset SEG010 --save_as_tensor --num_workers 4
+```
+
+Task 4 smoke test:
+
+```sh
+uv run asp_finetune_seg --config-name finetuning/smoke_test_seg_task_4.yaml
+```
+
+##### Notes on segmentation
+`SmriMaeSegBackbone` currently inherits asparagus/gardening-tools sliding-window
+inference. That path is compatible with asparagus eval, but it sums overlapping
+logits without overlap-count normalization or Gaussian/Hann weighting. Treat
+that as a known follow-up if segmentation quality near patch borders matters.
+
+Future segmentation variants worth testing:
+
+- Canonical Task 2 two-channel finetuning (`flair`, `dwi`) with explicit
+  multi-channel checkpoint stem adaptation.
+- MAE reconstruction decoder reuse instead of the current Primus-like patch
+  segmentation decoder.
+- Normalized sliding-window blending with overlap-count averaging, Gaussian
+  weighting, or Hann weighting.
 
 #### Linear probing
 

src/asparagus_bridge/configs/finetuning/smoke_test_seg_task_2.yaml

@@ -0,0 +1,38 @@
+# @package _global_
+defaults:
+  - /default_finetune_seg
+  - /model/smri_mae@model
+  - _self_
+
+# hydra:
+#   job_logging:
+#     root:
+#       level: DEBUG
+
+task: SEG009_FOMO26_Meningioma_FLAIR
+
+checkpoint_path: /Users/lukasecerovic/Documents/repos/sMRI/smri-fm/.scratch/pretrained_mae_checkpoint/checkpoint-last-asparagus-NEW.ckpt
+
+data:
+  train_split: split_80_10_10
+  test_split: TEST_80_10_10
+
+hardware:
+  num_workers: 2
+  precision: 32-true
+  compile_mode: null
+
+training:
+  epochs: 1
+  batch_size: 1
+  patch_size: [64, 64, 64]
+  train_batches_per_epoch_per_device: 2
+  val_batches_per_epoch_per_device: 1
+  check_val_every_n_epoch: 1
+
+model:
+  _seg_net:
+    patch_size: 8
+
+logger:
+  wandb_logging: false

src/asparagus_bridge/configs/finetuning/smoke_test_seg_task_4.yaml

@@ -0,0 +1,33 @@
+# @package _global_
+defaults:
+  - /default_finetune_seg
+  - /model/smri_mae@model
+  - _self_
+
+task: SEG010_FOMO26_TrigeminalNeuralgia
+
+checkpoint_path: /Users/lukasecerovic/Documents/repos/sMRI/smri-fm/.scratch/pretrained_mae_checkpoint/checkpoint-last-asparagus-NEW.ckpt
+
+data:
+  train_split: split_80_10_10
+  test_split: TEST_80_10_10
+
+hardware:
+  num_workers: 2
+  precision: 32-true
+  compile_mode: null
+
+training:
+  epochs: 1
+  batch_size: 1
+  patch_size: [64, 64, 64]
+  train_batches_per_epoch_per_device: 2
+  val_batches_per_epoch_per_device: 1
+  check_val_every_n_epoch: 1
+
+model:
+  _seg_net:
+    patch_size: 8
+
+logger:
+  wandb_logging: false

src/asparagus_bridge/configs/model/smri_mae.yaml

@@ -8,10 +8,14 @@ seg_net: smri_mae
 cls_net: smri_mae
 reg_net: smri_mae
 
-# Segmentation: placeholder.
 _seg_net:
   _target_: asparagus_bridge.models_smri_mae.SmriMaeSegBackbone
   dimensions: ${model.dimensions}
+  img_size: ${training.patch_size}
+  patch_size: 16
+  depth: 12
+  embed_dim: 768
+  num_heads: 12
 
 _cls_net:
   _target_: asparagus_bridge.models_smri_mae.SmriMaeClsRegBackbone

src/asparagus_bridge/models_smri_mae.py

@@ -7,6 +7,10 @@
     - weights loaded later by asparagus.BaseModule via load_state_dict(strict=False)
 """
 
+import math
+
+from gardening_tools.modules.networks.BaseNet import BaseNet
+from gardening_tools.modules.networks.components.transformer import PatchDecode
 import torch.nn as nn
 from torch import Tensor
 
@@ -62,18 +66,68 @@ def _features(self, x: Tensor) -> Tensor:
     def forward(self, x: Tensor) -> Tensor:
         """Encoder + head """
         return self.head(self._features(x))
-    
+
     def _encode(self, x: Tensor) -> Tensor:
         """ Encoder output in format used for linear probing"""
         feat = self._features(x)
         return feat[:, :, None, None, None]
 
 
-class SmriMaeSegBackbone(nn.Module):
-    """Placeholder for ViT-based segmentation backbone."""
+class SmriMaeSegBackbone(BaseNet):
+    """MAE ViT segmentation backbone with a Primus-like patch decoder."""
+
+    def __init__(
+        self,
+        input_channels: int,
+        output_channels: int,
+        img_size: int | tuple[int, int, int] = (160, 160, 160),
+        patch_size: int | tuple[int, int, int] = (16, 16, 16),
+        depth: int = 12,
+        embed_dim: int = 768,
+        num_heads: int = 12,
+        dimensions: str = "3D",
+        **_ignored,
+    ):
+        super().__init__()
+        assert dimensions == "3D", f"only 3D supported, got dimensions={dimensions}"
+
+        self.num_classes = output_channels
+        self.stem_weight_name = "encoder.patch_embed.weight"
 
-    def __init__(self, *args, **kwargs):
-        raise NotImplementedError("SmriMaeSegBackbone is not yet implemented")
+        self.encoder = MaskedViT(
+            img_size=img_size,
+            patch_size=patch_size,
+            in_chans=input_channels,
+            depth=depth,
+            embed_dim=embed_dim,
+            num_heads=num_heads,
+            class_token=True,
+        )
+        self.grid_size = self.encoder.patchify.grid_size
+        self.decoder = PatchDecode(
+            patch_size=self.encoder.patchify.patch_size,
+            embed_dim=embed_dim,
+            out_channels=output_channels,
+        )
 
     def forward(self, x: Tensor) -> Tensor:
-        raise NotImplementedError
+        _, _, patch_embeds, _, _ = self.encoder(x)
+        expected_tokens = math.prod(self.grid_size)
+        if patch_embeds.shape[1] != expected_tokens:
+            raise ValueError(
+                "unexpected MAE patch token count: "
+                f"got {patch_embeds.shape[1]}, expected {expected_tokens}"
+            )
+
+        features = patch_embeds.reshape(
+            x.shape[0],
+            *self.grid_size,
+            patch_embeds.shape[-1],
+        )
+        features = features.permute(0, 4, 1, 2, 3).contiguous()
+        return self.decoder(features)
+
+    # Inherits BaseNet.sliding_window_predict. That implementation sums
+    # overlapping logits without normalizing by an overlap-count map.
+    # TODO: replace inherited sliding-window accumulation with normalized overlap averaging.
+    # TODO: consider Gaussian/Hann weighting so patch centers contribute more than patch borders.