Managing Class Imbalance in Multi-Organ CT Segmentation in Head and Neck Cancer Patients

Radiotherapy planning of head and neck cancer patients requires an accurate delineation of several organs at risk (OAR) from planning CT images in order to determine a dose plan which reduces toxicity and salvages normal tissue. However training a single deep neural network for multiple organs is highly sensitive to class imbalance and variability in size between several structures within the head and neck region. In this paper, we propose a single-class segmentation model for each OAR in order to handle class imbalance issues during training across output classes (one class per structure), where there exists a severe disparity between 12 OAR. Based on a U-net architecture, we present a transfer learning approach between similar OAR to leverage common learned features, as well as a simple weight averaging strategy to initialize a model as the average of multiple models, each trained on a separate organ. Experiments performed on an internal dataset of 200 H & N cancer patients treated with external beam radiotherapy, show the proposed model presents a significant improvement compared to the baseline multi-organ segmentation model, which attempts to simultaneously train several OAR. The proposed model yields an overall Dice score of $0.75 \pm 0.12$, by using both transfer learning across OAR and a weight averaging strategy, indicating that a reasonable segmentation performance can be achieved by leveraging additional data from surrounding structures, limiting the uncertainty in ground-truth annotations.

[1]  J.D. Furst,et al.  Single Organ Segmentation Filters for Multiple Organ Segmentation , 2006, 2006 International Conference of the IEEE Engineering in Medicine and Biology Society.

[2]  Jimmy Ba,et al.  Adam: A Method for Stochastic Optimization , 2014, ICLR.

[3]  Thomas Brox,et al.  U-Net: Convolutional Networks for Biomedical Image Segmentation , 2015, MICCAI.

[4]  Max Dahele,et al.  Deep Learning-Based Delineation of Head and Neck Organs at Risk: Geometric and Dosimetric Evaluation. , 2019, International journal of radiation oncology, biology, physics.

[5]  Christian S. Perone,et al.  Spinal cord gray matter segmentation using deep dilated convolutions , 2017, Scientific Reports.

[6]  Yabo Fu,et al.  Deep Learning in Multi-organ Segmentation , 2020, ArXiv.

[7]  Sébastien Ourselin,et al.  Generalised Dice overlap as a deep learning loss function for highly unbalanced segmentations , 2017, DLMIA/ML-CDS@MICCAI.

[8]  Zhijian Song,et al.  Organ at Risk Segmentation in Head and Neck CT Images Using a Two-Stage Segmentation Framework Based on 3D U-Net , 2018, IEEE Access.

[9]  Klaus H. Maier-Hein,et al.  Brain Tumor Segmentation and Radiomics Survival Prediction: Contribution to the BRATS 2017 Challenge , 2017, BrainLes@MICCAI.

[10]  Joseph Paul Cohen,et al.  Deep semantic segmentation of natural and medical images: a review , 2019, Artificial Intelligence Review.

[11]  Kai Ma,et al.  Med3D: Transfer Learning for 3D Medical Image Analysis , 2019, ArXiv.

[12]  Jieun Kim,et al.  Abdominal multi-organ auto-segmentation using 3D-patch-based deep convolutional neural network , 2020, Scientific Reports.

[13]  Yuanjie Zheng,et al.  Multi-organ Segmentation via Co-training Weight-averaged Models from Few-organ Datasets , 2020, MICCAI.

[14]  Andrew Gordon Wilson,et al.  Averaging Weights Leads to Wider Optima and Better Generalization , 2018, UAI.