こんにちは。NRIデジタルの奥山です。 2022年10月までKaggleで開催されていた「RSNA 2022 Cervical Spine Fracture Detection」コンペ（以下、RSNA2022）に、NRIデジタルの大崎さんと奥山、そして社外の方1名を加えた計3名のチームで参加し、883チーム中10位で金メダル入賞となりました。 この記事ではRSNA2022コンペの概要と私たちチームの用いた解法を紹介していきたいと思います。また、KaggleのDiscussionにも解法を投稿していますので、併せてご覧いただければと思います。
コンペ概要
今回私たちが参加したRSNA2022はRSNA（北米放射線学会）が主催した医療画像の分類コンペティションです。RSNAは例年、CT画像やMRI画像を使った医療画像診断をテーマとしたコンペをホストしているのですが、2022年は頸椎のCT画像から骨折の有無を予測するというテーマでコンペが実施されました。
（頸椎を含む）脊椎の骨折は米国だけでも年間150万件以上発生しており、特に高齢者の脊椎骨折の発生率は年々増加傾向にあります。骨折の診断と言えばX線レントゲン写真を思い浮かべる方も多いと思いますが、高齢者の方の場合、レントゲン写真を使った脊椎骨折の診断は困難である事が多く、現在ではCT画像を使った骨折診断が主流となっています。機械学習を使った高速・高精度な脊椎骨折が可能となることで、早期の骨折の検出や予後の改善に繋がる可能性があります。
データ
今回は教師データとしてCT撮影約2000回分、予測データとして約3000回分のデータが提示され、1回1回のデータから椎骨（頸椎の骨）部位毎の骨折有無と椎骨全体での骨折有無を予測する事が求められました。この1回分のデータ中には身体の上下方向に輪切りにされた頸椎断面の画像が数百枚含まれており、このある1回のCT撮影画像全体のことをStudy、その中に含まれる1枚1枚の頸椎断面画像をImageと呼びます。椎骨はC1（最も頭側の骨）からC7（最もお尻側の骨）という7つの骨で構成されているため、椎骨全体と併せてひとつのStudyに対し8項目の2値分類問題を行うイメージです。
また、追加データとして一部のStudyには椎骨の位置をセグメンテーションしたマスク画像または骨折の位置をアノテーションしたバウンディングボックスがImage単位で付与されていました。ただし、椎骨マスク画像は学習データ約2000Study中87回分、骨折バウンディングボックスは約2000Study中235回分のみとStudy全体のうちごく少数のデータにしか付与されておらず、これらの追加データを効率良く扱う事がスコア向上を目指す上で重要なポイントでした。
個々のImageが512x512px、16bit深度と比較的高精細な画像データである点もコンペを難しくしているポイントでした。学習用の約2000Study分だけでも300GB近い容量があり、テスト用データはこれよりやや多い程度の画像が含まれています。一方、本コンペはCode Competitionであるため、テストデータの予測に掛けられる時間と使用可能なVMスペックが厳密に定められており、この大量の画像データを限られたリソース上で効率良く推論するようなモデルの開発が求められました。
自チームの解法
私たちのチームは、ImageレベルでCT画像の特徴量ベクトルの抽出を行うStage1モデルと、Stage1モデルが出力した特徴量ベクトルを奥行き方向の系列データと見做し、最終的な骨折有無を予測するStage2モデルを組合せた、2stage構成の解法を最終的に採用しました。
こうした2stage構成のパイプラインは過去のRSNAコンペでも良く用いられており、私たちもそれを素直に踏襲した形となります。Stage1は画像認識問題を、Stage2は系列分類問題をそれぞれ専念して解けば良いため、チームメンバー毎にタスクを分担して進めやすかった点も採用のポイントでした。
Stage1モデル
Stage1モデルは椎骨の位置をセグメンテーションするVertebraeモデルと、骨折の有無を検出するFractureモデルの2種類の2D画像モデルを直列に実行する構成になっています。どちらのモデルも、Image単位での予測を行うモデルとなっています。
データの項目でお話した通り、本コンペは推論時間の制限がかなり厳しかったため、Stage1ではあまり大規模なモデルアンサンブルは行わずに、スコアの良い少数のモデルに絞って採用しています。
- Vertebraeモデル: 1モデルのみ
- SEResNeXt50 + UNet + Classification Head
- 外部の公開CT画像データセット(Verse2020)を使用した事前学習
- Fractureモデル: 2モデルのアンサンブル
- SEResNeXt50 + UNet + Classification Head
- EfficientNet-B4 + UNet + Classification Head
- timmのImageNet事前学習済モデル
クラス分類問題にもかかわらずセグメンテーション用のUNetデコーダを使用しているのは、後述のマルチタスク学習を行うためです。下記がStage1のパイプラインです。
パイプライン
図の通り、VertebraeとFractureの両モデルで個別に学習・推論を行うのではなく
- Vertebraeモデルを学習
- Vertebraeモデルで推論を行い、椎骨の位置をセグメンテーション
- 予測したセグメンテーションマスクから椎骨の周囲のみクロップ
- クロップした画像でFractureモデルを学習
というシーケンシャルな手順で学習・推論を行っています。
この骨の付近のみをクロップするアイディアはコンペ中盤にチームメイトが考案・実装してくれたものなのですが、導入後は銀メダル圏から金メダル圏へとスコアを大きく押し上げる事に成功しました。骨折は本来椎骨の付近でのみ発生するはずですが、CT断面画像には筋肉や気道、脂肪といった椎骨以外の体組織も多く映りこんでいます。これらの影響を除去できた事が予測精度向上に繋がったものと思います。
マルチタスク学習
もう一つ、学習でスコア向上に効果的だった施策としてマルチタスク学習を導入した事が挙げられます。マルチタスク学習とはひとつのDNNモデルで複数のタスクを同時に解けるよう学習させる手法で、Kaggleではしばしば汎化性能向上の手段として採用されます。今回はVertebrae・Fractureどちらのモデルも下記の様にクラス分類問題とセグメンテーション問題を同時に解くような構成でマルチタスク学習を行いました。
Vertebraeモデル： 椎骨位置のセグメンテーションと椎骨の部位（C1-C7）を予測する多クラス分類問題
Fractureモデル： 骨折箇所のセグメンテーションと骨折有無の二値分類問題
上のモデル構成図では、ClassifierとSegmentationに枝分かれした箇所がマルチタスク学習に相当しています。
本コンペの予測対象はあくまでも骨折の有無であるため、本来骨折の位置までセグメンテーションで予測する必要はありません。一方、データの項でお話した通り本コンペはStage1の学習に利用できるデータ件数が全体のごく一部に限られていた事もあり、マルチタスク学習によって少数のデータでも効率的に学習を進めることができたと考えています。
Stage2モデル
Stage1で学習が完了したモデルはセグメンテーション・分類問題用のHeadを取り除き、Image単位の特徴量ベクトル抽出器として利用します。下図の通り、VertebraeモデルとFractureモデルから出力された特徴量ベクトルを一対一で結合し、それを体の上下方向に並べてやることで、Study単位の特徴量ベクトル列を作成します。
これをFractureの2モデル(SEResNeXt, EfficientNet)分だけ作成し、アンサンブルする
この特徴量ベクトル列全体にC1-C7、全体の骨折有無のラベルを付与してやることで、Study単位での骨折有無を学習していきます。特徴量ベクトル列はImage枚数*ベクトル次元数の1次元系列データとみなせることから、私達は下記のような系列データ分類問題で良く用いられるDNNモデルのアンサンブルを最終的に使用しました。
- Attention Bidirectional GRU
- Attention Bidirectional LSTM
- BERT-like Transformer Encoder
- DenseNet-like Bidirectional LSTM
４つとも過去のRSNAコンペのStage2で使用されていたモデルです。Stage2はStudy単位（データ件数は高々数千件）での学習・予測を行えば良いため、予測時間の制限にはほとんど影響を与えません。そこでこうした多様な系列モデルのアンサンブルを行う事で予測精度の向上を狙いました。
結局、いずれのモデルでも単体では互いにCV/LBスコアに大きな差は無かったのですが、TransformerモデルとRNNモデルをアンサンブルするとスコアが向上する傾向が見られました。
後処理
こうして得られたStage2の予測値を後処理で補正する事で、更にスコアを改善する事ができました。
上述のとおり、本コンペではC1-C7各部位の骨折有無に加え、椎骨全体で骨折が有ったか/無かったかというラベルを加えた8部位の分類が求められました。
椎骨全体での骨折有無については、C1-C7同様に2値分類問題としてStage2モデルに予測させる手法も考えられますが、一方でC1-C7の予測結果から確率的に計算する手法も考えられます。
椎骨全体のどこかが骨折している事象というのは、C1-C7がいずれも骨折していない事象の余事象となることから
$$ P_{overall} = 1 - \prod_{i=1}^n (1-P_i) $$
と表せそうですね。ここで、\(P_{overall}\)は椎骨全体の骨折確率、\(P_{i}\) はC1-C7各部位での予測確率です。この確率的に計算した骨折有無の確率をStage2モデルの予測結果とアンサンブルすることで、更にスコアをブーストすることが出来ました。
こうした機械学習に頼らない後処理テクニックはどのコンペにおいても非常に重視されていますが、本コンペでもやはり重要なファクターだったと思います。
他チーム動向と反省点
より上位チームの解法を見てみると、2Stage構成自体はどのチームもほぼ採用していましたが、Stage1に3Dモデルを採用しているチームが多数見受けられた点が印象的でした。3DモデルではStudy中のCT断面画像全体を1枚の3D画像とみなして取扱います。
例えば1位チームはStage1のVertebraeモデルで3D-UNetによるSegmentationを採用し、Study全体に対し一発で椎骨の位置をセグメンテーション＆クロップしています。これはImage単位で1枚1枚セグメンテーションしていた私たちのソリューションに比べ大幅に予測時間が短縮できていると思われます。実際、私たちは推論時間の制約からクロップ後のFractureモデルとして2モデルしか用意できなかったのに対し、1位チームでは6種類もの多様なモデルをアンサンブルに使えていた様です。
私たちのチームでもコンペ序盤は3Dモデルの開発にトライしていたのですが、3Dモデルへの経験が足らず良い予測モデルを作成する事ができませんでした。今回得た知見を活かし、次の機会があれば3Dモデルの開発に取り組んでみたいと思います。
おわりに
本コンペは全体8位まで賞金の対象で、私たちのチームは最終日の時点でPublic7位という状況でした。残念ながらPrivateでは10位とShake Downしていまい賞金は逃してしまったのですが、それでも金メダル圏にギリギリ踏みとどまれて安堵しています。
特にチームメンバーの大崎さん、HONGさんは初の金メダル獲得となり、HONGさんはこのメダルを以てKaggle Masterに昇格された点は大変嬉しく思います。私一人の力ではとても金メダルには手が届かなかったと思います。私も引き続きKaggle Grand Masterへの昇格と賞金圏入賞を目指してコンペに取り組んで行きたいと思います。
