Introduction

ShapeKit is a plug-and-play post-processing toolkit that enables researchers and clinicians to correct anatomical errors in AI-predicted segmentations without retraining models. It integrates seamlessly into existing pipelines and supports robust, anatomy-aware refinement across multiple organs and datasets.

Using a parallelized Python workflow, ShapeKit combines, calibrates, and refines multi-organ segmentations, leading to up to 15% improvement in Dice Similarity Coefficient (DSC) and producing consistent outputs suitable for downstream analysis.

Paper

ShapeKit
Junqi Liu*, Dongli He*, Wenxuan Li, Ningyu Wang, Alan Yuille, Zongwei Zhou
Johns Hopkins University
*Equal contribution.
MICCAI 2025 Workshop on Shape in Medical Imaging

Installation

To set up environment, see INSTALL.md for details.

git clone https://github.com/BodyMaps/ShapeKit.git
cd ShapeKit
pip install -r requirements.txt

Use ShapeKit

Organize your data

INPUT or OUTPUT
└── case_001
    ├── combined_labels.nii.gz (optional)
    └── segmentations
            ├── liver.nii.gz
            ...
            └── veins.nii.gz

export INPUT="/path/to/your/input/folder"
export OUTPUT="/path/to/your/output/folder"
export CPU_NUM=16
export LOG="logs/folder_named_after_your_task"

python -W ignore main.py --input_folder $INPUT --output_folder $OUTPUT --cpu_count $CPU_NUM --log_folder $LOG --continue_prediction

The processing process will be recorded as debug.log and postprocessing.log,and are stored under the directory LOG.

Plug-and-Play Configuration

Tell ShapeKit which anatomical structures you are interested in by modifying the config.yaml file.

Check for details 🔍

How to choose your interested anatomical structures:

Open the config.yamlfile and list the anatomical structures you want to process under target_organs. It’s as easy as checking boxes on a form.

# plug-and-play like Lego! choose organs for processing

target_organs: (example)
  - bladder
  - colon
  - duodenum
  - femur
  - intestine
  - kidney
  - liver
  - lung
  - pancreas
  - vertebrae

For detailed configuration setting, please check the config instructions 🌞..

Before running any commands, please ensure that config.yaml is properly configured. But don't worry! Most of the configurations do not need to be changed at all.

Evidence-Gated Vertebrae Engine (ShapeKit-Pro)

The default vertebrae module works from the masks alone. ShapeKit can now optionally repair vertebrae against the case CT image with an evidence-gated engine that recolors label errors inside the prediction envelope instead of deleting bone:

• fragments are re-attached through CT-certified bone corridors, never discarded when they are real bone;
• level-band mass misassignment (e.g. a collapsed L1 split between neighbors) is re-arbitrated at image-detected disc planes;
• the posterior arch is rebuilt from pedicle roots, and one-level-down spinous chains on fused spines are repaired by a caudal-flow re-derivation;
• every risky stage carries its own defect meter and reverts itself when it cannot prove improvement, so one parameter set is safe across clean and pathological cases at scale.

Enable it in config.yaml:

vertebrae_engine: shapekit_pro   # default: shapekit (existing module)
ct_file_name: ct.nii.gz          # looked up inside each input case folder
# ct_root: /path/to/ct/cases     # fallback root when CTs live elsewhere

No new dependencies (numpy, scipy, nibabel, scikit-image and connected-components-3d are already required). CPU only; ~2 min for a 2.5 mm case and ~30 min for a 0.7 mm whole-spine case on 2 cores, with a peak of roughly 9 GB on the latter — budget --cpu_count accordingly. When a case has no reachable CT the engine logs it and falls back to the default vertebrae module, so batch runs never stall.

Measured on the AbdomenAtlasDemo cases (identical parameters, per-stage QA and verification tooling in the ShapeKit-Pro repository): both cases reach zero structural audit flags (fragmentation, ordering, size, emptiness; exactly 24 components), the collapsed L1 is restored from 23.3 to 62.9 cm3 at detected disc planes, and every spinous process is re-attached to its own vertebra on the fused case.

Key Functions

In addition to these general utilities, anatomical-structures-specific correction functions are available in organs_postprocessing.py.

Please check the details in functions guide book 📖.

Related Articles

@article{liu2025shapekit,
  title={ShapeKit},
  author={Liu, Junqi and He, Dongli and Li, Wenxuan and Wang, Ningyu and Yuille, Alan L and Zhou, Zongwei},
  journal={arXiv preprint arXiv:2506.24003},
  year={2025}
}

Acknowledgement

This work was supported by the Lustgarten Foundation for Pancreatic Cancer Research, the Patrick J. McGovern Foundation Award, and the National Institutes of Health (NIH) under Award Number R01EB037669. We would like to thank the Johns Hopkins Research IT team in IT@JH for their support and infrastructure resources where some of these analyses were conducted; especially DISCOVERY HPC. Paper content is covered by patents pending.