CT organ segmentation using GPU data augmentation, unsupervised labels and IOU loss (1811.11226v1)

Abstract: Fully-convolutional neural networks have achieved superior performance in a variety of image segmentation tasks. However, their training requires laborious manual annotation of large datasets, as well as acceleration by parallel processors with high-bandwidth memory, such as GPUs. We show that simple models can achieve competitive accuracy for organ segmentation on CT images when trained with extensive data augmentation, which leverages existing graphics hardware to quickly apply geometric and photometric transformations to 3D image data. On 3 mm^3 CT volumes, our GPU implementation is 2.6-8X faster than a widely-used CPU version, including communication overhead. We also show how to automatically generate training labels using rudimentary morphological operations, which are efficiently computed by 3D Fourier transforms. We combined fully-automatic labels for the lungs and bone with semi-automatic ones for the liver, kidneys and bladder, to create a dataset of 130 labeled CT scans. To achieve the best results from data augmentation, our model uses the intersection-over-union (IOU) loss function, a close relative of the Dice loss. We discuss its mathematical properties and explain why it outperforms the usual weighted cross-entropy loss for unbalanced segmentation tasks. We conclude that there is no unique IOU loss function, as the naive one belongs to a broad family of functions with the same essential properties. When combining data augmentation with the IOU loss, our model achieves a Dice score of 78-92% for each organ. The trained model, code and dataset will be made publicly available, to further medical imaging research.

• The paper demonstrates that GPU-based augmentation with fully-convolutional networks accelerates training by 2.6-8 times compared to traditional CPU methods.
• The paper shows that using IOU loss instead of weighted cross-entropy effectively tackles class imbalance, achieving Dice scores between 78% and 92%.
• The paper introduces unsupervised label generation via morphological operations and 3D Fourier transforms to reduce manual annotation efforts in CT segmentation.

Overview of CT Organ Segmentation Using GPU Data Augmentation, Unsupervised Labels, and IOU Loss

This paper presents a methodology for improving CT organ segmentation through a combination of GPU-based data augmentation, automatic label generation, and an innovative loss function choice. The authors tackle the prevalent challenges in medical image segmentation where data annotation is cumbersome and computational costs are high due to the volumetric nature of the data. By leveraging a fully-convolutional neural network (FCN), the paper demonstrates that simpler models, when appropriately trained and augmented, can achieve competitive segmentation accuracy.

The key innovations in the paper include employing extensive data augmentation using GPUs, which accelerates the training process significantly compared to traditional CPU methods. The paper reports a throughput that is 2.6-8 times faster with GPU implementation, highlighting the computational advantage in managing large 3D image data sets. Additionally, the introduction of the intersection-over-union (IOU) loss function addresses the problem of class imbalance more effectively than the usual weighted cross-entropy loss. The authors propose a mathematically solid rationale for the use of IOU loss, underscoring its symmetry in penalizing false positives and false negatives, in contrast to cross-entropy.

Furthermore, the paper describes the generation of training labels through unsupervised techniques, utilizing morphological operations enhanced by 3D Fourier transforms. This facilitates automatic segmentation for organs with obvious density contrasts, such as lungs and bone, while reducing the manual labor typically required for these tasks. The dataset created includes 130 labeled CT scans and supports multi-class segmentation tasks encompassing organs like the liver, kidneys, and bladder.

The implications of this research extend into practical applications where such models, once trained, could be deployed in clinical settings requiring robust and quick segmentation outcomes. The methodology ensures that a wide variety of anatomical structures can be segmented accurately and with minimal manual intervention. The authors open-source their model, code, and dataset for further research, encouraging community contributions to expand organ labeling.

Experimentally, the model trained with IOU loss and GPU-accelerated augmentations achieved Dice scores ranging from 78-92% across different organs, comparable with more complex model architectures. The performance validated on the Liver Tumor Segmentation Challenge (LiTS) data aligns closely with established literature benchmarks, showing promise in reducing the complexity without sacrificing accuracy.

In conclusion, this research provides a methodologically sound and computationally efficient approach to organ segmentation in CT imaging, driven by GPU-enhanced techniques and a judiciously chosen loss function. The findings suggest potential for further enhancement and adaptation of this methodology to other complex segmentation challenges in medical imaging by building upon the foundational innovations presented. Future work in AI segmentation could benefit from exploring additional loss function families within the IOU framework, optimizing GPU frameworks further, and expanding the application to additional types of imaging data and conditions.