Skip to content

Add anatomy segmentation grounding step (MONAI / TotalSegmentator total_mr) #53

Description

@Liohtml

Part of #51.

Summary

Add an optional pipeline step that runs anatomical segmentation on the
preprocessed MRI volume and feeds the resulting structures/measurements into the
prompt as grounding context for the Vision-LLM, and into the validator (#52)
as an independent signal. Grounding the LLM in "here are the structures actually
present and their measurements" measurably reduces hallucination and lets the
report cite quantitative findings (volumes, presence/absence) instead of only
free-text impressions.

Why (vision)

The Vision-LLM currently sees only slice images. If we additionally tell it "a
segmentation model localized these structures, with these volumes,"
its findings
become anchored to measurable facts. The segmentation masks also become a strong
independent signal for the cross-validation step (#52): a finding about a
structure the segmenter couldn't even locate is suspect.

Recommended components (license-verified, Apache-2.0-clean)

  • MONAI (Apache-2.0) — primary framework; bundles in the MONAI model-zoo
    (check each bundle's license individually).
  • TotalSegmentatorcode Apache-2.0; use the total_mr task, whose
    weights are Apache-2.0 / commercially usable.
    ⚠️ Do NOT use the non-commercial tasks (e.g. brain_aneurysm is CC BY-NC,
    and appendicular_bones/tissue_types/heartchambers_highres require a
    commercial license) — restrict to Apache-2.0 weight tasks only.
  • TorchIO (MIT) — IO/resampling/augmentation helper.
  • Avoid: SynthSeg (TensorFlow stack mismatch) and the FreeSurfer surface path
    (separate license) unless brain-only and clearly gated.

Where

  • New optional extra in pyproject.toml, e.g. segmentation = ["monai>=1.4", "torch>=2.5", ...]
    (keep it out of the base install; torch is already only in local-models).
  • New file: src/medcheck/pipeline/segmentation.pySegmentationStep(PipelineStep)
  • Register as "segmentation"; runs after preprocess, before
    ml_analysis / vision_analysis.
  • Extend PipelineContext with segmentation: dict[str, Any] (per-structure masks
    summary, volumes, present/absent) — store summaries, not giant arrays, where possible.
  • vision_analysis.build_prompt() gains an optional grounding block listing
    detected structures + measurements when context.segmentation is populated.

Design constraints

  • Lazy import of monai/torch inside the step (mirror the existing
    LocalLLMProvider / provider pattern); if the extra isn't installed,
    check/validate returns False and the step is skipped with a clear log — the
    core pipeline still runs.
  • Model weights are downloaded on first use (like medcheck download-models,
    which is currently a stub) — wire this into that command.
  • CPU-capable path documented; GPU optional.

Acceptance criteria

  • SegmentationStep registered as "segmentation", behind the optional extra.
  • Skips gracefully (no crash) when deps/weights are unavailable.
  • Only Apache-2.0-licensed weights/tasks are used by default; NC tasks are
    not selectable without an explicit opt-in flag + warning.
  • Populates context.segmentation; build_prompt() includes grounding when present.
  • download-models can fetch the segmentation weights.
  • Tests with a mocked segmenter (no real model download in CI).
  • Docs: new extra, licensing note (which tasks are commercial-safe).

Effort

~1–2 days (incl. mocked tests + docs). Heavy optional deps.

Metadata

Metadata

Assignees

No one assigned

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions