VoxWhisper

Two-stage brain-MRI specialist. Phase 1 (this branch) pretrains VoxDense: a T1-only language-conditioned UNet on FreeSurfer dense structures. Phase 2 (later) freezes that T1 encoder, adds an FA encoder, and segments unseen nerves from spatial descriptions.

Do not train FA on FreeSurfer labels.

---

Quick start

# 1. Install
pip install -e .

# 2. Preprocess HCP T1 + wmparc onto the 1.25 mm grid
python scripts/preprocess.py --config config/voxdense.yaml

# 3. Train VoxDense 
python scripts/train.py --config config/voxdense.yaml

# 4. Evaluate
python scripts/evaluate.py --split test

Edit config/voxdense.yaml between runs. Structure names live in config/structures_dense.json.

---

Phase 1 model

T1 volume → encoder → PromptDecoder (frozen PubMedBERT names) → Decoder

The checkpoint stores encoder_state_dict so Phase 2 can reload the T1 encoder into dual VoxWhisper. The dual class stays in the tree but is not trained here.

Pipeline notes: docs/voxdense_training.md. Spatial-only nerve stage (design spec, not trained here): docs/spatial_nerve_segmentation.md.

Figures (docs/flowcharts/):

Figure	Content
VoxDense_architecture.drawio	Model (encoder / PromptDecoder / fusion)
VoxDense_full_forward.drawio	Tensor shapes through the forward
VoxDense_training_pipeline.drawio	Preprocess → batch → loss → checkpoint
VoxDense_inference.drawio	Sliding-window eval
VoxDense_metrics_debug.drawio	Which Dice number means what

---

Tunable knobs (config/voxdense.yaml)

Key	Description
training.learning_rate	AdamW peak LR after warm-up
training.warmup_epochs	Linear LR warm-up duration
training.epochs	Total training epochs (cosine T_max)
training.batch_size	Physical batch size
training.bce_weight / dice_weight	Loss term multipliers
training.exclude_background	Skip channel 0 in Dice and BCE (default true)
training.ignore_empty_targets	Skip empty-target FG channels (default true)
training.deep_supervision_weights	Per-scale weights [coarse, mid, fine]
training.checkpoint.monitor	dice (patch, V3) or later volume
training.checkpoint.keep	Top-k checkpoints
training.early_stopping.patience / min_delta	Stop if dice_patch gain ≤ min_delta
data.patch.prompts_per_crop	0 = all names; k>0 = k FG after background
data.patch.size / train_patches_per_subject	Crop geometry
logging.wandb.*	W&B project and tags

BCE positive weighting is always adaptive (n_neg/n_pos). Resolve eval checkpoints with --checkpoint / --run-dir.

---

Preprocessing

HCP FreeSurfer outputs already on disk under data/raw/{subject_id}/. Work in 1.25 mm space. Subjects with nerve_masks_1.25 are held out for Phase 2.

python scripts/preprocess.py

Writes data/processed_dense/{subject}/t1.nii.gz + mask.nii.gz and cache/prompts_dense.pt.

---

Project layout

VoxWhisper/
├── config/
│   ├── voxdense.yaml          ← Phase 1 baseline
│   ├── voxdense_smoke.yaml    ← 4-structure overfit gate
│   └── structures_dense.json
├── scripts/                   ← preprocess, train, evaluate
├── voxwhisper/
│   ├── models/vox_dense.py
│   ├── data/dataset.py        ← BasePatchDataset + phase subclasses
│   └── training/
├── docs/
│   ├── voxdense_training.md
│   └── flowcharts/
└── tests/

---

Citation

@misc{voxwhisper2026,
  title        = {VoxWhisper: Language-Grounded 3D Brain MRI Segmentation},
  author       = {Huppe, Maxime},
  year         = {2026},
}