[論文レビュー] 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

ひとこと要約

本論文は、チェストX線における気胸領域のセグメンテーションを目的として、ImageNetで事前学習されたResNet-34バックボーンを搭載した2段階学習U-Net（2ST-UNet）を提案する。この手法は、まず256×256解像度の画像で学習を行い、その後512×512解像度で微調整する。データオーグメンテーション、確率的重み平均化（SWA）、テスト時オーグメンテーション（TTA）を用い、平均Dice類似係数（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の有効性を評価し、セグメンテーションのロバスト性と汎化性能を向上させること。

提案手法

• モデルは、ImageNetで事前学習されたResNet-34をエンコーダーバックボーンとして使用したU-Netエンコーダーデコーダー構造を採用している。
• ネットワークは2段階で学習される：まず256×256解像度の画像で学習し、その後同じアーキテクチャと初期重みを用いて512×512解像度の画像で微調整する。
• データオーグメンテーション技術には、水平反転（A1）、ランダムな明るさ/コントラスト/ガンマ調整（A2）、エラスティックおよびグリッド歪み（A3）、ランダムサイズのクロップ（A4）が含まれる。
• 学習中に確率的重み平均化（SWA）を適用し、汎化性能とモデルのロバスト性を向上させた。
• 推論時にテスト時オーグメンテーション（TTA）を用い、複数のオーグメント済み入力に対する予測の一貫性を向上させた。
• モデルの評価には、SIIM-ACRコンテストデータセットにおけるDice類似係数（DSC）と交差率（IoU）が用いられた。

実験結果

リサーチクエスチョン

• RQ1 段階的解像度スケーリングを伴う2段階学習は、気胸検出におけるセグメンテーション性能と収束速度を向上させるか？
• RQ2 異なるデータオーグメンテーション戦略は、セグメンテーションモデルのロバスト性と汎化性能にどのように影響するか？
• RQ3 SWAとTTAを組み合わせることで、モデルの汎化性能と予測の一貫性はどの程度向上するか？
• RQ4 ResNet-34バックボーンを搭載したU-Netは、中程度の入力解像度（512×512）でも、気胸セグメンテーションで競争力のあるパフォーマンスを達成できるか？

主な発見

• 2ST-UNetは、プライベートテストセットで平均Dice類似係数（DSC）0.8356を達成し、SIIM-ACR気胸セグメンテーションコンテストで1,475チーム中124位にランクインした。
• A1およびA2のオーグメンテーション技術の組み合わせが、TTAとR-THを併用した際、最高のDSCスコア0.8212を記録した。
• TTAとR-THを併用することで、A1+A2グループのDSCスコアが0.8212に向上し、他のすべてのオーグメンテーション組み合わせを上回った。
• 平均推論時間は1サンプルあたり25.6ミリ秒であり、リアルタイム診断の可能性を示した。
• 2段階学習戦略により、高解像度画像での単一段階学習と比較して収束が速く、性能が向上した。
• R-THを用いた際、A4オーグメンテーショングループ（ランダムサイズのクロップ）がバリデーションセットで最高のIoUスコア0.7660を記録し、データオーグメンテーションとしての有効性を示した。