Alleviating Class-wise Gradient Imbalance for Pulmonary Airway Segmentation
This work provides a solution for medical image analysis, specifically improving automated pulmonary airway segmentation for clinicians, which is crucial for pre-operative diagnosis and intra-operative navigation.
The paper addresses the challenge of pulmonary airway segmentation, particularly for small peripheral bronchi, which suffer from severe class imbalance and gradient erosion. The authors propose a method using group supervision with WingsNet and a General Union loss function to enhance gradient flow and adaptively tune gradient ratios, resulting in higher accuracy and better morphological completeness compared to baselines.
Automated airway segmentation is a prerequisite for pre-operative diagnosis and intra-operative navigation for pulmonary intervention. Due to the small size and scattered spatial distribution of peripheral bronchi, this is hampered by severe class imbalance between foreground and background regions, which makes it challenging for CNN-based methods to parse distal small airways. In this paper, we demonstrate that this problem is arisen by gradient erosion and dilation of the neighborhood voxels. During back-propagation, if the ratio of the foreground gradient to background gradient is small while the class imbalance is local, the foreground gradients can be eroded by their neighborhoods. This process cumulatively increases the noise information included in the gradient flow from top layers to the bottom ones, limiting the learning of small structures in CNNs. To alleviate this problem, we use group supervision and the corresponding WingsNet to provide complementary gradient flows to enhance the training of shallow layers. To further address the intra-class imbalance between large and small airways, we design a General Union loss function which obviates the impact of airway size by distance-based weights and adaptively tunes the gradient ratio based on the learning process. Extensive experiments on public datasets demonstrate that the proposed method can predict the airway structures with higher accuracy and better morphological completeness than the baselines.