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
while read requirement; do
    pip install "$requirement" || echo "Failed to install $requirement, skipping..."
done < 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.

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.