nnQC - Segmentation Quality Control via Latent Diffusion

Project page: https://robustml-eurecom.github.io/nnQC/

nnQC is a quality-control model for medical image segmentation. It is a 2D latent diffusion model (LDM) that, given a CT/MR scan and a corrupted segmentation mask, reconstructs what it believes the correct mask should look like. The Dice score between an input mask and the model's reconstruction is the QC signal - calibrated to track the true Dice against the ground truth (see tutorials/TUTORIAL.md).

        scan + corrupted mask ──► nnQC LDM ──► predicted "correct" mask (pgt)
                                                    │
                              Dice(pgt, corrupted) ─┴─► QC score

The pipeline is built on top of MONAI's AutoencoderKL and DiffusionModelUNet, with CLIP-based scan conditioning (UniMedCLIP) and a slice-ratio embedding so the same 2D model handles apex / mid / base slices.

---

News

• 2026-08-24 - MCP server released: expose nnQC to MCP-capable agents (python -m nnqc.mcp_server). See MCP server below and the walkthrough demo.
• 2026-07-10 - Paper accepted at IEEE Transactions on Medical Imaging. 🎉🎊🥳
• 2025-12-29 - Paper submitted and under review at IEEE TMI.

Roadmap

• MONAI tutorial covering the full train-then-QC workflow.

---

Installation

nnqc is a Python package and is best used inside a fresh virtualenv. We recommend uv for fast resolution:

git clone https://github.com/robustml-eurecom/nnQC.git
cd nnQC
uv venv
source .venv/bin/activate
uv pip install -e .

(Plain pip install -e . also works.) Installing registers the nnqc command and ships the bundled task presets, so --task prostate works out of the box.

The package depends on PyTorch with CUDA. If you need a specific CUDA version, install torch first from the official wheel index, then uv pip install -e . will pick up the rest. To run a one-off command without activating the venv, prefix it with uv run (e.g. uv run nnqc list-tasks).

Pretrained weights

Trained checkpoints are not distributed in the git repo because they exceed 1 GB per task. They are published in two places:

• Zenodo as a single archive, nnQC_pretrained_weights.zip, with one weights_<task>/ folder per task (used by the built-in helper below);
• Hugging Face as a model repo with the same weights_<task>/ folders, for direct per-file download with huggingface_hub or wget.

Fetch them with the built-in helper:

import nnqc
nnqc.download_weights("prostate")     # -> trained_weights/prostate/

nnqc download prostate                 # same, from the CLI

The archive is downloaded once and cached under trained_weights/.cache/, so fetching a second task does not download it again. check and evaluate also auto-download a task's weights on first use if they are missing. The files per task (autoencoder.pt, diffusion_unet.pt, xa.pt, embed.pt, scale_factor.txt) are extracted under trained_weights/<task>/, where xa.pt already bundles the UniMedCLIP backbone.

From Hugging Face instead:

from huggingface_hub import snapshot_download
snapshot_download("sanbast/nnQC", allow_patterns="weights_prostate/*",
                  local_dir="trained_weights/hf")

The Zenodo record id is set in ZENODO_RECORD in nnqc/hub.py. To use a different record (e.g. a fork's weights), pass the id explicitly:

nnqc download prostate --record 1234567
# or
export NNQC_ZENODO_RECORD=1234567

Maintainer note: publishing weights on Zenodo

1. Create a single deposition on Zenodo.

2. Upload nnQC_pretrained_weights.zip, built as:

   cd trained_weights
   zip -r nnQC_pretrained_weights.zip nnQC_pretrained_weights

   where nnQC_pretrained_weights/ holds one weights_<task>/ folder per task (weights_acdc for cardiac), each containing autoencoder.pt, diffusion_unet.pt, xa.pt, embed.pt and scale_factor.txt.

3. Publish the deposition and copy its numeric record id.

4. Either add the id to ZENODO_RECORD in nnqc/hub.py, or export NNQC_ZENODO_RECORD=<id> when downloading.

5. Mirror the same weights_<task>/ folders to the Hugging Face model repo (huggingface_hub.upload_folder) so users can fetch single files without downloading the whole zip.

---

Pipeline

nnQC is trained in two stages:

1. Autoencoder - encodes/decodes one-hot segmentation masks into a low-dimensional latent.
2. Diffusion UNet - denoises the mask latent, conditioned on (a) the corrupted mask resized to latent resolution, (b) CLIP image features of the scan, and (c) a slice-ratio embedding.

Both stages read a config pair, either a bundled preset (--task) or an explicit JSON pair:

• env.json - paths and dataset settings (model dir, task name, modality, num_classes, image/label glob patterns, resume options).
• config.json - network architecture, training hyper-parameters, noise scheduler settings.

Any field can be overridden from the CLI or the Python API without editing the JSON. Bundled presets, all with pretrained weights on Zenodo (see above):

Task	Anatomy	num_classes	Modality	Preset
MSD Task03 + holdout	Liver	1	CT	liver
MSD Task05 + Prostate158	Prostate (peripheral + transition zone)	3	MRI T2	prostate
ACDC	Cardiac (LV, RV, myocardium)	4	MRI	cardiac
MSD Task09	Spleen	1	CT	spleen

Run nnqc list-tasks to see what is installed. The same JSON files also live under configs/ for you to copy and edit.

---

Quickstart (CLI)

A single nnqc command with subcommands:

# 1. autoencoder
nnqc train-autoencoder --task prostate --epochs 500 --lr 5e-5 --device 0

# 2. diffusion UNet (needs <model_dir>/autoencoder.pt)
nnqc train-diffusion --task prostate --epochs 4000 --lr 2.5e-5 \
    --scheduler cosine --warmup-dice-epochs 100 --device 0

# 3. visualize reconstructions
nnqc evaluate --task prostate --num-volumes 3 --num-steps 5 --device 0

Use an explicit config pair instead of a preset with --config configs/prostate/config.json --env configs/prostate/env.json. Multi-GPU uses torchrun: torchrun --nproc_per_node=2 -m nnqc.cli train-diffusion --task prostate -g 2.

EMA-smoothed UNet weights are written to diffusion_unet.pt (best by val loss) and diffusion_unet_last.pt (latest); cross-attention and slice-embedding weights follow the same convention. 5 DDIM steps give the best quality / latency trade-off; raise --num-steps to 20-50 for a finer schedule.

Quickstart (notebook / Python)

import nnqc

nnqc.train_autoencoder(task="prostate", epochs=500, lr=5e-5, device=0)
nnqc.train_diffusion(
    task="prostate", epochs=4000, lr=2.5e-5,
    scheduler="cosine", warmup_dice_epochs=100, device=0,
)
nnqc.evaluate(task="prostate", num_volumes=3, num_steps=5, device=0)

# resume a run, or point at your own data:
nnqc.train_diffusion(task="prostate", resume=True, start_epoch=1000, epochs=4000)
nnqc.train_diffusion(
    config="configs/prostate/config.json",
    env="configs/prostate/env.json",
    data_dir="/data/my_prostate", model_dir="/data/runs/prostate",
)

Common overrides (CLI flag / Python kwarg): epochs, lr, batch_size, patch_size, val_interval, scheduler (cosine|constant|step|exponential), warmup_dice_epochs, lambda_recon, ema_decay, num_train_timesteps, model_dir, output_dir, data_dir, resume, start_epoch.

---

Quality control on a new mask

Once a model is trained, check() is the one-call QC entry point. Give it a scan and a candidate segmentation; it preprocesses both (orientation, slice and foreground crop, resize, intensity scaling), reconstructs the mask the model believes is correct, and returns the Dice agreement as the QC score. The reconstruction is mapped back onto the input volume's grid (shape + affine), so you can save or overlay it directly.

import nnqc

result = nnqc.check("scan.nii.gz", "candidate_mask.nii.gz", task="prostate")
print(result.qc_score)            # volume Dice(candidate, reconstruction); low = suspect mask
print(result.qc_score_per_class)  # per-class Dice (multi-class models)
print(result.slice_scores)        # per-slice Dice, with result.slice_ratios
result.save("reconstruction.nii.gz")   # written on the input grid

nnqc check --task prostate --image scan.nii.gz --mask candidate_mask.nii.gz \
    --save reconstruction.nii.gz

With the default Dice metric, a high qc_score means the candidate agrees with what the model reconstructs (likely good); a low score flags a probable segmentation error.

Custom QC metric

The agreement metric is pluggable via metric=. Pass a built-in name, a bare callable fn(pred, ref) -> float (e.g. from medpy), or a Metric subclass. QCResult reports metric_name and higher_is_better so you know how to read the score.

import nnqc

# built-in by name
nnqc.check(img, mask, task="prostate", metric="iou")

# any third-party function (install medpy: pip install nnqc[metrics])
from medpy.metric.binary import hd95
nnqc.check(img, mask, task="prostate",
           metric=nnqc.FunctionMetric(hd95, name="hd95", higher_is_better=False))
# (or metric="hd95", which already knows lower-is-better)

# a custom metric class
from nnqc.metrics import Metric

class MyMetric(Metric):
    name = "my_metric"
    higher_is_better = True
    def score(self, pred, ref):       # pred, ref: binary numpy arrays
        return float((pred & ref).sum() / max(pred.sum(), 1))

nnqc.check(img, mask, task="prostate", metric=MyMetric())

Empty-mask edge cases are handled by the base class (empty_value / empty_both_value), so per-slice scoring never crashes on blank apex/base slices. From the CLI use --metric dice|iou|hd95|assd.

---

MCP server

nnqc/mcp_server.py exposes QC to MCP-capable agents (Claude Desktop, Codex, ...) over stdio:

python -m nnqc.mcp_server

Three tools are registered:

• list_tasks - which trained organ models are available, with modality and class names.
• check_mask - score a scan + candidate mask pair; returns the QC score, a per-class breakdown, and the worst-scoring slices.
• explain_qc_score - how to read a score, including the measured reliability limits.

A minimal client configuration (Claude Desktop style):

{
  "mcpServers": {
    "nnqc": {
      "command": "python",
      "args": ["-m", "nnqc.mcp_server"]
    }
  }
}

Two caveats: the server needs a GPU (nnqc/xa.py is CUDA-only), and stdout is the JSON-RPC transport, so nnqc's own progress prints are redirected to stderr. Set NNQC_MCP_STRICT_STDOUT=1 to make any non-JSON write to stdout raise instead of silently corrupting the protocol. A full walkthrough with example exchanges is in docs/mcp_demo.md.

---

Repository layout

nnQC/
├── nnqc/                       Importable package
│   ├── __init__.py             Public API: train_autoencoder/train_diffusion/evaluate
│   ├── cli.py                  `nnqc` command dispatcher
│   ├── config.py               JSON + kwargs config resolver, task presets
│   ├── train.py                Training loops (autoencoder + diffusion)
│   ├── evaluate.py             DDIM sampling + reconstruction panels
│   ├── infer.py                check(): one-call QC on a scan + mask pair
│   ├── hub.py                  download_weights(): fetch checkpoints from Zenodo
│   ├── metrics.py              Pluggable QC metrics (Dice, IoU, medpy adapters)
│   ├── xa.py                   CLIPCrossAttentionGrid (UniMedCLIP wrapper)
│   ├── corruptions.py          Morphologically realistic mask corruptions
│   ├── utils.py                Dataloaders, transforms, helpers
│   ├── visualize.py            TensorBoard image helpers
│   └── presets/                Bundled task configs (shipped in the wheel)
│       ├── prostate/{config,env}.json
│       └── spleen/{config,env}.json
├── configs/                    Editable copies of the presets
│   ├── prostate/{config,env}.json
│   └── spleen/{config,env}.json
├── tutorials/
│   └── TUTORIAL.md             End-to-end walkthrough
├── pyproject.toml
├── LICENSE                     (MIT)
└── README.md

---

Citation

If you use nnQC, please cite:

@ARTICLE{11614032,
  author={Marcianò, Vincenzo and Chaptoukaev, Hava and Fernandez, Virginia and Cardoso, M. Jorge and Antonelli, Michela and Ourselin, Sébastien and Zuluaga, Maria A.},
  journal={IEEE Transactions on Medical Imaging}, 
  title={Diffusion-Based Quality Control of Medical Image Segmentations across Organs}, 
  year={2026},
  volume={},
  number={},
  pages={1-1},
  keywords={Modeling;Biomedical imaging;Training;Biological systems;Learning (artificial intelligence);Labeling;Image segmentation;Magnetic resonance imaging;Layered division multiplexing;Conferences;Quality Control;Generative Modeling;Self-adapting Framework;Medical Image Segmentation},
  doi={10.1109/TMI.2026.3714697}}

See CITATION.cff for a machine-readable version.

---

License

MIT - see LICENSE.