An open source, PyTorch-based segmentation framework for 3D medical image
3 min read
EfficientSegmentation
Introduction
EfficientSegmentation is an open source, PyTorch-based segmentation framework for 3D medical image.
Features
A whole-volume-based coarse-to-fine segmentation framework. The segmentation network is decomposed into different components, including encoder, decoder and context module. Anisotropic convolution block and anisotropic context block are designed for efficient and effective segmentation.
Pre-process data in multi-process. Distributed and Apex training support. Postprocess is performed asynchronously in inference stage.
git clone https://github.com/NVIDIA/apex
cd apex
pip install -v --no-cache-dir ./
3.Install dependencies
pip install -r requirements.txt
Get Started
preprocessing
Download FLARE21, resulting in 361 training images and masks, 50 validation images.
Copy image and mask to ‘FlareSeg/dataset/’ folder.
Edit the ‘FlareSeg/data_prepare/config.yaml’. ‘DATA_BASE_DIR'(Default: FlareSeg/dataset/) is the base dir of databases. If set the ‘IS_SPLIT_5FOLD'(Default: False) to true, 5-fold cross-validation datasets will be generated.
Run the data preprocess with the following command:
python FlareSeg/data_prepare/run.py
The image data and lmdb file are stored in the following structure:
Remark: Coarse segmentation is trained on Nvidia GeForce 2080Ti(Number:8) in the experiment, while fine segmentation on Nvidia A100(Number:4). If you use different hardware, please set the “ENVIRONMENT.NUM_GPU”, “DATA_LOADER.NUM_WORKER” and “DATA_LOADER.BATCH_SIZE” in ‘FlareSeg/coarse_base_seg/config.yaml’ and ‘FlareSeg/fine_efficient_seg/config.yaml’ files.
Coarse segmentation:
Edit the ‘FlareSeg/coarse_base_seg/config.yaml’
Train coarse segmentation with the following command:
cd FlareSeg/coarse_base_seg
sh run.sh
Fine segmentation:
Edit the ‘FlareSeg/fine_efficient_seg/config.yaml’.
Edit the ‘FlareSeg/fine_efficient_seg/run.py’, set the ‘tune_params’ for different experiments.
Train fine segmentation with the following command:
cd FlareSeg/fine_efficient_seg
sh run.sh
Inference:
The model weights are stored in ‘FlareSeg/model_weights/’.
John was the first writer to have joined pythonawesome.com. He has since then inculcated very effective writing and reviewing culture at pythonawesome which rivals have found impossible to imitate.
Previous Post
LibFewShot: A Comprehensive Library for Few-shot Learning
Next Post
Web service using Flask framework for Shortener URLs