Source-linked AI summary
Residual Attention U-Net for Automated Multi-Class Segmentation of COVID-19 Chest CT Images
Xiaocong Chen, Lina Yao, Yu Zhang
TL;DR
COVID-19 infection quantification from CT images lacks sufficient effective segmentation methods, despite segmentation's importance for diagnosis. The paper proposes a residual-attention U-Net and reports more than 10% improvement over U-Net and other baselines in multi-class segmentation. It presents the model as a tool for quantifying lung infection and supporting COVID-19 diagnosis.
Problem
Effective quantification of COVID-19 lung infection remains insufficiently studied, although CT segmentation is important for accurate infection measurement.
Method
The paper modifies U-Net with residual-network feature extraction and a soft attention mechanism in the decoder for multi-class COVID-19 CT segmentation.
Results
More than 10% improvement in multi-class segmentation was achieved over U-Net and a set of baselines.
Takeaways & Limitations
The study provides a tool for radiologists to determine lung infection percentage and diagnose COVID-19 progression.
Takeaways & Limitations
Future work identifies limited ground-truth data for new diseases and the need to improve model interpretability in medical applications.
Abstract
from arXiv · showhide
The novel coronavirus disease 2019 (COVID-19) has been spreading rapidly around the world and caused significant impact on the public health and economy. However, there is still lack of studies on effectively quantifying the lung infection caused by COVID-19. As a basic but challenging task of the diagnostic framework, segmentation plays a crucial role in accurate quantification of COVID-19 infection measured by computed tomography (CT) images. To this end, we proposed a novel deep learning algorithm for automated segmentation of multiple COVID-19 infection regions. Specifically, we use the Aggregated Residual Transformations to learn a robust and expressive feature representation and apply the soft attention mechanism to improve the capability of the model to distinguish a variety of symptoms of the COVID-19. With a public CT image dataset, we validate the efficacy of the proposed algorithm in comparison with other competing methods. Experimental results demonstrate the outstanding performance of our algorithm for automated segmentation of COVID-19 Chest CT images. Our study provides a promising deep leaning-based segmentation tool to lay a foundation to quantitative diagnosis of COVID-19 lung infection in CT images.
I. INTRODUCTION
COVID-19 CT imaging offers a basis for evaluating infection, but accurately segmenting infection regions remains difficult and underexplored. The paper therefore proposes an enhanced U-Net for automated multi-class segmentation and quantitative assessment.
- Accurate infection-region segmentation is crucial for reliably quantifying COVID-19 infection in chest CT images.
- Manual annotation by expert radiologists burdens labeling efforts, motivating automated segmentation of infection regions.
- Variation in infection-region position and shape, together with low CT contrast, makes automated delineation challenging.
- COVID-19 infection-region segmentation with deep learning remains underexplored despite successful U-Net applications to other medical imaging tasks.
- The proposed model combines residual networks and attention with U-Net and is evaluated on a public CT dataset from patients with COVID-19.
II. METHODOLOGY
The methodology builds on U-Net's encoder-decoder design, skip connections, and convolutional processing for segmentation-map construction. Its described baseline structure uses downsampling in the encoder and upsampling with a final one-channel convolutional output in the decoder.
- U-Net uses a symmetric encoder-decoder architecture, with the encoder extracting spatial features and the decoder constructing the segmentation map.
- The encoder performs repeated convolution and max-pooling operations, doubling the number of filters after each downsampling step.
- The decoder upsamples feature maps through transposed convolution, convolutional blocks, and a final 1 × 1 convolutional layer.
- The final decoder layer uses Sigmoid activation, whereas the other layers use ReLU activation.
- Skip connections concatenate encoder features before pooling with decoder features, helping recover information lost during max-pooling.
B. Aggregated Residual Block
The proposed encoder replaces U-Net’s conventional feature extractor with aggregated residual transformations, while attention-enhanced decoding preserves relevant features for multi-class segmentation.
- Motivation: U-Net’s shallow convolutional encoder may struggle to distinguish small COVID-19 infection regions, motivating a deeper feature-extraction network.Increasing depth can introduce degradation, where accuracy rises and then rapidly decreases.
- Residual learning: ResNet mitigates degradation by adding an identity shortcut to the output of stacked layers, supporting residual learning.The shortcut avoids estimating a large number of parameters generated by convolutional layers.
- Aggregated residual transformations: ResNeXt aggregates C residual transformations T_j(i), with cardinality representing the number of transformations combined.The component transformations can project the input into a lower-dimensional embedding before transforming it.
- Aggregated residual transformations: The ResNeXt block uses multiple topologically equivalent small blocks and a smaller weight layer than ResNet while maintaining performance.The paper illustrates cardinality = 3 in the ResNeXt block.
- Proposed encoder-decoder: The proposed model uses ResNeXt blocks instead of the U-Net encoder to prevent excessive depth while retaining performance and reducing training cost relative to ResNet.Decoder blocks combine encoder features with deconvolution outputs followed by LSH attention, then propagate filtered features through skip connections.
C. Locality Sensitive Hashing Attention
The model applies soft attention in the decoder to emphasize relative pixel positions, using locality-sensitive hashing to reduce the memory demands of standard attention.
- Soft attention: Soft attention lets the model focus on each pixel’s relative position, unlike hard attention, which focuses only on absolute position.The paper considers scaled dot-product and multi-head attention as soft-attention variants.
- Scaled dot-product attention: Scaled dot-product attention computes attention from queries Q, keys K, and values V, with softmax applied after scaling by the key dimension.The scaling factor is √n, and the softmax input is an n-dimensional vector.
- Multi-head attention: Multi-head attention linearly projects queries, keys, and values into separate learnable dimensions, calculates attention in parallel, and concatenates the resulting heads.The output is multiplied by the learnable matrix W^O.
- LSH attention: Standard multi-head attention is memory-inefficient because QK^T produces a tensor whose length × length dimensions dominate memory use for images.The paper introduces LSH attention to address this issue.
- LSH attention: LSH attention assigns queries and keys to hash buckets so each query attends to a restricted set P_i, while multi-round hashing addresses mismatched similar items.The paper describes parallel calculation of the resulting LSH attention scores.
D. Training Strategy
Training treats multi-class lung CT segmentation as pixel-wise classification and uses multi-class cross entropy to compare predicted and correct classes.
- Task formulation: The task predicts whether each image pixel belongs to a predefined class or the background, making it a multi-class pixel-wise classification problem.This extends the binary formulation used in traditional medical image segmentation.
- Loss function: The model uses multi-class cross entropy as its loss function for training the segmentation network.The loss compares the correct class with the observed class using class probabilities across M classes.
A. Data Description
The experiments use COVID-19 chest CT images from the Italian Society of Medical and Interventional Radiology dataset, comprising 110 axial images from 60 patients.
- Dataset: The study uses 110 axial CT images collected from 60 patients in the SIRM COVID-19 CT image dataset.The images were intensity-normalized using RGB values from air and fat reference regions to establish a unified Hounsfield Unit scale.
B. Data Preprocessing and Augmentation
The preprocessing pipeline resizes and converts CT images to grayscale, reduces image artifacts, and augments training data through geometric transformations.
- CT images are resized from 512 × 512 to 369 × 369 and converted to grayscale.
- The preprocessing reduces effects from black frames, random noise, and embedded words on segmentation.
- Data augmentation rotates existing images by 90, 180, and 270 degrees to generate 300 additional examples.
C. Experiments setting and Measure Metrics
The experiments compare the proposed model with U-Net under controlled training conditions and assess segmentation using the Dice score and 10-fold cross-validation.
- The proposed model and U-Net are trained for 100 epochs with default parameters, using both augmented and non-augmented data.
- The Dice score measures overlap between predicted and correct segmentation sets X and Y.
- Performance assessment uses 10-fold cross-validation.
D. Results
The results compare the proposed model with U-Net using quantitative metrics and visual examples, showing stronger multi-class segmentation by the proposed approach.
- Figure 5 presents two augmented-data examples with preprocessed scans, ground-truth masks, proposed-model outputs, and U-Net outputs.
- The proposed model and U-Net are compared in Table I, whose values are averages from 10-fold cross-validation.
- The original U-Net almost failed to segment one example, which the authors associate with the small region of interest and limited discrimination capability.
- The visualization distinguishes consolidation, pleural effusion, and ground-glass opacities using class-specific colors in the masks and proposed outputs.
E. Ablation Study
The paper examines ResNeXt and attention components through ablation experiments and concludes that the proposed model improves COVID-19 CT segmentation, while identifying broader generalization and interpretability as future work.
- E. Ablation Study: The ablation study separately evaluates variants without Attention and without ResNeXt under the same experimental environment as the main experiments.
- IV. Discussion and Conclusions: The proposed model uses residual feature extraction and an attention mechanism in the decoding process to produce multi-class segmentation results.
- IV. Discussion and Conclusions: The authors state that the method can help radiologists determine lung infection percentage and diagnose COVID-19 progression.
- IV. Discussion and Conclusions: Future work targets broader practical scenarios, limited-ground-truth settings, knowledge-based generative modeling, and deeper interpretability.