[论文解读] The 2ST-UNet for Pneumothorax Segmentation in Chest X-Rays using ResNet34 as a Backbone for U-Net

Ayat Abedalla, Malak Abdullah+2 more|arXiv (Cornell University)|Sep 6, 2020

COVID-19 diagnosis using AI被引用 6

一句话总结

本文提出一种基于ImageNet预训练ResNet-34主干网络的两阶段训练U-Net（2ST-UNet），用于胸部X光片中气胸的分割。该方法首先在256×256分辨率图像上进行训练，随后在512×512分辨率图像上微调，结合数据增强、随机权重平均（SWA）和测试时增强（TTA），在2019年SIIM-ACR气胸分割挑战赛中取得了0.8356的平均Dice相似系数（DSC），位列前9%。

ABSTRACT

Pneumothorax, also called a collapsed lung, refers to the presence of the air in the pleural space between the lung and chest wall. It can be small (no need for treatment), or large and causes death if it is not identified and treated on time. It is easily seen and identified by experts using a chest X-ray. Although this method is mostly error-free, it is time-consuming and needs expert radiologists. Recently, Computer Vision has been providing great assistance in detecting and segmenting pneumothorax. In this paper, we propose a 2-Stage Training system (2ST-UNet) to segment images with pneumothorax. This system is built based on U-Net with Residual Networks (ResNet-34) backbone that is pre-trained on the ImageNet dataset. We start with training the network at a lower resolution before we load the trained model weights to retrain the network with a higher resolution. Moreover, we utilize different techniques including Stochastic Weight Averaging (SWA), data augmentation, and Test-Time Augmentation (TTA). We use the chest X-ray dataset that is provided by the 2019 SIIM-ACR Pneumothorax Segmentation Challenge, which contains 12,047 training images and 3,205 testing images. Our experiments show that 2-Stage Training leads to better and faster network convergence. Our method achieves 0.8356 mean Dice Similarity Coefficient (DSC) placing it among the top 9% of models with a rank of 124 out of 1,475.

研究动机与目标

• 通过深度学习提高在低对比度、噪声较大的胸部X光片中气胸分割的准确性。
• 通过多阶段训练和数据增强，解决标注医学数据有限及模型泛化能力不足的问题。
• 利用基于U-Net的架构与迁移学习，在SIIM-ACR气胸分割挑战赛中实现最先进性能。
• 评估数据增强、SWA和TTA在提升分割鲁棒性和泛化能力方面的有效性。

提出的方法

• 模型采用U-Net编码器-解码器架构，以ResNet-34作为编码器主干网络，并在ImageNet上进行预训练以实现迁移学习。
• 网络采用两阶段训练：首先在256×256分辨率图像上训练，然后使用相同架构和初始权重在512×512分辨率图像上进行微调。
• 数据增强技术包括水平翻转（A1）、随机亮度/对比度/伽马调整（A2）、弹性形变与网格扭曲（A3）以及随机尺寸裁剪（A4）。
• 在训练过程中应用随机权重平均（SWA）以提升模型泛化能力与鲁棒性。
• 在推理阶段使用测试时增强（TTA）以提升多个增强输入下预测结果的一致性。
• 模型在SIIM-ACR挑战赛数据集上通过Dice相似系数（DSC）和交并比（IoU）进行评估。

实验结果

研究问题

• RQ1 采用渐进式分辨率扩展的两阶段训练是否能提升气胸检测中的分割性能与收敛速度？
• RQ2 不同数据增强策略对分割模型鲁棒性与泛化能力有何影响？
• RQ3 SWA与TTA联合使用在多大程度上提升了模型泛化能力与预测一致性？
• RQ4 基于ResNet-34主干网络的U-Net能否在中等输入分辨率（512×512）下实现具有竞争力的气胸分割性能？

主要发现

• 2ST-UNet在私有测试集上取得了0.8356的平均Dice相似系数（DSC），在SIIM-ACR气胸分割挑战赛中排名第124位（共1,475支队伍）。
• 当同时应用TTA和R-TH时，A1与A2增强技术组合取得了最高的DSC得分0.8212。
• 结合TTA与R-TH后，A1+A2组合的DSC提升至0.8212，优于所有其他增强组合。
• 模型平均推理时间为每样本25.6毫秒，表明具备实时诊断潜力。
• 两阶段训练策略相比在高分辨率图像上的一阶段训练，实现了更快收敛与更优性能。
• 在使用R-TH时，A4增强组（随机尺寸裁剪）在验证集上取得了最高的IoU得分0.7660，证明其在数据增强中的有效性。