This repository contains our proposed image segmentation frameworks applicable to both 2D and 3D segmentation. These include architectures and losses of:
-
HartleyMHA
Ken C. L. Wong, Hongzhi Wang, and Tanveer Syeda-Mahmood, “HartleyMHA: self-attention in frequency domain for resolution-robust and parameter-efficient 3D image segmentation,” in International Conference on Medical Image Computing and Computer-Assisted Intervention (MICCAI), 2023, pp. 364–373. [pdf]
@inproceedings{Conference:Wong:MICCAI2023:hartleymha, title = {{HartleyMHA}: self-attention in frequency domain for resolution-robust and parameter-efficient {3D} image segmentation}, author = {Wong, Ken C. L. and Wang, Hongzhi and Syeda-Mahmood, Tanveer}, booktitle = {International Conference on Medical Image Computing and Computer-Assisted Intervention (MICCAI)}, pages = {364--373}, year = {2023}, } -
FNOSeg3D
Ken C. L. Wong, Hongzhi Wang, and Tanveer Syeda-Mahmood, “FNOSeg3D: resolution-robust 3D image segmentation with Fourier neural operator,” in IEEE International Symposium on Biomedical Imaging (ISBI), 2023, pp. 1–5. [pdf]
@inproceedings{Conference:Wong:ISBI2023:fnoseg3d, title = {{FNOSeg3D}: resolution-robust {3D} image segmentation with {Fourier} neural operator}, author = {Wong, Ken C. L. and Wang, Hongzhi and Syeda-Mahmood, Tanveer}, booktitle = {IEEE International Symposium on Biomedical Imaging (ISBI)}, pages = {1--5}, year = {2023}, } -
V-Net-DS (V-Net with deep supervision)
Ken C. L. Wong, Mehdi Moradi, Hui Tang, and Tanveer Syeda-Mahmood, “3D segmentation with exponential logarithmic loss for highly unbalanced object sizes,” in International Conference on Medical Image Computing and Computer-Assisted Intervention (MICCAI), 2018, pp. 612–619. [pdf]
@inproceedings{Conference:Wong:MICCAI2018:3d, title = {{3D} segmentation with exponential logarithmic loss for highly unbalanced object sizes}, author = {Wong, Ken C. L. and Moradi, Mehdi and Tang, Hui and Syeda-Mahmood, Tanveer}, booktitle = {International Conference on Medical Image Computing and Computer-Assisted Intervention (MICCAI)}, pages = {612--619}, year = {2018}, } -
Pearson’s Correlation Coefficient (PCC) loss
Ken C. L. Wong and Mehdi Moradi, “3D segmentation with fully trainable Gabor kernels and Pearson’s correlation coefficient,” in Machine Learning in Medical Imaging, 2022, pp. 53–61. [pdf]
@inproceedings{Workshop:Wong:MLMI2022:3d, title = {{3D} segmentation with fully trainable {Gabor} kernels and {Pearson's} correlation coefficient}, author = {Wong, Ken C. L. and Moradi, Mehdi}, booktitle = {Machine Learning in Medical Imaging}, pages = {53--61}, year = {2022}, }
The code is developed with Python 3.10.12 and Keras in TensorFlow 2.16.1, and the "channel-last" format is assumed. If you are only interested in the architectures, the nets module is all you need, though the hyperparameters are stored under experiments/config_files as they are dataset specific. The experimental setups such as data splits and training procedure are in the experiments module. The nets module is dataset independent, while some functions in the experiments module (e.g., dataset partitioning) are written exclusively for BraTS'19.
In experiments, parameters or arguments are provided through a config file using the Python's module ConfigParser. The config file is saved to the output directory for future reference. Examples of the config files used in our experiments are provided under experiments/config_files for reproducibility.
For your convenience, we include an example of our experimental setup in BraTS2019_example.zip which contains the necessary folders and config files for testing. To perform experiments on the BraTS'19 dataset, please obtain the dataset from CBICA. In our example, we assume that the images are downsampled by the nearest neighbor interpolation to 120x120x78 for training, and the original data hierarchy remains unchanged. For inference, the official validation dataset with the original image size of 240x240x155 should be used.
There are multiple Python packages required to run the code. You can install them by the following steps:
-
Create a virtual environment (https://docs.python.org/3/library/venv.html). Note that the default
pythonin your system may be Python 2. You can use the following command to ensure Python 3 is used:python3 -m venv /path/to/new/virtual/environment -
Upgrade
pipin the activated virtual environment:pip install --upgrade pipIt is important to upgrade
pipas the installed version can be outdated and the next step may fail. -
Install the required Python packages using:
pip install tensorflow[and-cuda] natsort SimpleITK matplotlib pandas pydotNote: The Linux (not Python) library
graphvizis required by the functionkeras.utils.plot_model. If you encounter the corresponding runtime error, you can either installgraphvizbysudo apt-get install graphvizif you have sudo privileges, or setis_plot_model = Falsein the training config file to skipplot_model.
For more information on troubleshooting, see Troubleshooting.
The experiments/data_split folder contains the script and config files for partitioning the BraTS'19 dataset. The program goes through the dataset folders to extract the patient IDs and groups them into training, validation, and testing sets. The resulted lists of file paths are saved as txt files. To run the script, we first modify the config_partitioning.ini config file, then use the command line:
python partitioning.py /path/to/config_partitioning.ini
The split examples used in our experiments are provided under split_examples. They are also included in BraTS2019_example.zip.
To perform training, we first modify the config_<arch>.ini file, then run:
python run.py /path/to/config_<arch>.ini
where <arch> stands for an architecture (e.g., fnoseg). The config files of different architectures are only different in the [model] section and output_dir. Note that in our example, the validation and testing sets are the same. The config files used in our experiments can be found under experiments/config_files. They are also included in BraTS2019_example.zip.
To perform inference, we first modify the config_inference_<arch>.ini file, then run:
python inference.py /path/to/config_inference_<arch>.ini
where <arch> stands for an architecture (e.g., fnoseg). The segmentation results can be uploaded to CBICA for the official performance validation. The config files used in our experiments can be found under experiments/config_files. They are also included in BraTS2019_example.zip.
We find that in the official validation results, the "enhancing tumor" (ET) region has sensitivity of NaN. We also find that Hausdorff95_ET = NaN when Sensitivity_ET = 1. These indicate that there may be no positives for ET for some cases. Therefore, when computing the means and variances for the ET region (e.g., Dice_ET, Hausdorff95_ET), those cases with Sensitivity_ET equals to NaN or 1 are ignored.
- Updated codes for the most recent version of TensorFlow (2.16.1).
- The
datagenerator.pymodule is replaced by thedataset.pymodule that usesPyDatasetin Keras 3. AsPyDatasetis new in Keras 3 and thus TensorFlow 2.16.1, classInputDatais not backward compatible.
Ken C. L. Wong (clwong@us.ibm.com)