Lightweight Deep Learning for Biomedical Image Segmentation
- Author(s): Uys, Dylan
- Advisor(s): Cottrell, Garrison W
- et al.
Many techniques for analyzing cardiovascular health rely on cardiac magnetic resonance images that have been segmented to identify various components of the heart. Manually segmenting these images is cumbersome and prone to variability, which calls for the development of accessible automation tools for cardiac researchers. In order to benefit the developing symbiosis between machine learning and medicine, such tools must be accurate, efficient and inferentially transparent. This paper introduces a U-Net-based pipeline for left ventricular segmentation of short-axis CMRs. The U-Net, a Fully Convolutional Network known for its success in biomedical image segmentation, is a natural candidate for our task. This work constitutes the core of a larger-scale project focused on improving human disease models through the acceleration of animal cardiac research. Accordingly, experiments discussed here leverage both human and animal data to explore the efficacy of image processing and model training strategies.
This paper focuses on optimizing our U-Nets for resource constrained environments, and demonstrates that these models require only a fraction of their typical convolutional filters. This reduction affords efficiency with the added benefit of explainability by improving the practicality of visualizing learned features. Inference can also be further optimized by pruning trained models without any loss of accuracy or the need for retraining. Specifically, we show that U-Nets with less than 2% of their original parameters train in minutes on a single GPU and achieve Dice scores above 0.95 on multiple CMR datasets. Furthermore, inference on hundreds of images can be performed in seconds on a laptop.