Add files using upload-large-folder tool

.gitattributes

@@ -33,3 +33,45 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+checkpoints/pretrain_full_90_10_h100/eval__fomo_val__mask_pred__00096.png filter=lfs diff=lfs merge=lfs -text
+checkpoints/pretrain_full_90_10_h100/eval__fomo_val__mask_pred__00095.png filter=lfs diff=lfs merge=lfs -text
+checkpoints/pretrain_full_90_10_h100/eval__fomo_val__mask_pred__00099.png filter=lfs diff=lfs merge=lfs -text
+checkpoints/pretrain_full_90_10_h100/eval__fomo_val__mask_pred__00098.png filter=lfs diff=lfs merge=lfs -text
+checkpoints/pretrain_full_90_10_h100/eval__fomo_val__mask_pred__00097.png filter=lfs diff=lfs merge=lfs -text
+checkpoints/pretrain_full_90_10_h100/eval__fomo_val__mask_pred__00000.png filter=lfs diff=lfs merge=lfs -text
+checkpoints/pretrain_full_90_10_h100/eval__fomo_val__mask_pred__00001.png filter=lfs diff=lfs merge=lfs -text
+checkpoints/pretrain_full_90_10_h100/eval__fomo_val__mask_pred__00002.png filter=lfs diff=lfs merge=lfs -text
+checkpoints/pretrain_full_90_10_h100/eval__fomo_val__mask_pred__00004.png filter=lfs diff=lfs merge=lfs -text
+checkpoints/pretrain_full_90_10_h100/eval__fomo_val__mask_pred__00003.png filter=lfs diff=lfs merge=lfs -text
+checkpoints/pretrain_full_90_10_h100/eval__fomo_val__mask_pred__00005.png filter=lfs diff=lfs merge=lfs -text
+checkpoints/pretrain_full_90_10_h100/eval__fomo_val__mask_pred__00007.png filter=lfs diff=lfs merge=lfs -text
+checkpoints/pretrain_full_90_10_h100/eval__fomo_val__mask_pred__00006.png filter=lfs diff=lfs merge=lfs -text
+checkpoints/pretrain_full_90_10_h100/eval__fomo_val__mask_pred__00008.png filter=lfs diff=lfs merge=lfs -text
+checkpoints/pretrain_full_90_10_h100/eval__fomo_val__mask_pred__00011.png filter=lfs diff=lfs merge=lfs -text
+checkpoints/pretrain_full_90_10_h100/eval__fomo_val__mask_pred__00009.png filter=lfs diff=lfs merge=lfs -text
+checkpoints/pretrain_full_90_10_h100/eval__fomo_val__mask_pred__00010.png filter=lfs diff=lfs merge=lfs -text
+checkpoints/pretrain_full_90_10_h100/eval__fomo_val__mask_pred__00012.png filter=lfs diff=lfs merge=lfs -text
+checkpoints/pretrain_full_90_10_h100/eval__fomo_val__mask_pred__00013.png filter=lfs diff=lfs merge=lfs -text
+checkpoints/pretrain_full_90_10_h100/eval__fomo_val__mask_pred__00014.png filter=lfs diff=lfs merge=lfs -text
+checkpoints/pretrain_full_90_10_h100/eval__fomo_val__mask_pred__00015.png filter=lfs diff=lfs merge=lfs -text
+checkpoints/pretrain_full_90_10_h100/eval__fomo_val__mask_pred__00016.png filter=lfs diff=lfs merge=lfs -text
+checkpoints/pretrain_full_90_10_h100/eval__fomo_val__mask_pred__00017.png filter=lfs diff=lfs merge=lfs -text
+checkpoints/pretrain_full_90_10_h100/eval__fomo_val__mask_pred__00018.png filter=lfs diff=lfs merge=lfs -text
+checkpoints/pretrain_full_90_10_h100/eval__fomo_val__mask_pred__00019.png filter=lfs diff=lfs merge=lfs -text
+checkpoints/pretrain_full_90_10_h100/eval__fomo_val__mask_pred__00020.png filter=lfs diff=lfs merge=lfs -text
+checkpoints/pretrain_full_90_10_h100/eval__fomo_val__mask_pred__00021.png filter=lfs diff=lfs merge=lfs -text
+checkpoints/pretrain_full_90_10_h100/eval__fomo_val__mask_pred__00023.png filter=lfs diff=lfs merge=lfs -text
+checkpoints/pretrain_full_90_10_h100/eval__fomo_val__mask_pred__00024.png filter=lfs diff=lfs merge=lfs -text
+checkpoints/pretrain_full_90_10_h100/eval__fomo_val__mask_pred__00022.png filter=lfs diff=lfs merge=lfs -text
+checkpoints/pretrain_full_90_10_h100/eval__fomo_val__mask_pred__00026.png filter=lfs diff=lfs merge=lfs -text
+checkpoints/pretrain_full_90_10_h100/eval__fomo_val__mask_pred__00025.png filter=lfs diff=lfs merge=lfs -text
+checkpoints/pretrain_full_90_10_h100/eval__fomo_val__mask_pred__00027.png filter=lfs diff=lfs merge=lfs -text
+checkpoints/pretrain_full_90_10_h100/eval__fomo_val__mask_pred__00029.png filter=lfs diff=lfs merge=lfs -text
+checkpoints/pretrain_full_90_10_h100/eval__fomo_val__mask_pred__00028.png filter=lfs diff=lfs merge=lfs -text
+checkpoints/pretrain_full_90_10_h100/eval__fomo_val__mask_pred__00032.png filter=lfs diff=lfs merge=lfs -text
+checkpoints/pretrain_full_90_10_h100/eval__fomo_val__mask_pred__00030.png filter=lfs diff=lfs merge=lfs -text
+checkpoints/pretrain_full_90_10_h100/eval__fomo_val__mask_pred__00031.png filter=lfs diff=lfs merge=lfs -text
+checkpoints/pretrain_full_90_10_h100/eval__fomo_val__mask_pred__00033.png filter=lfs diff=lfs merge=lfs -text
+checkpoints/pretrain_full_90_10_h100/eval__fomo_val__mask_pred__00034.png filter=lfs diff=lfs merge=lfs -text
+checkpoints/pretrain_full_90_10_h100/eval__fomo_val__mask_pred__00035.png filter=lfs diff=lfs merge=lfs -text
+checkpoints/pretrain_full_90_10_h100/eval__fomo_val__mask_pred__00037.png filter=lfs diff=lfs merge=lfs -text

checkpoints/pretrain_full_90_10_h100/config.yaml

@@ -0,0 +1,69 @@
+name: pretrain_full_90_10_h100
+notes: null
+output_dir: checkpoints/pretrain_full_90_10_h100
+img_size:
+- 208
+- 240
+- 208
+patch_size: 8
+mask_ratio: 0.8
+pred_mask_ratio: null
+model: mae_vit_large
+model_kwargs:
+  decoding: attn
+  target_norm: none
+  no_decode_pos: false
+  mask_drop_scale: false
+  class_token: true
+  reg_tokens: 0
+  no_embed_class: false
+  decoder_depth: 4
+  drop_path_rate: 0.0
+datasets:
+  fomo_train:
+    url: datasets/FOMO300/wds/shard.{000000..001020}.tar
+    samples_per_epoch: 153000
+    shuffle: true
+    buffer_size: 8000
+    drop_last: true
+  fomo_val:
+    url: datasets/FOMO300/wds/shard.{001021..001134}.tar
+    samples_per_epoch: 17000
+    shuffle: false
+    buffer_size: 1000
+    drop_last: true
+train_dataset: fomo_train
+eval_datasets:
+- fomo_val
+num_workers: 16
+prefetch_factor: 8
+presend_cuda: false
+epochs: 100
+batch_size: 64
+accum_iter: 1
+base_lr: 0.001
+min_lr: 1.0e-06
+warmup_epochs: 10
+weight_decay: 0.05
+betas:
+- 0.9
+- 0.95
+clip_grad: 1.0
+amp: true
+amp_dtype: bfloat16
+compile: false
+ckpt: null
+resume: false
+auto_resume: true
+start_epoch: 0
+checkpoint_period: 5
+max_checkpoints: 5
+r2_sync: null
+device: cuda
+distributed: false
+seed: 7338
+eval_seed: 7338
+debug: false
+wandb: true
+wandb_entity: null
+wandb_project: smri-fm

finetune/fomo_tune_baseline/output/task1/build/Apptainer.def

@@ -0,0 +1,31 @@
+Bootstrap: docker
+From: python:3.11-slim
+
+# NOT buildable where it sits: the %files paths below are relative to the build cwd, which is the
+# staging dir `build.py` writes. Build it with `python -m fomo_tune.build <run_dir>`, not by
+# pointing apptainer at this file.
+#
+# Versions are pinned to the training environment: numpy, scikit-learn and joblib because they
+# unpickle `head.joblib`, torch because that is what the checkpoint was written by.
+
+%files
+    fomo_tune /app/fomo_tune
+    smri_mae /app/smri_mae
+    model /app/model
+    predict.py /app/predict.py
+
+%post
+    pip install --no-cache-dir \
+        torch==2.8.0 \
+        numpy==2.4.6 \
+        nibabel==5.4.2 \
+        einops==0.8.2 \
+        jaxtyping==0.3.10 \
+        timm==1.0.27 \
+        huggingface-hub==0.36.2 \
+        scikit-learn==1.8.0 \
+        joblib==1.5.3 \
+        omegaconf==2.3.0
+
+%runscript
+    exec python /app/predict.py "$@"

finetune/fomo_tune_baseline/output/task1/build/fomo_tune/README.md

@@ -0,0 +1,239 @@
+# fomo_tune
+
+The five FOMO26 challenge tasks, one script each, tuned independently.
+
+This is a spinoff of `nanobrain.eval`, which scored every backbone on every task through one fixed
+probe. That was the right shape for a benchmark and the wrong shape for a competition: here we care
+about one backbone (sMRI MAE) and five scores, and each task wants a different method. **Nothing
+here imports `nanobrain.eval`, and it should stay that way** — this package may be shared with
+people who won't get the eval suite.
+
+## Layout
+
+| File | |
+|---|---|
+| `datasets.py` | core, **frozen**. One `load_fomo_task<k>()` per task, streaming the challenge zips into an HF dataset. Raw niftis, no resampling — the backbone transform does that. |
+| `backbone.py` | core, **frozen**. `load_backbone(ckpt_path) -> (SmriMaeBackbone, SmriMaeTransform)`. Frozen sMRI MAE encoder; the transform canonicalizes to RAS, rescales to 1mm, fits to the pretraining shape, z-scores in a mean-threshold brain mask. |
+| `utils.py` | core. `set_seed`, `git_sha`, `setup_logging`. |
+| `main_task<k>.py` | shell. One task, end to end. Task 1 is the worked example; copy it. |
+| `build.py` + `Apptainer.def` | shell. Package a run dir into the challenge `.sif`. Shared by every task. |
+
+`datasets.py` and `backbone.py` are settled and their caches are warm. Treat them as read-only:
+new work goes in `main_task<k>.py`. If one of them genuinely needs to change, that is a
+conversation first, because it invalidates every score already recorded.
+
+## The pattern
+
+`main_task1.py` is in three sections, and the split is the point of the whole design.
+
+**`Task1Method` — the part we tune.** Features, head, hyperparameters, anything that might move
+the score. Its interface is:
+
+```python
+method.fit(rows)              # rows are dataset records: subject, label, images
+method.predict(images)        # -> the challenge's output for one subject
+method.save(model_dir)        # config.yaml + head.joblib
+Task1Method.load(model_dir, **overrides)
+```
+
+**The protocol — fixed.** Pool out-of-fold predictions over all subjects, bootstrap subjects for
+the CI. No repeats, no stratification; the bootstrap is the only variance estimate. Splitting is
+per-task but fixed within a task — leave-one-out where n is tiny (task 1, n=21), **20-fold** where
+it isn't (tasks 3 and 5), which is close enough to LOO without paying for 494 refits. Once a task's
+scheme is set, hold it or scores stop being comparable across iterations. That is also why
+`cross_validate` seeds its shuffle at 0 rather than from `cfg.seed`: the folds are part of the
+protocol, so tuning the run's seed must not silently redraw them.
+
+**Two entrypoints.** `train` runs the protocol then fits a head on all subjects and saves it;
+`predict` is the challenge CLI. Both go through `Method.predict`, which is why every fold
+exercises the code the submission will run.
+
+That last point is the load-bearing one. `predict` is not a wrapper written at packaging time — it
+is the same call cross-validation already made once per held-out subject. When you add a task, keep
+that property.
+
+```bash
+uv run python -m fomo_tune.main_task1 train modalities=[dwi_b1000,flair] name=task1_dwi_flair
+uv run python -m fomo_tune.main_task1 predict --model-dir output/fomo_tune/task1_dwi/model \
+    --adc adc.nii.gz --dwi dwi.nii.gz --flair flair.nii.gz --output prob.txt
+```
+
+`train` takes omegaconf dotlist overrides against the `Config` dataclass at the top of the file.
+It writes `config.yaml`, `log.txt`, `metrics.json`, and `model/` into `{output_root}/{name}/`.
+
+## Status
+
+Tasks 1, 5 and 3 are drafted and verified. Task 1 is also packaged — its container passes the
+challenge validator; 5 and 3 have not been built yet. **Tasks 2 and 4 are tabled** — both are
+segmentation, both need `predict` to emit a nifti on the input grid, and neither is worth opening
+until the classification and regression tasks are settled.
+
+All three on `vitl_fomo300`, one H100, wall being the cross-validation loop:
+
+| run | result | wall |
+|---|---|---|
+| `task1_dwi`, dwi_b1000, n=21, LOO | AUROC **0.990** [0.944, 1.000] | 25s |
+| `task5_t1w`, t1w, n=48, 20-fold | AUROC **0.984** [0.953, 1.000] | 73s |
+| `task3_t1w`, t1w, n=494, 20-fold | r **0.962** [0.956, 0.968], MAE **3.71y** [3.45, 3.97] | 260s |
+
+**Task 3's row is one fold-seed stale.** It was measured before `cross_validate` froze its shuffle
+at 0, so it is a 20-fold run with `random_state=4466`. Task 1 (LOO) and task 5 are current. The
+re-run is cheap — 260s on a GPU — it just has not been done. Expect a shift of the same order task
+5 saw when its folds moved (0.948 → 0.984, i.e. inside the CI but not negligible).
+
+Task 1's earlier probe sweep got 0.954 [0.861, 1.000] on the same checkpoint
+(`experiments/eval_global_0728`), so it roughly reproduces — the gap is LOO vs 5×5 stratified CV,
+one interpolation instead of two, and a head selected on AUROC instead of balanced accuracy.
+
+Two checks worth repeating per task — `.claude/scratch/verify_task1.py` and
+`.claude/scratch/verify_task35.py <k>` do both:
+- features are **bit-identical** whether the nifti comes from the HF dataset wrapper or from
+  `nib.load` off disk, so CV numbers transfer to the container
+- the `predict` CLI agrees with the in-process method
+
+## What changes per task
+
+Counts and modalities, read from the local zips:
+
+| Task | n | Inputs | Output | Split | Notes |
+|---|---|---|---|---|---|
+| 1 infarct | 21 | adc, dwi_b1000, flair (+t2s/swi) | probability | LOO | done |
+| 5 polymicrogyria | 48 | t1w | probability | 20-fold | done |
+| 3 brain age | 494 | t1w | age in years | 20-fold | done — RidgeCV head, scored by **Pearson r and MAE**, each with its own bootstrap CI |
+| 2 meningioma | 23 | dwi_b1000, flair (+t2s/swi) | mask, input grid | — | tabled |
+| 4 trigeminal | 40 | t2w | mask, labels 1=nerve 2=vessel | — | tabled |
+
+Tasks 5 and 3 diverge from task 1 only where that table says. `cross_validate` over a shuffled
+`KFold` replaces `leave_one_out`; both take one modality, so `features` loses the
+concat-over-modalities loop and `Config` loses `modalities`; the challenge CLI flag is `--t1` for
+both, and it is `--t1` for task 3 too even though the file in the zip is `t1w.nii.gz`.
+
+Task 3 is the first regression, so its `score` loops over the two metrics rather than returning
+one, and it drops task 1's guard against bootstrap resamples with fewer than two distinct labels.
+The analogous degenerate case for regression is a resample with no spread in `y`, where Pearson r
+is undefined rather than merely unstable — at n=494 it does not happen.
+
+When tasks 2 and 4 come back: `predict` must write a nifti on the input's grid, and the method
+needs localized features rather than a pooled vector — `backbone.forward` returns `patch_coords`
+in world mm for exactly that. Task 4's label order (1=nerve, 2=vessel) is still a guess and needs
+confirming against the challenge data before per-class numbers mean anything.
+
+## Gotchas
+
+**Raw niftis are on disk** at `data/fomo_eval/Task_<k>/preprocessed/<sub>/ses-01/`, which is the
+easy way to exercise `predict` on a real file rather than one written out of the dataset:
+
+```bash
+uv run python -m fomo_tune.main_task1 predict \
+    --model-dir output/fomo_tune/task1_dwi/model \
+    --adc  data/fomo_eval/Task_1/preprocessed/sub-20/ses-01/adc.nii.gz \
+    --dwi  data/fomo_eval/Task_1/preprocessed/sub-20/ses-01/dwi_b1000.nii.gz \
+    --flair data/fomo_eval/Task_1/preprocessed/sub-20/ses-01/flair.nii.gz \
+    --output /tmp/prob.txt
+```
+
+Task 5 breaks the naming: `Task_5/preprocessed/sub_01/ses_01/t1.nii.gz` — underscores throughout,
+and `t1` not `t1w`. `datasets.py` already handles it; anything you write by hand won't.
+
+```bash
+uv run python -m fomo_tune.main_task5 predict --model-dir output/fomo_tune/task5_t1w/model \
+    --t1 data/fomo_eval/Task_5/preprocessed/sub_01/ses_01/t1.nii.gz --output /tmp/prob.txt
+```
+
+**Volumes are wildly anisotropic.** Task 1's DWI is 0.46×0.46×**5.6**mm, so the transform
+upsamples z by 5.6× to reach 1mm iso. Nothing is wrong, but don't read the 1mm grid as real
+resolution.
+
+**The backbone never saw skull or neck.** Pretraining used a SynthSeg brain mask; the transform
+substitutes a mean-intensity threshold, which keeps both. Known fidelity gap — see
+`.claude/memory/smri-mae-preprocessing-gap.md`.
+
+**Probabilities are not calibrated.** `LogisticRegressionCV` on ~20 samples × 1024 features shrinks
+hard; task 1's out-of-fold probabilities all land in 0.48–0.52 with near-perfect ranking. Fine for
+AUROC, which is what the challenge scores, but don't read them as probabilities. Task 5's do span
+0–1, which is n=48 rather than n=21 and not evidence of calibration.
+
+**n is tiny, so the CI is the result.** Task 1's is ~0.06 wide at the top of the range. Most tuning
+deltas you chase will be inside it. `.claude/NOTES.md` thread 1 has the longer argument.
+
+**GPUs need an allocation** — the login node has no driver. See the `gpu-session` skill.
+
+## Submission
+
+`build.py` packages a run dir into the `.sif` the challenge wants. One command, taking the run dir
+the shipped head was saved into:
+
+```bash
+uv run python -m fomo_tune.build output/fomo_tune/task1_dwi
+```
+
+It stages `/app`, then builds from there:
+
+```
+/app/predict.py          # shim: calls fomo_tune.main_task<k> predict
+/app/model/config.yaml   # from the run dir
+/app/model/head.joblib   # from the run dir
+/app/model/backbone.pth  # stripped checkpoint, --ckpt-path points here
+```
+
+**Both `build.py` and `Apptainer.def` are shared across tasks**, which they can be because nothing
+in staging or in the dependency list is task-specific. The one thing that does vary is the module
+the shim imports, and that comes from `task` in the run's saved config — so a run dir knows which
+task it belongs to, and `build.py` never needs telling.
+
+`predict.py` is **generated at build time** rather than checked in. It is eight lines whose whole
+meaning is the container layout staged around it, so there is nowhere outside a container to run
+it. This does not weaken the point above about `predict` not being written at packaging time: the
+logic still lives in `main_task<k>.py`, exercised once per fold, and the shim only picks the
+subcommand and two paths.
+
+**`Apptainer.def` is not buildable where it sits.** Its `%files` paths are relative to the build
+cwd, which is the staging dir. Pointing `apptainer build` at it in the repo fails confusingly; go
+through `build.py`.
+
+The run dir deliberately does *not* carry backbone weights — that checkpoint is 3.9G and would be
+copied on every run. `--ckpt-path` overrides what `config.yaml` recorded, so the saved config stays
+a faithful record of what trained rather than being rewritten at build time.
+
+**The staged checkpoint is stripped to `model` and `args`**, which is 3.9G → 1.3G because the rest
+is optimizer state inference never reads. `load_backbone` needs no change for this, and `predict`
+gives a bit-identical probability either way (0.524739 on `sub-20`, checked on GPU).
+
+**The base image is `python:3.11-slim`, not a CUDA image.** The PyPI torch wheel *is* the cu128
+build and vendors the whole CUDA userspace as `nvidia-*` packages, so all the container needs from
+the host is the driver, which `--nv`/`--nvccli` binds in. That keeps the SIF at 5.3G (4.0G of
+image, 1.3G of weights) against roughly double for `pytorch/pytorch` and far more for NGC.
+Versions are pinned to the training environment
+mostly so `head.joblib` unpickles against the numpy/sklearn that wrote it.
+
+Apptainer caches the bootstrap layers but **always re-runs `%post`**, so every build re-downloads
+~3G of wheels. If that gets annoying, bake a deps-only base SIF and `Bootstrap: localimage` off it.
+
+### Validating
+
+`third_party/container-validator` is the challenge's own validator, test niftis included:
+
+```bash
+python third_party/container-validator/container_validator/validate.py \
+    --task task1 --sif output/fomo_tune/task1_dwi/task1.sif
+```
+
+It runs `python /app/predict.py --flair /input/… --adc … --dwi … --swi … --output /output/<sid>.txt`
+inside an `apptainer instance` with `/input`, `/output` and `/tmp` bound — which is exactly the
+shim's contract, so nothing in `predict.py` is guessing at the interface.
+
+One thing it does that is easy to miss: it takes GPU via `--nvccli` rather than `--nv`, and one of
+its tests runs `nvidia-smi -L` **inside** the container. `python:3.11-slim` ships no `nvidia-smi`,
+so that test passes only because `--nvccli` injects the host one — a CUDA base image would hide
+that dependency rather than remove it.
+
+**The `task1_dwi` container passes all 20 validator tests**, and `predict` inside it returns
+0.524739 on `sub-20`, identical to the same call outside the container. So the packaging is
+verified end to end, not just built.
+
+**Run it on a compute node with apptainer, which as of 2026-08-11 means `n-6`** — `salloc
+--nodelist=n-6`. The other nodes fail the validator's preflight. The login node has apptainer but
+no driver, and
+`predict` there dies inside `can_use_cudnn_attention` — the jagged-SDPA path reaches into CUDA even
+when the tensors are on CPU, so a driver-less host fails at the forward pass rather than falling
+back. That is the CPU gap worth remembering; it is not a container problem.

finetune/fomo_tune_baseline/output/task1/build/fomo_tune/backbone.py

@@ -0,0 +1,153 @@
+import inspect
+
+import nibabel as nib
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+from torch import Tensor
+
+import smri_mae.model_mae as models_mae
+
+
+class SmriMaeBackbone(nn.Module):
+    grid_coords: Tensor
+
+    def __init__(self, encoder: models_mae.MaskedEncoder):
+        super().__init__()
+        self.encoder = encoder
+        self.img_size = self.encoder.patchify.img_size
+
+        grid_size = self.encoder.patchify.grid_size
+        patch_size = np.array(self.encoder.patchify.patch_size)
+        grid_coords = rearrange(np.indices(grid_size), "c x y z -> (x y z) c")
+        grid_coords = grid_coords * patch_size + (patch_size - 1) / 2
+        grid_coords = torch.as_tensor(grid_coords, dtype=torch.float32)
+        self.register_buffer("grid_coords", grid_coords)
+
+    def forward(self, batch: dict[str, Tensor]) -> dict[str, Tensor]:
+        images = batch["image"]
+        mask = batch["mask"]
+        affine = batch["affine"]
+
+        B, C, X, Y, Z = images.shape
+        assert (X, Y, Z) == self.img_size, f"expected {self.img_size}, got {(X, Y, Z)}"
+
+        _, _, patch_embeds, _, patch_ids, token_mask = self.encoder(images, mask=mask)
+
+        # [B, L, 3] world xyz coords of embeddings
+        patch_coords = self.grid_coords[patch_ids, :]
+        rot = affine[:, :3, :3]
+        trans = affine[:, :3, 3]
+        patch_coords = patch_coords @ rot.transpose(1, 2) + trans[:, None, :]
+
+        return {
+            "patch_embeds": patch_embeds,
+            "patch_ids": patch_ids,
+            "token_mask": token_mask,
+            "patch_coords": patch_coords,
+        }
+
+
+class SmriMaeTransform:
+    def __init__(
+        self,
+        img_size: tuple[int, int, int] = (208, 240, 208),
+        spacing: tuple[float, float, float] = (1.0, 1.0, 1.0),
+    ):
+        self.img_size = img_size
+        self.spacing = spacing
+
+    def __call__(self, img: nib.Nifti1Image) -> dict[str, Tensor]:
+        # repack image to handle incomplete hf Nifti interface
+        img = nib.Nifti1Image(img.dataobj, img.affine, img.header)
+        img = nib.as_closest_canonical(img)
+
+        data = torch.from_numpy(np.ascontiguousarray(img.get_fdata(dtype=np.float32)))
+        affine = np.asarray(img.affine)
+
+        spacing = img.header.get_zooms()
+        if max(abs(s - s_) for s, s_ in zip(spacing, self.spacing)) > 0.05:
+            data, affine = rescale(data, affine, spacing, self.spacing)
+
+        data, affine = fit_to_shape(data, affine, target_shape=self.img_size)
+
+        # mean threshold, not the SynthSeg mask used in pretraining, so skull and neck stay in
+        mask = data > data.mean()
+        brain = data[mask]
+        mean = brain.mean()
+        # population std (correction=0) to match the pretraining normalization
+        std = brain.std(correction=0).clamp_min(1e-6)
+        data = torch.where(mask, (data - mean) / std, 0.0)
+
+        return {
+            "image": data.unsqueeze(0),
+            "mask": mask.unsqueeze(0),
+            "affine": torch.as_tensor(affine, dtype=torch.float32),
+        }
+
+
+def rescale(
+    x: torch.Tensor,
+    affine: np.ndarray,
+    spacing: tuple[float, ...],
+    target_spacing: tuple[float, ...] = (1.0, 1.0, 1.0),
+) -> tuple[torch.Tensor, np.ndarray]:
+    scales = tuple([current / target for current, target in zip(spacing, target_spacing)])
+    resampled = F.interpolate(x[None, None], scale_factor=scales, mode="trilinear").squeeze(0, 1)
+
+    # align_corners=False reads output voxel j from input voxel (j + 0.5) / scale - 0.5
+    scale = np.asarray(scales, dtype=float)
+    step = np.diag([*(1 / scale), 1.0])
+    step[:3, 3] = 0.5 / scale - 0.5
+    return resampled, affine @ step
+
+
+def fit_to_shape(
+    x: torch.Tensor, affine: np.ndarray, target_shape: tuple[int, ...]
+) -> tuple[torch.Tensor, np.ndarray]:
+    """Centre the volume in `target_shape`, padding the short axes and cropping the long ones."""
+    pads = [target - size for size, target in zip(x.shape, target_shape)]
+    padding = [side for pad in reversed(pads) for side in (pad // 2, pad - pad // 2)]
+
+    # a crop is a negative pad, so output voxel k came from input voxel k - pad // 2 either way
+    step = np.eye(4)
+    step[:3, 3] = [-(pad // 2) for pad in pads]
+    return F.pad(x, padding), affine @ step
+
+
+def resolve_ckpt(ckpt_path: str) -> str:
+    """A local path for a checkpoint, downloading it if it is an hf://<org>/<repo>/<file> URI."""
+    from huggingface_hub import hf_hub_download
+
+    if ckpt_path.startswith("hf://"):
+        org, repo, *rest = ckpt_path.removeprefix("hf://").split("/")
+        return hf_hub_download(f"{org}/{repo}", "/".join(rest))
+
+    return ckpt_path
+
+
+def load_backbone(ckpt_path: str) -> tuple[SmriMaeBackbone, SmriMaeTransform]:
+    path = resolve_ckpt(ckpt_path)
+    ckpt = torch.load(path, map_location="cpu", weights_only=True, mmap=True)
+    args = ckpt["args"]
+
+    model_fn = models_mae.__dict__[args["model"]]
+    model: models_mae.MaskedAutoencoderViT = model_fn(
+        img_size=args["img_size"],
+        in_chans=args.get("in_chans", 1),
+        patch_size=args["patch_size"],
+        # older checkpoints carry training flags the current model_mae no longer takes
+        **filter_kwargs(models_mae.MaskedAutoencoderViT, args.get("model_kwargs") or {}),
+    )
+    model.load_state_dict(ckpt["model"])
+    backbone = SmriMaeBackbone(model.encoder)
+    transform = SmriMaeTransform(img_size=args["img_size"])
+    return backbone, transform
+
+
+def filter_kwargs(func, kwargs):
+    signature = inspect.signature(func)
+    kwargs = {k: v for k, v in kwargs.items() if k in signature.parameters}
+    return kwargs

finetune/fomo_tune_baseline/output/task1/build/fomo_tune/datasets.py

@@ -0,0 +1,205 @@
+import os
+import shutil
+import tempfile
+import zipfile
+from collections.abc import Generator
+from contextlib import contextmanager
+from pathlib import Path
+
+import fsspec
+from datasets import Dataset, Features, Nifti, Value
+
+FOMO_EVAL_BASE_URL = os.getenv(
+    "FOMO_EVAL_BASE_URL",
+    "https://sid.erda.dk/share_redirect/fmeuvo1EdF",
+)
+FOMO_EVAL_TASK5_URL = os.getenv(
+    "FOMO_EVAL_TASK5_URL",
+    "https://huggingface.co/datasets/medarc/smri-fm/resolve/main/fomo_eval/Task_5.zip",
+)
+
+
+@contextmanager
+def open_zip(url: str) -> Generator[zipfile.ZipFile, None, None]:
+    """Open a task zip, copying a remote url to a temp file first."""
+    with tempfile.TemporaryDirectory() as tmp:
+        local = Path(url)
+        if not local.exists():
+            local = Path(tmp) / "task.zip"
+            with fsspec.open(url) as src, local.open("wb") as dst:
+                shutil.copyfileobj(src, dst)
+        with zipfile.ZipFile(local) as zf:
+            yield zf
+
+
+def subject_ids(zf: zipfile.ZipFile) -> list[str]:
+    return sorted({name.split("/")[2] for name in zf.namelist() if name.endswith(".nii.gz")})
+
+
+# ---- Task 1: acute infarct (classification; positives also carry a lesion mask) --------
+
+
+def load_fomo_task1() -> Dataset:
+    # No 4th modality: it is swi on 16 subjects and t2s on the other 5.
+    suffixes = ("adc", "dwi_b1000", "flair")
+    features = Features(
+        {
+            "subject": Value("string"),
+            "label": Value("int32"),
+            **{suffix: Nifti() for suffix in suffixes},
+        }
+    )
+    dataset = Dataset.from_generator(
+        _fomo_task1_generator,
+        features=features,
+        gen_kwargs={"suffixes": suffixes},
+        writer_batch_size=16,
+    )
+    return dataset
+
+
+def _fomo_task1_generator(suffixes: tuple[str, ...]):
+    url = f"{FOMO_EVAL_BASE_URL}/Task_1.zip"
+    with open_zip(url) as zf:
+        for sub in subject_ids(zf):
+            label = int(zf.read(f"Task_1/labels/{sub}/ses-01/label.txt").strip())
+            sample = {"subject": sub, "label": label}
+            for suffix in suffixes:
+                name = f"Task_1/preprocessed/{sub}/ses-01/{suffix}.nii.gz"
+                sample[suffix] = {"path": None, "bytes": zf.read(name)}
+            yield sample
+
+
+# ---- Task 2: meningioma segmentation ---------------------------------------------------
+
+
+def load_fomo_task2() -> Dataset:
+    # No 4th modality: it is t2s on 15 subjects and swi on the other 8.
+    suffixes = ("dwi_b1000", "flair")
+    features = Features(
+        {
+            "subject": Value("string"),
+            **{suffix: Nifti() for suffix in suffixes},
+            "seg": Nifti(),
+        }
+    )
+    dataset = Dataset.from_generator(
+        _fomo_task2_generator,
+        features=features,
+        gen_kwargs={"suffixes": suffixes},
+        writer_batch_size=16,
+    )
+    return dataset
+
+
+def _fomo_task2_generator(suffixes: tuple[str, ...]):
+    url = f"{FOMO_EVAL_BASE_URL}/Task_2.zip"
+    with open_zip(url) as zf:
+        for sub in subject_ids(zf):
+            sample = {"subject": sub}
+            for suffix in suffixes:
+                name = f"Task_2/preprocessed/{sub}/ses-01/{suffix}.nii.gz"
+                sample[suffix] = {"path": None, "bytes": zf.read(name)}
+            # Seg is on the image grid (shapes match) but its affine differs by up to 0.03mm.
+            name = f"Task_2/labels/{sub}/ses-01/seg.nii.gz"
+            sample["seg"] = {"path": None, "bytes": zf.read(name)}
+            yield sample
+
+
+# ---- Task 3: brain age regression ------------------------------------------------------
+
+
+def load_fomo_task3() -> Dataset:
+    features = Features(
+        {
+            "subject": Value("string"),
+            "age": Value("int32"),
+            "t1w": Nifti(),
+        }
+    )
+    dataset = Dataset.from_generator(
+        _fomo_task3_generator,
+        features=features,
+        writer_batch_size=16,
+    )
+    return dataset
+
+
+def _fomo_task3_generator():
+    url = f"{FOMO_EVAL_BASE_URL}/Task_3.zip"
+    with open_zip(url) as zf:
+        for sub in subject_ids(zf):
+            age = int(zf.read(f"Task_3/labels/{sub}/ses-01/labels.txt").strip())
+            image_gz = zf.read(f"Task_3/preprocessed/{sub}/ses-01/t1w.nii.gz")
+            sample = {
+                "subject": sub,
+                "age": age,
+                "t1w": {"path": None, "bytes": image_gz},
+            }
+            yield sample
+
+
+# ---- Task 4: trigeminal nerve/vessel segmentation --------------------------------------
+
+
+def load_fomo_task4() -> Dataset:
+    # Volumes are uncropped 0.5mm near-iso, ~360x512x512; crop before feeding a model.
+    features = Features(
+        {
+            "subject": Value("string"),
+            "t2w": Nifti(),
+            "seg": Nifti(),
+        }
+    )
+    dataset = Dataset.from_generator(
+        _fomo_task4_generator,
+        features=features,
+        writer_batch_size=16,
+    )
+    return dataset
+
+
+def _fomo_task4_generator():
+    url = f"{FOMO_EVAL_BASE_URL}/Task_4.zip"
+    with open_zip(url) as zf:
+        for sub in subject_ids(zf):
+            image_gz = zf.read(f"Task_4/preprocessed/{sub}/ses-01/t2w.nii.gz")
+            seg_gz = zf.read(f"Task_4/labels/{sub}/ses-01/seg.nii.gz")
+            sample = {
+                "subject": sub,
+                "t2w": {"path": None, "bytes": image_gz},
+                "seg": {"path": None, "bytes": seg_gz},
+            }
+            yield sample
+
+
+# ---- Task 5: polymicrogyria classification ---------------------------------------------
+
+
+def load_fomo_task5() -> Dataset:
+    features = Features(
+        {
+            "subject": Value("string"),
+            "label": Value("int32"),
+            "t1w": Nifti(),
+        }
+    )
+    dataset = Dataset.from_generator(
+        _fomo_task5_generator,
+        features=features,
+        writer_batch_size=16,
+    )
+    return dataset
+
+
+def _fomo_task5_generator():
+    with open_zip(FOMO_EVAL_TASK5_URL) as zf:
+        for sub in subject_ids(zf):
+            label = int(zf.read(f"Task_5/labels/{sub}/ses_01/labels.txt").strip())
+            image_gz = zf.read(f"Task_5/preprocessed/{sub}/ses_01/t1.nii.gz")
+            sample = {
+                "subject": sub,
+                "label": label,
+                "t1w": {"path": None, "bytes": image_gz},
+            }
+            yield sample

finetune/fomo_tune_baseline/output/task1/build/fomo_tune/main_task1.py

@@ -0,0 +1,253 @@
+"""FOMO task 1: acute infarct classification, scored by AUROC as the challenge scores it.
+
+`Task1Method` is the part we tune -- features, head, hyperparameters. The protocol below it is
+fixed so scores stay comparable across iterations: leave one subject out, pool the out-of-fold
+predictions, bootstrap subjects for the CI.
+
+`train` runs that protocol then fits and saves a head; `predict` is the challenge contract,
+modality paths in and one probability out. Both go through `Task1Method.predict`, so every fold
+exercises the path the submission will run.
+"""
+
+import argparse
+import json
+import logging
+import time
+from dataclasses import dataclass, field
+from pathlib import Path
+
+import joblib
+import nibabel as nib
+import numpy as np
+import torch
+from omegaconf import OmegaConf
+from sklearn.linear_model import LogisticRegressionCV
+from sklearn.metrics import roc_auc_score
+from sklearn.pipeline import make_pipeline
+from sklearn.preprocessing import StandardScaler
+
+from fomo_tune.backbone import load_backbone
+from fomo_tune.utils import git_sha, set_seed, setup_logging
+
+logger = logging.getLogger("fomo_tune")
+
+Images = dict[str, nib.Nifti1Image]
+
+
+@dataclass
+class Config:
+    task: str = "task1"
+    ckpt_path: str = (
+        "/data/mihir-stuff/smri-pretrained/pretrain_full_90_10_h100/checkpoint-last.pth"
+    )
+    modalities: list[str] = field(default_factory=lambda: ["dwi_b1000"])
+    output_root: str = "output/fomo_tune"
+    name: str = "task1"
+    device: str = "cuda"
+    seed: int = 4466
+
+
+# ---- method: the part we tune -----------------------------------------------------------
+
+
+class Task1Method:
+    """Frozen sMRI MAE, mean-pooled tokens per modality concatenated, logistic head."""
+
+    def __init__(self, cfg: Config):
+        self.cfg = cfg
+        self.backbone, self.transform = load_backbone(cfg.ckpt_path)
+        self.device = torch.device(cfg.device)
+        self.backbone.to(self.device).eval().requires_grad_(False)
+        self.modalities = list(cfg.modalities)
+        self.cache: dict[str, np.ndarray] = {}
+        self.head = None
+
+    @torch.inference_mode()
+    def features(self, images: Images) -> np.ndarray:
+        """(D,) per subject. A pure function of the images, so training and inference agree."""
+        pooled = []
+        for modality in self.modalities:
+            sample = self.transform(images[modality])
+            batch = {key: value[None].to(self.device) for key, value in sample.items()}
+
+            with torch.autocast("cuda", torch.bfloat16, enabled=self.device.type == "cuda"):
+                out = self.backbone(batch)
+
+            patch_embeds = out["patch_embeds"]
+            token_mask = out["token_mask"].bool().unsqueeze(-1)
+            embed = (patch_embeds * token_mask).sum(dim=1) / token_mask.sum(dim=1)
+            pooled.append(embed[0].float().cpu())
+
+        return torch.cat(pooled).numpy()
+
+    def cached_features(self, row: dict) -> np.ndarray:
+        if row["subject"] not in self.cache:
+            self.cache[row["subject"]] = self.features(row)
+        return self.cache[row["subject"]]
+
+    def fit(self, rows: list[dict]) -> None:
+        X = np.stack([self.cached_features(row) for row in rows])
+        y = np.array([row["label"] for row in rows])
+
+        clf = LogisticRegressionCV(
+            Cs=10,
+            class_weight="balanced",
+            scoring="roc_auc",
+            max_iter=1000,
+            l1_ratios=(0,),
+            use_legacy_attributes=False,
+        )
+        self.head = make_pipeline(StandardScaler(), clf)
+        self.head.fit(X, y)
+        self.positive = list(self.head.classes_).index(1)
+
+    def predict(self, images: Images) -> float:
+        """Positive-class probability. Indexes `classes_` rather than assuming column 1, which
+        would silently score the wrong class if the label order differed."""
+        X = self.features(images)[None]
+        probs = self.head.predict_proba(X)[0]
+        return float(probs[self.positive])
+
+    def save(self, model_dir: Path) -> None:
+        """Everything `load` needs but the backbone weights, which stay wherever `ckpt_path`
+        points -- a few hundred KB, so a run saves one without copying a 3.7G checkpoint."""
+        model_dir.mkdir(parents=True, exist_ok=True)
+        OmegaConf.save(self.cfg, model_dir / "config.yaml")
+        joblib.dump({"head": self.head, "positive": self.positive}, model_dir / "head.joblib")
+
+    @classmethod
+    def load(cls, model_dir: Path, **overrides) -> "Task1Method":
+        """Rebuild a fitted method from `save`. Overrides are Config fields, for what differs
+        between here and the container -- the backbone path, the device."""
+        cfg = OmegaConf.merge(
+            OmegaConf.structured(Config), OmegaConf.load(model_dir / "config.yaml"), overrides
+        )
+        method = cls(cfg)
+        state = joblib.load(model_dir / "head.joblib")
+        method.head, method.positive = state["head"], state["positive"]
+        return method
+
+
+# ---- protocol: the part we hold fixed ---------------------------------------------------
+
+# Every image the task ships. The method picks which of them it wants, as at inference, where
+# the challenge hands over all four modalities whether or not a model uses them.
+IMAGE_COLS = ("adc", "dwi_b1000", "flair")
+
+
+def leave_one_out(rows: list[dict], method: Task1Method) -> tuple[np.ndarray, np.ndarray]:
+    """Out-of-fold score for every subject, each predicted by a head fit on the other n-1."""
+    y = np.array([row["label"] for row in rows])
+    oof = np.zeros(len(rows), dtype=float)
+    start = time.perf_counter()
+    for held_out, row in enumerate(rows):
+        method.fit([r for r in rows if r["subject"] != row["subject"]])
+        oof[held_out] = method.predict({key: row[key] for key in IMAGE_COLS})
+        logger.info(
+            f"fold {held_out + 1}/{len(rows)} {row['subject']} "
+            f"y={y[held_out]} p={oof[held_out]:.3f} ({time.perf_counter() - start:.0f}s)"
+        )
+    return y, oof
+
+
+def score(
+    y: np.ndarray, oof: np.ndarray, seed: int = 0, n_boot: int = 2000, alpha: float = 0.05
+) -> dict:
+    """AUROC, the challenge metric, plus a percentile CI resampling subjects with replacement."""
+    rng = np.random.default_rng(seed)
+    samples = []
+    for _ in range(n_boot):
+        rows = rng.integers(0, len(y), size=len(y))
+        if len(np.unique(y[rows])) < 2:
+            continue
+        samples.append(roc_auc_score(y[rows], oof[rows]))
+
+    low, high = np.percentile(samples, [100 * alpha / 2, 100 * (1 - alpha / 2)])
+    return {
+        "auroc": float(roc_auc_score(y, oof)),
+        "auroc_ci_low": float(low),
+        "auroc_ci_high": float(high),
+    }
+
+
+# ---- entrypoints ------------------------------------------------------------------------
+
+
+def train(args: argparse.Namespace) -> None:
+    # imported here, not at the top, so the container needs no dataset stack to run `predict`
+    from fomo_tune.datasets import load_fomo_task1
+
+    cfg = OmegaConf.merge(OmegaConf.structured(Config), OmegaConf.from_dotlist(args.overrides))
+    run_dir = Path(cfg.output_root) / cfg.name
+    run_dir.mkdir(parents=True, exist_ok=True)
+
+    setup_logging(run_dir)
+    set_seed(cfg.seed)
+    logger.info(f"run {cfg.name} (git {git_sha()})")
+    logger.info(f"config:\n{OmegaConf.to_yaml(cfg).rstrip()}")
+    OmegaConf.save(cfg, run_dir / "config.yaml")
+
+    # decoded once: leave-one-out revisits every subject n times, and the niftis are small
+    rows = list(load_fomo_task1())
+    logger.info(f"dataset: {len(rows)} subjects, {sum(r['label'] for r in rows)} positive")
+
+    method = Task1Method(cfg)
+    start = time.perf_counter()
+    y, oof = leave_one_out(rows, method)
+    run_time = time.perf_counter() - start
+    summary = score(y, oof)
+
+    # the shipped head sees all n subjects, so it is not any of the models scored above
+    method.fit(rows)
+    method.save(run_dir / "model")
+
+    record = {"name": cfg.name, **summary, "run_time": round(run_time, 1)}
+    (run_dir / "metrics.json").write_text(json.dumps(record) + "\n")
+    scores = "  ".join(f"{k}={v:.4f}" for k, v in summary.items())
+    logger.info(f"result: {scores}  ({run_time:.0f}s)")
+
+
+def predict(args: argparse.Namespace) -> None:
+    """The challenge contract: modality paths in, one probability written to `--output`.
+
+    `/app/predict.py` in the container is a shim over this, so what scores the submission is the
+    code leave-one-out already ran, not something generated at build time.
+    """
+    overrides = {"device": args.device}
+    if args.ckpt_path:
+        overrides["ckpt_path"] = args.ckpt_path
+    method = Task1Method.load(args.model_dir, **overrides)
+
+    # every image the challenge hands over, as in `leave_one_out`; the method takes what it uses
+    paths = {"adc": args.adc, "dwi_b1000": args.dwi, "flair": args.flair}
+    probability = method.predict({key: nib.load(path) for key, path in paths.items()})
+
+    args.output.write_text(f"{probability:.6f}\n")
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    modes = parser.add_subparsers(required=True)
+
+    train_parser = modes.add_parser("train", help="leave-one-out over the task, then fit and save")
+    train_parser.add_argument("overrides", nargs="*", help="config overrides, e.g. device=cpu")
+    train_parser.set_defaults(run=train)
+
+    predict_parser = modes.add_parser("predict", help="one subject, one probability")
+    for flag in ("--flair", "--adc", "--dwi"):
+        predict_parser.add_argument(flag, type=Path, required=True)
+    # accepted and ignored: the 4th modality is swi on some subjects and t2s on others
+    for flag in ("--t2s", "--swi"):
+        predict_parser.add_argument(flag, type=Path)
+    predict_parser.add_argument("--output", type=Path, required=True)
+    predict_parser.add_argument("--model-dir", type=Path, default=Path("/app/model"))
+    predict_parser.add_argument("--ckpt-path", help="overrides the trained config's backbone path")
+    predict_parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
+    predict_parser.set_defaults(run=predict)
+
+    args = parser.parse_args()
+    args.run(args)
+
+
+if __name__ == "__main__":
+    main()

finetune/fomo_tune_baseline/output/task1/build/fomo_tune/main_task3.py

@@ -0,0 +1,241 @@
+"""FOMO task 3: brain age regression, scored by Pearson r and MAE as the challenge scores it.
+
+`Task3Method` is the part we tune -- features, head, hyperparameters. The protocol below it is
+fixed so scores stay comparable across iterations: 20-fold over the 494 subjects, pool the
+out-of-fold predictions, bootstrap subjects for the CI.
+
+`train` runs that protocol then fits and saves a head; `predict` is the challenge contract, one t1
+path in and one age out. Both go through `Task3Method.predict`, so every fold exercises the path
+the submission will run.
+"""
+
+import argparse
+import json
+import logging
+import time
+from dataclasses import dataclass
+from pathlib import Path
+
+import joblib
+import nibabel as nib
+import numpy as np
+import torch
+from omegaconf import OmegaConf
+from sklearn.linear_model import RidgeCV
+from sklearn.model_selection import KFold
+from sklearn.pipeline import make_pipeline
+from sklearn.preprocessing import StandardScaler
+
+from fomo_tune.backbone import load_backbone
+from fomo_tune.utils import git_sha, set_seed, setup_logging
+
+logger = logging.getLogger("fomo_tune")
+
+Images = dict[str, nib.Nifti1Image]
+
+
+@dataclass
+class Config:
+    task: str = "task3"
+    ckpt_path: str = (
+        "/data/mihir-stuff/smri-pretrained/pretrain_full_90_10_h100/checkpoint-last.pth"
+    )
+    output_root: str = "output/fomo_tune"
+    name: str = "task3"
+    device: str = "cuda"
+    seed: int = 4466
+
+
+# ---- method: the part we tune -----------------------------------------------------------
+
+
+class Task3Method:
+    """Frozen sMRI MAE, mean-pooled tokens over the t1w, ridge head."""
+
+    def __init__(self, cfg: Config):
+        self.cfg = cfg
+        self.backbone, self.transform = load_backbone(cfg.ckpt_path)
+        self.device = torch.device(cfg.device)
+        self.backbone.to(self.device).eval().requires_grad_(False)
+        self.cache: dict[str, np.ndarray] = {}
+        self.head = None
+
+    @torch.inference_mode()
+    def features(self, images: Images) -> np.ndarray:
+        """(D,) per subject. A pure function of the images, so training and inference agree."""
+        sample = self.transform(images["t1w"])
+        batch = {key: value[None].to(self.device) for key, value in sample.items()}
+
+        with torch.autocast("cuda", torch.bfloat16, enabled=self.device.type == "cuda"):
+            out = self.backbone(batch)
+
+        patch_embeds = out["patch_embeds"]
+        token_mask = out["token_mask"].bool().unsqueeze(-1)
+        embed = (patch_embeds * token_mask).sum(dim=1) / token_mask.sum(dim=1)
+        return embed[0].float().cpu().numpy()
+
+    def cached_features(self, row: dict) -> np.ndarray:
+        if row["subject"] not in self.cache:
+            self.cache[row["subject"]] = self.features(row)
+        return self.cache[row["subject"]]
+
+    def fit(self, rows: list[dict]) -> None:
+        X = np.stack([self.cached_features(row) for row in rows])
+        y = np.array([row["age"] for row in rows], dtype=float)
+
+        # RidgeCV picks alpha by its own efficient leave-one-out, so the fold's own split is
+        # never touched by model selection
+        self.head = make_pipeline(StandardScaler(), RidgeCV(alphas=np.logspace(-3, 6, 19)))
+        self.head.fit(X, y)
+
+    def predict(self, images: Images) -> float:
+        """Age in years."""
+        X = self.features(images)[None]
+        return float(self.head.predict(X)[0])
+
+    def save(self, model_dir: Path) -> None:
+        """Everything `load` needs but the backbone weights, which stay wherever `ckpt_path`
+        points -- a few hundred KB, so a run saves one without copying a 3.7G checkpoint."""
+        model_dir.mkdir(parents=True, exist_ok=True)
+        OmegaConf.save(self.cfg, model_dir / "config.yaml")
+        joblib.dump(self.head, model_dir / "head.joblib")
+
+    @classmethod
+    def load(cls, model_dir: Path, **overrides) -> "Task3Method":
+        """Rebuild a fitted method from `save`. Overrides are Config fields, for what differs
+        between here and the container -- the backbone path, the device."""
+        cfg = OmegaConf.merge(
+            OmegaConf.structured(Config), OmegaConf.load(model_dir / "config.yaml"), overrides
+        )
+        method = cls(cfg)
+        method.head = joblib.load(model_dir / "head.joblib")
+        return method
+
+
+# ---- protocol: the part we hold fixed ---------------------------------------------------
+
+# Every image the task ships. The method picks which of them it wants, as at inference, where the
+# challenge hands over the modalities whether or not a model uses them.
+IMAGE_COLS = ("t1w",)
+
+
+def cross_validate(
+    rows: list[dict], method: Task3Method, seed: int = 0, n_folds: int = 20
+) -> tuple[np.ndarray, np.ndarray]:
+    """Out-of-fold age for every subject, each predicted by a head fit on the other folds."""
+    y = np.array([row["age"] for row in rows], dtype=float)
+    oof = np.zeros(len(rows), dtype=float)
+    folds = KFold(n_splits=n_folds, shuffle=True, random_state=seed)
+    start = time.perf_counter()
+    for fold, (train, test) in enumerate(folds.split(rows)):
+        method.fit([rows[i] for i in train])
+        for i in test:
+            oof[i] = method.predict({key: rows[i][key] for key in IMAGE_COLS})
+        logger.info(
+            f"fold {fold + 1}/{n_folds} n={len(test)} mae={np.abs(y[test] - oof[test]).mean():.2f} "
+            f"({time.perf_counter() - start:.0f}s)"
+        )
+    return y, oof
+
+
+def metrics(y: np.ndarray, oof: np.ndarray) -> dict:
+    return {
+        "pearson_r": float(np.corrcoef(y, oof)[0, 1]),
+        "mae": float(np.abs(y - oof).mean()),
+    }
+
+
+def score(
+    y: np.ndarray, oof: np.ndarray, seed: int = 0, n_boot: int = 2000, alpha: float = 0.05
+) -> dict:
+    """Both challenge metrics, each with a percentile CI resampling subjects with replacement."""
+    rng = np.random.default_rng(seed)
+    resamples = rng.integers(0, len(y), size=(n_boot, len(y)))
+
+    summary = {}
+    for name, point in metrics(y, oof).items():
+        samples = [metrics(y[rows], oof[rows])[name] for rows in resamples]
+        low, high = np.percentile(samples, [100 * alpha / 2, 100 * (1 - alpha / 2)])
+        summary[name] = point
+        summary[f"{name}_ci_low"] = float(low)
+        summary[f"{name}_ci_high"] = float(high)
+    return summary
+
+
+# ---- entrypoints ------------------------------------------------------------------------
+
+
+def train(args: argparse.Namespace) -> None:
+    # imported here, not at the top, so the container needs no dataset stack to run `predict`
+    from fomo_tune.datasets import load_fomo_task3
+
+    cfg = OmegaConf.merge(OmegaConf.structured(Config), OmegaConf.from_dotlist(args.overrides))
+    run_dir = Path(cfg.output_root) / cfg.name
+    run_dir.mkdir(parents=True, exist_ok=True)
+
+    setup_logging(run_dir)
+    set_seed(cfg.seed)
+    logger.info(f"run {cfg.name} (git {git_sha()})")
+    logger.info(f"config:\n{OmegaConf.to_yaml(cfg).rstrip()}")
+    OmegaConf.save(cfg, run_dir / "config.yaml")
+
+    rows = list(load_fomo_task3())
+    ages = np.array([row["age"] for row in rows])
+    logger.info(
+        f"dataset: {len(rows)} subjects, age {ages.min()}-{ages.max()} mean {ages.mean():.1f}"
+    )
+
+    method = Task3Method(cfg)
+    start = time.perf_counter()
+    y, oof = cross_validate(rows, method)
+    run_time = time.perf_counter() - start
+    summary = score(y, oof)
+
+    # the shipped head sees all n subjects, so it is not any of the models scored above
+    method.fit(rows)
+    method.save(run_dir / "model")
+
+    record = {"name": cfg.name, **summary, "run_time": round(run_time, 1)}
+    (run_dir / "metrics.json").write_text(json.dumps(record) + "\n")
+    scores = "  ".join(f"{k}={v:.4f}" for k, v in summary.items())
+    logger.info(f"result: {scores}  ({run_time:.0f}s)")
+
+
+def predict(args: argparse.Namespace) -> None:
+    """The challenge contract: a t1 path in, one age written to `--output`.
+
+    `/app/predict.py` in the container is a shim over this, so what scores the submission is the
+    code cross-validation already ran, not something generated at build time.
+    """
+    overrides = {"device": args.device}
+    if args.ckpt_path:
+        overrides["ckpt_path"] = args.ckpt_path
+    method = Task3Method.load(args.model_dir, **overrides)
+
+    age = method.predict({"t1w": nib.load(args.t1)})
+
+    args.output.write_text(f"{age:.6f}\n")
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    modes = parser.add_subparsers(required=True)
+
+    train_parser = modes.add_parser("train", help="cross-validate over the task, then fit and save")
+    train_parser.add_argument("overrides", nargs="*", help="config overrides, e.g. device=cpu")
+    train_parser.set_defaults(run=train)
+
+    predict_parser = modes.add_parser("predict", help="one subject, one age in years")
+    predict_parser.add_argument("--t1", type=Path, required=True)
+    predict_parser.add_argument("--output", type=Path, required=True)
+    predict_parser.add_argument("--model-dir", type=Path, default=Path("/app/model"))
+    predict_parser.add_argument("--ckpt-path", help="overrides the trained config's backbone path")
+    predict_parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
+    predict_parser.set_defaults(run=predict)
+
+    args = parser.parse_args()
+    args.run(args)
+
+
+if __name__ == "__main__":
+    main()