[论文解读] COVID-CAPS: A Capsule Network-based Framework for Identification of COVID-19 cases from X-ray Images

Parnian Afshar, Shahin Heidarian+4 more|arXiv (Cornell University)|Apr 6, 2020

COVID-19 diagnosis using AI被引用 5

一句话总结

本文提出 COVID-CAPS，一种基于胶囊网络的深度学习框架，用于从胸部X光片中检测COVID-19，通过注意力机制和空间层次结构保留空间关系，并在小样本数据集上提升性能。在未进行预训练的情况下，该模型实现95.7%的准确率、90%的敏感度、95.8%的特异度和0.97的AUC；通过在相似X光数据集上进行预训练，准确率进一步提升至98.3%，特异度提升至98.6%。

ABSTRACT

Novel Coronavirus disease (COVID-19) has abruptly and undoubtedly changed the world as we know it at the end of the 2nd decade of the 21st century. COVID-19 is extremely contagious and quickly spreading globally making its early diagnosis of paramount importance. Early diagnosis of COVID-19 enables health care professionals and government authorities to break the chain of transition and flatten the epidemic curve. The common type of COVID-19 diagnosis test, however, requires specific equipment and has relatively low sensitivity. Computed tomography (CT) scans and X-ray images, on the other hand, reveal specific manifestations associated with this disease. Overlap with other lung infections makes human-centered diagnosis of COVID-19 challenging. Consequently, there has been an urgent surge of interest to develop Deep Neural Network (DNN)-based diagnosis solutions, mainly based on Convolutional Neural Networks (CNNs), to facilitate identification of positive COVID-19 cases. CNNs, however, are prone to lose spatial information between image instances and require large datasets. The paper presents an alternative modeling framework based on Capsule Networks, referred to as the COVID-CAPS, being capable of handling small datasets, which is of significant importance due to sudden and rapid emergence of COVID-19. Our results based on a dataset of X-ray images show that COVID-CAPS has advantage over previous CNN-based models. COVID-CAPS achieved an Accuracy of 95.7%, Sensitivity of 90%, Specificity of 95.8%, and Area Under the Curve (AUC) of 0.97, while having far less number of trainable parameters in comparison to its counterparts. To further improve diagnosis capabilities of the COVID-CAPS, pre-training based on a new dataset constructed from an external dataset of X-ray images. Pre-training with a dataset of similar nature further improved accuracy to 98.3% and specificity to 98.6%.

研究动机与目标

• 解决COVID-19诊断中医学影像数据集有限且不平衡的挑战。
• 克服传统CNN在数据稀缺条件下难以保留空间关系和泛化能力的局限性。
• 开发一种深度学习框架，在参数更少、计算需求更低的情况下保持高性能。
• 通过在疾病特异性X光数据集上进行预训练，而非自然图像数据集，提升诊断准确性和鲁棒性。
• 提供一个公开可用、高效且准确的模型，用于从胸部X光片中实现COVID-19的早期检测。

提出的方法

• 该框架采用胶囊网络架构，结合动态路由机制，以保留空间层次结构，并提升在图像变换下的泛化能力。
• 集成一种改进的损失函数，以缓解训练数据中的类别不平衡问题，增强模型鲁棒性。
• 模型在包含正常、细菌性、非COVID-19病毒性及COVID-19病例的二值化X光数据集上进行端到端训练。
• 在来自NIH胸部X光数据集的大型外部X光数据集（94,323张图像）上进行预训练，按五类疾病分组。
• 通过微调预训练模型在目标COVID-19分类任务上应用迁移学习。
• 最终模型使用标准指标（准确率、敏感度、特异度和AUC）进行评估，并辅以ROC曲线分析。

实验结果

研究问题

• RQ1 在小样本、类别不平衡的数据集上训练时，基于胶囊网络的模型是否能优于CNN在X光片中分类COVID-19的表现？
• RQ2 与在自然图像数据集上预训练相比，在疾病特异性X光数据集上预训练是否能提升诊断性能？
• RQ3 胶囊网络在保留空间关系方面的能力在医学影像中在多大程度上增强了泛化能力并减少了过拟合？
• RQ4 所提模型的可训练参数数量与现有基于CNN的模型相比如何？对计算效率有何影响？
• RQ5 误分类的假阳性病例中，错误分类的模式是什么？哪些非COVID-19病例最常被误分类？

主要发现

• 在未进行预训练的二值化X光数据集上，COVID-CAPS实现了95.7%的准确率、90%的敏感度、95.8%的特异度和0.97的AUC。
• 在大型疾病特异性X光数据集上进行预训练后，模型准确率显著提升至98.3%，特异度提升至98.6%。
• 预训练后AUC从0.97提升至0.99，表明模型判别能力显著增强。
• 假阳性病例中，绝大多数为正常X光片（54%），其次是细菌性（27%）和非COVID-19病毒性（19%）病例。
• COVID-CAPS仅使用295,488个可训练参数，远低于先前工作中类似CNN模型的2300万参数。
• 尽管数据稀缺，模型性能依然稳健，证明胶囊网络在低数据场景下的有效性。