[논문 리뷰] The 2ST-UNet for Pneumothorax Segmentation in Chest X-Rays using ResNet34 as a Backbone for U-Net

Ayat Abedalla, Malak Abdullah외 2명|arXiv (Cornell University)|2020. 09. 06.

COVID-19 diagnosis using AI인용 수 6

한 줄 요약

이 논문은 흉부 X선에서 기흉 분할을 위해 ImageNet에서 미리 훈련된 ResNet-34 백본을 사용하는 2단계 훈련 U-Net(2ST-UNet)을 제안한다. 방법은 먼저 256×256 해상도 이미지에서 훈련한 후 512×512로 미세조정하며, 데이터 증강, 확률적 가중치 평균화(SWA), 테스트 시 증강(TTA)을 사용하여 평균 DSC(유사도 계수) 0.8356을 달성하여 2019년 SIIM-ACR 기흉 분할 챌린지 상위 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의 효과를 평가하여 분할 정확도 향상과 일반화 능력 향상 가능성 분석.

제안 방법

• 모델은 ResNet-34를 인코더 백본으로 사용하는 U-Net 인코더-디코더 아키텍처를 사용하며, 전이 학습을 위해 ImageNet에서 사전 훈련된 가중치를 활용한다.
• 네트워크는 두 단계로 훈련되며, 먼저 256×256 해상도 이미지에서 훈련한 후 동일한 아키텍처와 초기 가중치를 사용해 512×512 해상도 이미지로 미세조정한다.
• 데이터 증강 기법으로는 수평 뒤집기(A1), 무작위 밝기/대trast/감마 조정(A2), 탄성 및 격자 왜곡(A3), 무작위 크기의 자르기(A4)가 포함된다.
• 일반화 능력 향상과 모델의 강인성을 높이기 위해 훈련 중에 확률적 가중치 평균화(SWA)를 적용한다.
• 추론 시에 테스트 시 증강(TTA)을 사용하여 다수의 증강된 입력에 대한 예측 일致성 향상.
• 모델 평가는 SIIM-ACR 챌린지 데이터셋에서 DSC(유사도 계수)와 IoU(교차율)를 사용하여 평가한다.

실험 결과

연구 질문

• RQ1 진행적 해상도 스케일링을 통한 2단계 훈련이 기흉 검출에서 분할 성능 향상과 수렴 속도 향상에 기여하는가?
• RQ2 다양한 데이터 증강 전략이 분할 모델의 강인성과 일반화 능력에 어떤 영향을 미치는가?
• RQ3 SWA와 TTA를 결합할 경우 모델의 일반화 능력과 예측 일치성 향상 정도는 어느 정도인가?
• RQ4 512×512의 중간 수준의 입력 해상도로도 ResNet-34 백본을 갖춘 U-Net이 경쟁적인 성능을 달성할 수 있는가?

주요 결과

• 2ST-UNet는 사전 테스트 세트에서 평균 DSC(유사도 계수) 0.8356을 기록하여 SIIM-ACR 기흉 분할 챌린지에서 1,475개 팀 중 124위를 차지했다.
• A1과 A2 증강 기법의 조합은 TTA와 R-TH를 함께 적용했을 때 최고의 DSC 점수 0.8212를 기록했다.
• TTA와 R-TH를 함께 사용한 A1+A2 조합은 다른 모든 증강 조합보다 높은 DSC 점수를 기록하여 우수한 성능을 보였다.
• 모델의 평균 추론 시간은 1개 샘플당 25.6 밀리초로, 실시간 진단 잠재력을 보였다.
• 2단계 훈련 전략은 고해상도 이미지에서 단일 단계 훈련보다 더 빠른 수렴과 향상된 성능을 이끌어냈다.
• R-TH를 사용했을 때 A4 증강 그룹(무작위 크기 자르기)은 검증 세트에서 최고의 IoU 점수 0.7660을 기록하여 데이터 증강의 효과를 입증했다.