UNetがさらなる高みへ！TransformerでUNetを強化する

segmentation 2022年02月18日

3つの要点
✔️ UNetとTransformerを組み合わせたモデルTransUNetを提案
✔️ CNNの局所性とTransformerの長期依存性の組み合わせが重要
✔️ 2つの医用画像データセットで、従来の手法を超えるセグメンテーション精度を実現

TransUNet: Transformers Make Strong Encoders for Medical Image Segmentation
written by Jieneng Chen, Yongyi Lu, Qihang Yu, Xiangde Luo, Ehsan Adeli, Yan Wang, Le Lu, Alan L. Yuille, Yuyin Zhou
(Submitted on 8 Feb 2021)
Comments: Published on arxiv.
Subjects: Computer Vision and Pattern Recognition (cs.CV)

code：

本記事で使用している画像は論文中のもの、紹介スライドのもの、またはそれを参考に作成したものを使用しております。

はじめに

医用画像のセグメンテーションは、医療アプリケーションの前処理として非常に重要であるため日々研究が進められています。近頃は、深層学習を使用したモデルが高いセグメンテーション精度を実現しています。

医用画像のセグメンテーションを行うモデルとして最も成功を収めているのが、UNetと呼ばれるU字型のアーキテクチャを持つCNNモデルです。しかしながら、そんなUNetにも、セグメンテーションを行う上での弱点が存在します。UNetの弱点とは、セグメンテーションを行う上で長期依存性を捉えることが苦手であるということです。原因は、UNetを構成するCNNが局所的な特徴を捉えることが得意なのに対して、長期的な特徴を捉えることには限界があるためだと言われています。

そんなUNetの弱点を克服させてくれるのが、Transformerです。Transformerの強みは、長期依存性を捉えることです。したがって、TransformerがUNetの弱点を補い、セグメンテーション精度を向上させてくれることが期待できます。

本論文では、UNetとTransformerを組み合わせたTransUNetというモデルを提案しています。TransUNetは、局所的な特徴を捉えることに長けているCNNと、長期的な特徴を捉えることに長けているTransformerを上手く組み合わせることで、従来の手法よりも正確なセグメンテーションを行うことが可能になりました。

その結果、2つの医用画像データセットで従来の手法を超えるセグメンテーション精度を実現しました。また、CNNとTransformerを単体で使うよりも、CNNとTransformerを組み合わせることで、より正確なセグメンテーションが行えることが実験によりわかりました。

本記事では、TransUNetの概要と、医用画像データセットを使用した実験の結果について紹介していきます。

TransUNet

上の図は、TransUNetのアーキテクチャを表しています。TransUNetを簡単に説明すると、UNetのエンコーダにTransformer (ViT)を組み込んだモデルになっています。以下では、TransUNetのエンコーダとデコーダについて説明します。

TransUNetのエンコーダでは、はじめにCNNで特徴抽出を行い、局所的な特徴を捉えます。そのあと、Transformerで特徴抽出を行い、長期的な特徴を捉えていきます。また、TransUNetでは、CNNとTransformerとして、それぞれImageNetで学習済みのResNet50及びViTを使用しています。

TransUNetのデコーダでは、UNetと同様にアップサンプリングを行い、最終的にセグメンテーションの結果を出力します。また、エンコーダのCNNとデコーダの対応する層をskip-connectionでつないでいます。

実験

医用画像データセット

本論文では、以下の2つの医用画像データセットを使用して、セグメンテーション実験を行います。

Synapse multi-organ segmentation dataset
- 腹部CT画像のデータセット
- 8部位のセグメンテーション
Automated cardiac diagnosis challenge (ACDC)
- 心臓部のMRIのデータセット
- 3部位のセグメンテーション

評価

モデルの評価には、Dice係数 (DSC、単位は%)とハウスドルフ距離 (HD、単位はmm)を使用しています。DSCは大きいほどセグメンテーション精度が高く、HDは小さいほどセグメンテーション精度が高いことを表しています。

Synapse multi-organ segmentation datasetでのセグメンテーション精度は次のようになりました。

TransUNet (DSC:77.48 %、HD:31.69 mm)で、従来の手法 (V-Net、DARR、U-Net、AttnUNet)を超えるセグメンテーション精度を実現しています。また、TransUNetが、エンコーダがTransformerのみのモデル (EncoderがViT、DecoderがCUPのモデル)よりもセグメンテーション精度が高いことから、CNNとTransformerを組み合わせることが重要であることがわかりました。

ACDCデータセットでのセグメンテーション精度は次のようになりました。

ACDCデータセットでも、従来の手法 (R50-U-Net、R50-AttnUNet)およびTransformerのみのモデルと比較して、TransUNetのセグメンテーション精度 (DSC:89.71 %)が最も高くなっています。

セグメンテーションの可視化

上の図には、Synapse multi-organ segmentation datasetにおいて、実際にセグメンテーションを行ったときのセグメンテーション画像を示しています。TransUNetは他のモデルと比較すると、より正確なセグメンテーションを行えていることがわかります。

例えば、2行目のセグメンテーション画像で比較してみます。UNetはleft kidney (赤)の部分を誤って spleen (水色)にセグメンテーションしていたり、AttnUNetは、spleen (水色)の部分を誤ってliver (紫)にセグメンテーションしています。一方で、TransUNetは、spleen (水色)の部分を正確にセグメンテーションできていることがわかります。