二次元画像のセグメンテーション

2024年12月30日更新

二次元画像のセグメンテーション

IEEE International Symposium on Biomedical Imaging(ISBI) が開催しているコンペティションの中の一つであるISBI 細胞追跡チャレンジ 2015を大差で勝ち取った U-Net: Convolutional Networks for Biomedical Image Segmentation(2015) は、現在でも広範囲に活用されている二次元画像セグメンテーション技術です。一方、医療用画像や自動運転の分野では、より高精度なセグメンテーション技術が望まれているのも事実です。

二次元画像のセグメンテーション技術は 2024 年現在でも進化を続けています。本ページでは、2015年～2024年の二次元画像セグメンテーション技術の進展を整理すると共に、画像セグメンテーションに取り組む際の注意点についてまとめましたので、参考にして頂ければ幸いです。

特筆すべき点は、 2023年～2024年に発表された Mamba 系の二次元画像セグメンテーションへの応用が急速に進展している点でしょう。Mamba は、LLM モデルのデファクトスタンダードになっている Attention ベースの Transformer - Attention Is All You Need(2017) の牙城を崩す可能性を秘めた新たな技術です。大規模な GPU メモリーを必要とし、計算量が多いけれども、CNN や RNN などの従来の状態空間モデルの性能を劇的に上回った Transformer が、今度は、GPU メモリーも計算量も少なくて済む選択的状態空間モデル Mamba により大逆転されるかも知れません。

二次元画像セグメンテーションの現状

図1 は、"MSVM-UNet: Multi-Scale Vision Mamba UNet for Medical Image Segmentation(2024) による、Synapse 多臓器データセットのセグメンテーション例です。真値(ground truth)に対して、UNet、TransUNet、Swin-UMamba、VM-UNet、および MSVM-UNet のセグメンテーション予測画像が示されています。特に、 MSVM-UNet の優位性が、赤い長方形で強調表示された臓器領域ではっきりと確認できます。

図1 Synapse 多臓器データセットセグメンテーションの予測画像比較例

表1 は、Synapse 多臓器データセットセグメンテーションの定量的比較です。 MSVM-UNet の優位性が示されていますが、どの手法もセグメンテーションする臓器により、その性能に差があることがわかります。特に、Gallbladder（胆嚢）、Kidney（腎臓）、 Pancreas（膵臓）など、小さく、人間の目でも見分けがつきにくい臓器であるほどセグメンテーションの精度が低くなるようです。

表1 Synapse 多臓器データセットセグメンテーションの定量的比較

ここで、DSC (Dice Similarity Coefficient)は、画像セグメンテーションの精度評価によく使用される指標であり、以下のように定義されます。

\[ \mathrm{DSC} = \frac{2 \cdot |G \cap P|}{|G| + |P|} \]

ここで：

$G$ は Ground Truth (真値)
$P$ はモデルによる予測値
$|G \cap P|$ は Ground Truth と予測値の重なり合う領域の面積（ピクセル数）
$|G|$ は Ground Truth の面積（ピクセル数）
$|P|$ は予測値の面積（ピクセル数）

また、HD95 (95% Hausdorff Distance; ハウスドルフ距離)は、ground truthと予測値の距離の近さを表す指標の一つであり、 Ground truthのセグメンテーション境界点集合を

G_{B}

、予測値境界点集合を

P_{B}

とすれば、以下の方法で計算されます。

$G_{B}$ の各点から $P_{B}$ に対する最近傍点の距離を計算します。
逆に、 $P_{B}$ の各点から $G_{B}$ に対する最近傍点の距離を計算します。
全ての距離を昇順に並べ、95% の位置の距離が HD95 になります。

この HD95 は、「95%のセグメンテーション境界点がこの距離以内にある」ことを意味し、残りの5%該当する外れ値の影響を軽減した、より安定した距離指標を与えます。数式で表すと：

\[ \mathrm{HD95} = \mathrm{percentile}_{95} \left(\left\{ \max\left(\min_{p_j \in P_B} d(g_i, p_j), \min_{g_i \in G_B} d(p_j, g_i)\right)\right\}\right) \]

ここで、$d(\cdot,\cdot)$ は2点間の距離を表します。

二次元画像分類・セグメンテーション技術の進展

CNN 系の画像分類・セグメンテーション技術

二次元画像セグメンテーション技術は、ニューロン構造のセグメント化に関する ISBI 細胞追跡チャレンジ 2015 において、他の手法に大差でつけて UNet(2015) が勝ち取ったことで大きく進展しました。UNetは、4×(畳み込み層２層 + プーリング層１層) の encoder で構成され、パラメータ数も 4.1MB とコンパクトな構造であるにも拘わらず、少ない画像からエンドツーエンドでトレーニングでき、512 x 512 画像のセグメントが、当時の GPU で 1 秒以内でできるようになりました。

その後、CNN 系列のセグメンテーションは、encoder 部の層が深くなった ResNet (2015) 、 encoder-decoder がU字構造がネストして網目状になった UNet++(2018) などへと発展しました。また、画像分類の分野では、桁違いに少ないパラメータで高速な推論で最先端の精度を達成した efficientnet (2019) も出現しました。

Transformer 系の画像分類・セグメンテーション技術

画像の局所的な特長量を捉えることが得意な CNN は、どうしても画像全体を捉えることができないため、汎化性能を向上することが難しい傾向にあります。この弱点を解決したのが、 Attention Is All You Need(2017) の発表以降、LLM モデルのデファクトスタンダードになっていた Transformer を画像に応用した Vision Transformer(ViT) - An Image is Worth 16×16 Words: Transformers for Image Recognition at Scale(2020) です。ViT では、入力画像を16×16ピクセルの小さなパッチに分割し、これを自然言語処理のトークンと同様、マルチヘッド自己アテンションによって処理することにより、高性能な画像分類ができるようになり注目を集めました。

しかし、ViT の汎化性能を高めるためには、大規模なデータセットで学習する必要がある。BEiT(2022) のように、ImageNet-1K や ImageNet-22Kといった大規模な画像データセットを用いて、 BERT(2018) で事前トレーニングする方法も提案されているが、やはり高い計算コストの制約から大画像での実用化には制限がある。また、ViT は全体的な関係性を捉える能力は高いけれども、局所的な特徴を捉える能力が CNN に比べて劣ることがあるという課題もある。

このような ViT の課題を解決したのが、 SegFormer (2021) や Swin Transformer など、Hierarchical Transformer（階層的 Transformer）です。階層的 Transformer では、複数の抽象化レベルやスケールを扱います。ViT では、16×16 のパッチを扱いますが、例えば Swin Transformer では、初期の層で小さなパッチ（例えば4×4）を使用して局所的な特徴を抽出し、次第に大きなパッチ（例えば8×8や16×16）にパッチを統合することにより、広範囲な情報も捉えることができます。そして、各層で異なるレベルの抽象化を行うことにより、画像全体のコンテキスト情報ばかりではなく、細部の情報も同時に学習できるので精度が高まります。また、階層的トランスフォーマーは、 ViT に比較して、階層的な情報処理やメモリ階層構造の最適化などにより、 GPU メモリー使用量を削減する工夫が施されています。

Mamba 概要

長くなりますので Mamba の解説は別ページに記述し、このページでは Mamba の要点のみまとめます。2023年～2024年に発表された選択的状態空間モデル Mamba: Linear-Time Sequence Modeling with Selective State Spaces(2023) は、従来のAI研究の主流であった「状態空間」を巧みに使用した「状態空間モデル」の一種です。

ここで、状態空間モデルとは、観測できない「状態」と観測される「観測値」を分けて考える数学的なフレームワークです。このモデルでは、時間と共に変化するシステムの内部状態である「状態空間」を推定し、それが観測値にどのように影響を与えるかを分析します。「状態空間」は、潜在的な変数を扱いますので「潜在空間」と呼ばれることもあり、潜在変数で構成された層は「隠れ層」と呼ばれます。

Mamba は、図2に示す通り、"S6":選択的状態空間モデル（S4 with a Selection mechanism and computed with a Scan: Selective SSMs）とも呼ばれ、"S4":構造化状態空間モデル（Structured State Space Sequence models: Structured SSMs）の発展形です。図2 で (Δ, B, C) が時間に依存するのが "S6" であり、時間に依存しないのが "S4" です。ただ、"S4" も "S6" も、 CNN と RNN の両方を巧みに使用していますので、具体的にはもう少し複雑です。

図2 Mamba: 選択的状態空間モデル（Selective SSMs）

Mamba 系の画像分類・セグメンテーション技術

ViT と同様、入力画像を16×16ピクセルの小さなパッチに分割して自然言語処理のトークンと同様に扱う手法を Mamba に展開したのが Vision Mamba (Vim:2024) です。この手法は、従来の ViT を高速化し、メモリを節約することができましたが、ViT と同様、全体的な関係性を捉える能力は高いけれども、局所的な特徴を捉える能力が劣るという課題を抱えていますので、その性能にも限界があるようです。

その後、このような課題を解決するため、Mamba-UNet (2024)、SMM-UNet (2024)、MSVM-UNet (2024) などの新技術が矢継ぎ早に開発され続けていますが、これらの技術を系統的に評価した例はまだありません。また、 SMP や mmdetection のように、事前学習重みを利用して簡単に実問題に応用できるモジュールも、2024年現在ではまだ出現していないようです。

二次元画像分類・セグメンテーション技術の進展

表2 二次元画像分類・セグメンテーション技術の進展

分類	基本モデル	Encoder (backbone)	(Encoder) Params	事前学習重み利用の可否	計算量推論速度（論文）	精度（論文）				備考
分類	基本モデル	Encoder (backbone)	(Encoder) Params	事前学習重み利用の可否	計算量推論速度（論文）	COCO _PQ (*1)	COCO _AP (*2)	ADE20K _mIoU (*3)	ImageNet Top-1 (*4)	備考
CNN	UNet(2015)	4層のFCN(*5)	4.1M	-	-	セグメンテーションの基本モデル				ISBI細胞追跡チャレンジ(*6)2015で1位
	UNet	ResNet50	23M	SMPで可	-	-	-	-	74.9	-
	UNet	ResNet152 (2015)	58M	SMPで可	-	-	-	-	77.8	ILSVRC 2015の画像分類部門で優勝
	UNet++(2018)	最大4層のネストしたFCN	23M	SMPで可	-	結腸ポリープ、細胞核、肝臓セグメンテーションで、UNet の mIoU 対比 +3.9%				-
	UNet++	ResNet50	23M	SMPで可	-	-	-	-	-	-
	Mask R-CNN (2017)	ResNet-101-FPN	50-FPN で44M	mmdetection で可	5fps (Frames/秒)	-	(35.7:test-dev)	-	-	ICCV2017で最優秀論文賞を受賞,トレーニングが簡単
	Mask R-CNN	ResNeXt-101-FPN	50-FPN で44M	mmdetectionで可	-	-	(37.1:test-dev)	-	-	ICCV2017で最優秀論文賞を受賞,トレーニングが簡単
	UNet	efficientnet (2019) -b5	30M	SMPで可	9.9 GFLOPs	-	-	-	83.6	-
	UNet	efficientnet-b7	66M	SMPで可	37 GFLOPs	-	-	-	84.3	ILSVRCで2019年当時最高レベル、既存の ConvNet よりも 8.4 倍小さく、推論が 6.1 倍高速
Transformer	UNet	MiT(*7) mit_b2	24M	SMPで可	-	-	-	-	-	-
	SegFormer (2021)	SegFormer-B2 (SMPのmit_b2)	24M	SMPで可	62.4 GFLOPs	-	(44.6:test-dev)	47.5	-	SegFormer-B5のゼロショット堅牢性をPR
	SegFormer	SegFormer-B4(SMPのmit_b4)	60M	SMPで可	95.7 GFLOPs 15.4fps	-	(46.5:test-dev)	51.1	-	-
	Mask2Former (2022)	ResNet50	44M	mmdetectionで可	226 GFLOPs 8.6fps	51.9	43.7	49.2	-	トレーニングメモリがMaskFormerの1/3の6GBに減少。パノプティック、インスタンス、セマンティックに使用可
	Mask2Former	Swin Transformer -L	216M	mmdetectionで可	868 GFLOPs 4.0fps	57.8	50.1	57.7	-	トレーニングメモリがMaskFormerの1/3の6GBに減少。パノプティック、インスタンス、セマンティックに使用可
	TransNeXt-Base（2024）	TransNeXt	89.7M	-	56.3GFLOPs	?	?	53.0	86.2	時間の経過と共に予測精度が低下する「深度劣化」を回避し、グローバルな視覚と局所的な視覚を併せ持つ Aggregated Attention を導入
	UperNet	DeiT(2020) -S	22M	-	-	-	-	44.0	79.8	ViTベース
	UperNet	DeiT-B	86M	-	-	-	-	-	81.8	ViTベース
	UperNet	BEiT(2022) -base	87M	BEiT で可	-	-	-	45.58	83.2	ViT(Vision Transformer)ベース
	UperNet	BEiT-large	305M	BEiTで可	-	-	-	-	85.2	ViTベース
Mamba	UperNet	Vision Mamba (Vim:2024) -S	46M	-	13M の Vim-Ti で 2.52fps	-	-	44.9	80.3	大画像に対しVRAM消費量少ない
	UNet	Mamba-UNet (2024)	-	-		CT腹部 segmentation で UNet を2%程度改善				-
	UNet	SMM-UNet (2024)	0.038M?	-	-	皮膚病変セグメンテーションタスクでIoUで UNet対比 3.48%-3.91%改善				-
	UNet	MSVM-UNet (2024)	36M	-	15.5 GFLOPs	Synapse多臓器セグメンテーションのDSCでUnet比+10.18%, VM-UNet比+2.62%				-

(*1) COCO_PQ = SQ × RQ - COCO val2017 with 80 categories データセットで、パノプティックセグメンテーションの性能を測る指標であり、SQはセマンティックセグメンテーションの品質（正しく分類されたピクセルの割合）, RQはインスタンスセグメンテーションによって正しく認識されたオブジェクトインスタンスの割合。COCO val2017 with 80 categories には、人物、動物、乗り物、家具など 80 カテゴリに分類された、5,000枚の画像が含まれている。スコアを向上するためには高い汎化性能が要求されます。

【セグメンテーションの種類】

セマンティックセグメンテーション - 画像内のすべてのピクセルをクラス単位で分類する手法
インスタンスセグメンテーション - 前景にある同一クラス内の個々のオブジェクトを区別して分類する手法
パノプティック（Panoptic）セグメンテーション - 画像内のすべてのピクセルをクラス単位で分類し、さらに同一クラス内の個々のオブジェクトを区別して分類する手法

(*2) COCO_AP(Average Precision)= mAP@0.50:0.95 - COCO val2017 with 80 categories データセットで、インスタンスセグメンテーションの性能を評価するための指標であり、各IoU閾値(0.50, 0.55, ..., 0.95)に対するPrecision-Recall（成功率）カーブの下の面積の平均です。スコアを向上するためには高い汎化性能が要求される。古い論文は COCO test-dev データセットを使用（test-dev と表記し括弧書き）。

(*3) ADE20K_mIoU - MITのADE20Kデータセットの mIoU(Mean Intersection over Union)。ADE20K には、人物、空、道路、車、家具など 20,000枚以上の画像が150以上の異なるカテゴリで含まれている。カテゴリあたりの画像数が少ないため、スコアを向上するためには高い汎化性能が要求されます。ここで、IoU (Intersection over Union) は、Ground Truth (G) と予測値 (P) の重なりを評価する指標で、以下のように定義されます。

\[ \mathrm{IoU} = \frac{|G \cap P|}{|G \cup P|} \]

$|G \cap P|$ は Ground Truth と予測値の共通部分（重なっている部分）の面積
$|G \cup P|$ は Ground Truth と予測値の和集合の面積

(*4) ImageNet Top-1: 参考のため Encoder 部の画像分類の性能を示す。ImageNet の validation set（「ImageNet-1K」の一部）で、224x224の画像を使用し、モデルが予測した最も確率の高いクラスが正解クラスと一致した割合。

(*5) fully convolutional networks (FCN): CNNのように全結合層を使用せず、全ての層を畳み込み層とプーリング層で構成したネットワーク。UNetでは 4×（畳み込み層２層　+　プーリング層１層）の構成になっている。

(*6) "IEEE International Symposium on Biomedical Imaging (ISBI) が開催しているコンペティション":https://biomedicalimaging.org/2024/accepted-challenges-tutorials/ の中の "Cell Tracking Challenge":https://celltrackingchallenge.net/ であり、細胞のセグメンテーションと追跡のアルゴリズム開発に焦点を当てている

(*7) Mix Vision Transformer encoders (MiT): SegFormer の階層的Transformerエンコーダ部であり、MiT-B0 ～ MiT-B5がある。データ拡張の一環として、Mixup, CutMixという2種類のデータ拡張手法を統合する Mix-ViT と混同しないよう注意。

画像セグメンテーションのデータセットについて

COCO(Common Objects in Context): Microsoft が提供する物体検出やセグメンテーションなどのタスクに広く使用される大規模な画像データセット(Creative Commons Attribution 4.0 License なので著作権表示すれば商用利用可能）
ADE20K(A Dataset for Semantic Segmentation and Scene Parsing): セマンティックセグメンテーションに特化したMITの大規模なデータセット（商用利用不可）。
ImageNet: スタンフォード大学が運営する画像分類や物体検出用の約1400万枚の高解像度画像を含む大規模なデータセットで、21,000カテゴリに分類されている。ILSVRC で使用された1,000クラス版は「ImageNet-1K」であり、完全版は「ImageNet-21K」（商用利用不可)。