Bubenpo/AnguinusSculpturae

MAMA-SYNTH — 2D pre-to-post contrast synthesis (final configuration)

Trained weights for our MAMA-SYNTH challenge entry: given a pre-contrast breast-MRI slice, generate the matching post-contrast slice.

This repository holds weights only. The code lives in the nnUNet-Mama-Synth repository (a fork of nnU-Net v2), which provides the prediction entry point used below. These are the checkpoints bundled in the submitted container mama-synth-foreground-stitched-lesion-folds-synthfolds-v1.0.0, copied verbatim except that optimizer and grad-scaler state have been stripped (inference-irrelevant, halves the download).

Everything is 2D. All networks are slice-wise, the only configuration is 2d, the plans identifier is mamaSynthPlans, and slices sit on a fixed 512×512 canvas whose zero padding is excluded from every loss and metric.

What the pipeline does, per slice

Two translation models cover different parts of the image, and two segmentations decide where each applies:

image_synth  = mean over folds of the outside-breast translation model
lesion_synth = mean over folds of the inside-breast translation model
stitched     = breast_soft     * lesion_synth + (1 - breast_soft)     * image_synth
final        = foreground_soft * stitched     + (1 - foreground_soft) * pre
final        = final * (1 + (gain - 1) * lesion_soft)

breast_soft and foreground_soft are signed-distance ramps across the mask boundaries, so the transitions are gradual rather than hard steps. Outside the tissue support the real pre-contrast image is kept unchanged. The lesion region comes from the fold-averaged lesion segmentation, with the threshold lowered per slice until at least one in-breast voxel passes; scaling the intensity that is already there preserves the lesion's internal texture.

Composite settings for this configuration: feather 4, lesion_feather 2, lesion_gain 1.25, lesion_threshold 0.5, lesion_step 0.05, air_from pre, flip TTA on (mirror axes (0, 1)).

Contents

Slot	Role	Trainer	Source checkpoint	Folds
image	outside-breast translation	nnUNetTrainerMamaSynthTranslationLPIPS_BS_48_MAEFinetune_ep_200	checkpoint_best	0–3
lesion_image	inside-breast translation	nnUNetTrainerMamaSynthTranslationLPIPSSSIMDiceTverskyA02B08_LPIPS05_BS_48_MAEFinetune	checkpoint_best_dice	0–3
lesion_seg	lesion segmentation	nnUNetTrainerMamaSynthLesionTverskyA02B08_BS_64_epoch_1000	checkpoint_best	0–3
breast	breast segmentation	nnUNetTrainerMamaSynthBreast_BS_32	checkpoint_final	0
foreground	tissue-support segmentation	nnUNetTrainerMamaSynthForeground_BS_32	checkpoint_final	0

14 networks in total, ~106 M parameters each (ResidualEncoderUNet, 7 stages, features 32→512, 512×512 patch), ~424 MB per checkpoint, 5.6 GB total. The translation models are initialised from a masked-autoencoder pretraining run; the inside-breast model's objective includes a term computed through the frozen lesion segmenter. image, lesion_image and lesion_seg were trained on Dataset625_Pre_Seg, foreground on Dataset627_foreground.

Layout

Per-fold checkpoints are renamed to a uniform checkpoint.pth; plans.json and dataset.json sit at each slot root and are required, since every network is rebuilt from them.

config.json                 # pipeline summary: slots, trainers, composite settings
assets/
├── image/        {plans.json, dataset.json, fold_0..3/checkpoint.pth}
├── lesion_image/ {plans.json, dataset.json, fold_0..3/checkpoint.pth}
├── lesion_seg/   {plans.json, dataset.json, fold_0..3/checkpoint.pth}
├── breast/       {plans.json, dataset.json, checkpoint.pth}
└── foreground/   {plans.json, dataset.json, checkpoint.pth}

The root config.json describes the configuration in machine-readable form; it is not consumed by the inference code, which reads each slot's plans.json instead. It is also the Hub's default download-counting query file, so its presence is what makes this repository's download statistics register at all.

Usage

Download the weights:

hf download Bubenpo/AnguinusSculpturae --local-dir mama-synth-weights

Then run the pipeline on a folder of .mha / .tif slices, from a checkout of nnUNet-Mama-Synth installed with pip install -e .:

W=mama-synth-weights/assets

nnUNetv2_predict_mamasynth_translate_foreground_stitched_lesion_folds_synthfolds \
  -i <input_dir> -o <output_dir> \
  --image-dir        $W/image        --image-folds        0 1 2 3 --image-checkpoint-name        checkpoint.pth \
  --lesion-image-dir $W/lesion_image --lesion-image-folds 0 1 2 3 --lesion-image-checkpoint-name checkpoint.pth \
  --lesion-seg-dir   $W/lesion_seg   --lesion-seg-folds   0 1 2 3 --lesion-seg-checkpoint-name   checkpoint.pth \
  --breast-weights     $W/breast/checkpoint.pth \
  --plans-breast       $W/breast/plans.json     --dataset-json-breast     $W/breast/dataset.json \
  --foreground-weights $W/foreground/checkpoint.pth \
  --plans-foreground   $W/foreground/plans.json --dataset-json-foreground $W/foreground/dataset.json \
  --feather 4 --lesion-feather 2 --lesion-gain 1.25 --lesion-threshold 0.5 \
  --dont-norm -device cuda

Two details are easy to get wrong:

• --plans-breast / --plans-foreground are required with this layout. For the fold-ensembled slots the run directory is inferred as the checkpoint's grandparent, so image/fold_0/checkpoint.pth finds image/plans.json on its own. The two single-checkpoint slots are flat, so their plans and dataset files have to be passed explicitly.
• --dont-norm assumes already z-scored inputs, as the challenge validation inputs are; the output then stays in that same space. For raw inputs, drop it for per-slice normalisation or pass --pre-stats for dataset-wide statistics. The segmenters always z-score each slice internally, independent of this flag.

--save-masks, --save-intermediates and --save-lesion-mask write the masks, the soft weights and the two intermediate syntheses next to each output — the quickest way to see where a result went wrong.

Runtime is dominated by the 12 ensembled networks × 4-flip TTA per slice; a single 512×512 slice takes a few seconds on a 24 GB GPU. Checkpoints were saved with pickle_protocol=2 and load under torch 2.3.1 (the submission container's version) and newer.

Limitations

Research artifact from a challenge entry — not a medical device, and not for clinical use. The models were trained on the challenge's breast-MRI data and expect single 2D pre-contrast slices on the 512×512 canvas described above; behaviour on other anatomy, other field strengths, 3D volumes, or non-z-scored inputs is untested. The lesion-gain step assumes each input slice contains a lesion (it lowers its threshold until one voxel passes), so on lesion-free slices it will brighten whatever the segmenter ranks highest.

Citation

Please cite nnU-Net when using this code:

Isensee, F., Jaeger, P. F., Kohl, S. A., Petersen, J., & Maier-Hein, K. H. (2021).
nnU-Net: a self-configuring method for deep learning-based biomedical image segmentation.
Nature Methods, 18(2), 203-211.

nnU-Net is developed by the Applied Computer Vision Lab of Helmholtz Imaging and the Division of Medical Image Computing at the German Cancer Research Center (DKFZ).