Lung Nodule Segmentation

A research project exploring the effects of data augmention on the LIDC dataset.

Our research paper follows IEEE-CPVR style. Our poster is in standard 32x24 format and was presented to our professors and TAs. We received full marks on this project for UW CSE’s graduate-level deep learning course.

Research Highlights

Motivation

Lung cancer has a mere 18% survival rate in the United States, largely due to the lack of preventive screening. The current method for screening lung cancer relies on radiologists — who are already in short supply — to manually segment lung nodules on CT scans. Computer-aided diagnosis (CAD) systems can help alleviate the shortage of screening capacity, but are still lacking in accuracy.

“the global implementation of lung cancer screening is of utmost importance” ~ IASLC (the International Association for the Study of Lung Cancer).

Given the extremely small size of biomedical datasets (LIDC only includes 1018 CT scans), data augmentations intuitively seem like a potent tool. Existing research confirms this intuition, but notes that each biomedical task has different optimal augmentations. Crucially, there is no existing work in finding the best augmentations for lung nodule segmentation. Our paper explores this open question and provides recommendations for future scientists working with the LIDC dataset.

Pipeline

We introduce the first open-source “plug-and-play” pipeline for the LIDC dataset, written entirely in PyTorch. This contribution enables the research community to easily try out different approaches to the LIDC dataset (whether it be expanding the neural network architecure, utilizing different augmentations, training with a new loss function, or really anything supported by PyTorch) without having to deal with the confusing world that is medical imaging.

In essence, our pipeline abstracts away the medical part of problem and lets researchers treat it as a normal deep learning task.

Under the hood, our workflow follows this training procedure:

  1. Download the 125 GB LIDC-IDRI dataset.
  2. Utilize pylidc and Pidicom to extract 10,005 2D slices of 512 x 512 pixel arrays and convert them to PNG images.
  3. Store each nodule’s bounding box cooordinates in MS-COCO format.
  4. Crop each nodule image and mask with a 64 x 64 pixel bounding box.
  5. Extract regions of interest (ROIs) using either Faster R-CNN or YOLO (our pipeline supports both).
  6. Augment images using any of PyTorch’s native augmentations. We also added support for RandAugment, AutoAugment, and TrivialAugment.
  7. Utilize our custom-designed U-Net or another PyTorch-compatible architecture for the LIDC task.

Models

The LIDC task can be broken down into two sub-tasks: a “detector” model which extracts ROIs where nodules may lie, and a “segmentation” model which does pixel-level classification on the ROIs. We focused on the segmentation task because general-purpose detection models transfer well to the LIDC task while general-purpose segmentation models underperform in medical contexts. The table below compares our models to existing work.

a comparison of our models to existing literature

Our best-performing U-Net model outperformed all published models except Lung_PAYNet. Our architecture is described below.

architecture of our u-net model

Augmentations

examples of augmentations

We tried a variety of augmentations on our models, pictured above. Their peformance is described by the table below. AutoAug (1), RandAug, and TrivialAug were pre-trained on ImageNet. AutoAug (2) was custom trained by us on the LIDC dataset using the Albumentations library. A more detailed discussion can be found in our paper in section 2.6 Augmentation Techniques.

comparison of augmentation techniques effectiveness on LIDC

Conclusion

Contrary to the suggestions of existing literature, augmentations provide only minor benefits for lung nodule segmentation. Future models should simply use horizontal and vertical flipping for augmentations and instead focus their efforts on improving the underlying network architecture.

We also believe that the development of a comprehensive automatic augmentation suite tailored to medical images would be of great utility, as current ImageNet-trained augmentation libraries struggle to adapt to grayscale medical imaging. Similarly, efforts to expand the LIDC dataset by using GANs should also be explored further.

Contributors

Acknowledgements

This work would not have been possible without the guidance of the CSE 493G1 teaching team and the Allen Center prost and senna undergraduate GPUs.