Classification of COVID-19 cases from chest CT volumes using hybrid model of 3D CNN and 3D MLP-Mixer

Technical Summary: Classification of COVID-19 Cases from Chest CT Volumes Using a Hybrid 3D CNN and 3D MLP-Mixer

Problem Statement
The rapid global spread of Novel Coronavirus Disease 2019 (COVID-19) has created a significant manpower shortage in medical institutions, necessitating efficient diagnostic workflows. While Reverse Transcriptase Polymerase Chain Reaction (RT-PCR) is the standard diagnostic tool, its sensitivity is reported to range only between 42% and 71%. In contrast, chest CT-based diagnosis demonstrates a reported sensitivity of 97%, making it a critical modality for identifying viral pneumonia. However, manual interpretation of CT volumes is time-consuming. Existing automated Computer-Aided Diagnosis (CAD) systems often rely on Convolutional Neural Networks (CNNs), which are limited by their receptive fields and primarily utilize local image features. Conversely, Vision Transformers (ViT) consider global features but require vast amounts of training data to achieve high performance, which is often unavailable in medical imaging contexts. The challenge lies in developing an automated classification method that effectively leverages both local and global features from 3D CT volumes while remaining robust with limited training data.

Methodology
The authors propose an automated classification method for chest CT volumes using a hybrid model combining a 3D Convolutional Neural Network (3D CNN) and a 3D version of the MLP-Mixer.

• Pre-processing: The input chest CT volumes undergo lung segmentation to isolate lung regions. Axial slices lacking lung tissue are removed, and the remaining sub-volume is scaled to a fixed resolution of $192 \times 192 \times 128$ voxels.
• 3D CNN Feature Extraction: The scaled sub-volume is processed by a 3D CNN to extract local image features. This section employs dense pooling connections, utilizing a combination of max and average pooling (mixed pooling) to reduce information loss and capture spatial information across multiple scales. The output is a 3D feature volume.
• 3D Patch Generation: Non-overlapping 3D patches of size $4 \times 4 \times 4$ voxels are extracted from the 3D feature volume. Each patch is reshaped into a 1D vector with $C=64$ components.
• 3D MLP-Mixer Classification: The authors extend the original 2D MLP-Mixer architecture to handle 3D volumes. The set of 3D patches serves as the input ($X \in \mathbb{R}^{S \times C}$) to the 3D MLP-Mixer. The model processes these inputs through $L$ layers, where each layer consists of two Multi-Layer Perceptron (MLP) blocks:
  1. Token-mixing: Performs cross-patch feature mixing to integrate global context.
  2. Channel-mixing: Performs intra-patch feature mixing.
     The final output is passed through global average pooling and fully connected layers to produce a binary classification (High vs. Low likelihood of COVID-19).

Key Contributions

1. Hybrid Architecture: The paper introduces a novel hybrid model that integrates 3D CNNs for local feature extraction with a 3D MLP-Mixer for global feature integration. This addresses the limitation of pure CNNs (local focus) and pure Transformers/MLP-Mixers (data hunger).
2. 3D MLP-Mixer Adaptation: The authors successfully modified the MLP-Mixer, originally designed for 2D images, to process 3D volumetric data by adapting the patching and mixing mechanisms to 3D space.
3. Efficiency with Limited Data: The hybrid approach is designed to achieve high classification performance even with a limited number of training samples, a common constraint in medical imaging, by leveraging the 3D CNN to pre-extract robust local features before global mixing.

Experimental Results
The proposed method was evaluated on a dataset comprising 1,205 chest CT volumes (COVID-19 and non-COVID-19 cases) obtained from multiple medical institutions in Japan. The dataset was split into 80% for training and 20% for testing.

• Performance: The proposed hybrid model achieved a classification accuracy of 79.5%.
• Comparison: A baseline 3D CNN model (consisting of 3D CNN layers and simple MLP layers) trained on the same dataset achieved an accuracy of 74.8%.
• Conclusion: The 4.7% improvement in accuracy demonstrates that the 3D MLP-Mixer component effectively contributes to the classification performance by utilizing global image features.

Significance and Claims
The paper claims that the hybrid use of 3D CNN and 3D MLP-Mixer is effective for 3D volume classification, particularly for viral pneumonia where ground glass opacities and consolidations occur over large areas, requiring both local and global feature analysis. The authors emphasize that this hybrid structure allows for high accuracy without the massive training datasets required by pure MLP-Mixer or ViT models.

However, the authors maintain a modest tone regarding clinical deployment. They state that further improvements in classification accuracy are required before the model can be applied in clinical settings. Future work identified by the authors includes the use of data augmentation to enhance anatomical and intensity variations, the incorporation of international datasets, and the continued development of a full COVID-19 CAD system. The paper does not claim the current model is ready for immediate clinical use but rather presents it as a promising step toward that goal.