[論文レビュー] Prototype Learning for Out-of-Distribution Polyp Segmentation

Nikhil Kumar Tomar, Debesh Jha+1 more|arXiv (Cornell University)|Aug 7, 2023

Colorectal Cancer Screening and DetectionMedicine被引用数 3

ひとこと要約

本論文では、WLI、BLI、LCI、FICEの多様な腸内視鏡画像モality（WLI、BLI、LCI、FICE）において、分布外（OOD）ポリープ分類のためのロバスト性を向上させる、プロトタイプ学習に基づくセグメンテーションモデル、PrototypeLabを提案する。粗いマスク生成モジュール（CMGM）、プロトタイプ生成モジュール（PGM）、プロトタイプマスク生成モジュール（PMGM）を活用し、形状、テクスチャ、色の特徴を捉えるクラス固有のプロトタイプを学習することで、OODデータセットでDSC ≥ 90%、mIoU ≥ 85%を達成し、ほぼリアルタイムの推論を実現する。

ABSTRACT

Existing polyp segmentation models from colonoscopy images often fail to provide reliable segmentation results on datasets from different centers, limiting their applicability. Our objective in this study is to create a robust and well-generalized segmentation model named PrototypeLab that can assist in polyp segmentation. To achieve this, we incorporate various lighting modes such as White light imaging (WLI), Blue light imaging (BLI), Linked color imaging (LCI), and Flexible spectral imaging color enhancement (FICE) into our new segmentation model, that learns to create prototypes for each class of object present in the images. These prototypes represent the characteristic features of the objects, such as their shape, texture, color. Our model is designed to perform effectively on out-of-distribution (OOD) datasets from multiple centers. We first generate a coarse mask that is used to learn prototypes for the main object class, which are then employed to generate the final segmentation mask. By using prototypes to represent the main class, our approach handles the variability present in the medical images and generalize well to new data since prototype capture the underlying distribution of the data. PrototypeLab offers a promising solution with a dice coefficient of $\geq$ 90\% and mIoU $\geq$ 85\% with a near real-time processing speed for polyp segmentation. It achieved superior performance on OOD datasets compared to 16 state-of-the-art image segmentation architectures, potentially improving clinical outcomes. Codes are available at https://github.com/xxxxx/PrototypeLab.

研究の動機と目的

• 異なる臨床施設の分布外（OOD）データセットにおいて、既存のポリープ分類モデルの一般化性能の低さを解決すること。
• 小規模で平坦、部分的にしか見えない、または camouflage されたポリープが示す明確でない境界と変動する画像条件に起因する課題を克服すること。
• WLI、BLI、LCI、FICEを含む多様な画像モダリティおよびマルチセンターデータセットにわたって、良好に一般化するロバストなセグメンテーションモデルを開発すること。
• 内在するデータ分布の特徴を捉えるためにプロトタイプ学習を統合し、臨床現場におけるモデルの一般化性能と信頼性を向上させること。
• 実用的なコンピュータ支援診断（CADx）システムへの導入に適した、高い精度とほぼリアルタイムの推論速度を達成すること。

提案手法

• 入力された腸内視鏡画像からマルチスケール特徴を抽出するために、ピラミッドビジョントランスフォーマー（PVT）エンコーダーを活用する。
• 初期の粗いマスクを生成し、プロトタイプ学習をガイドするために、粗いマスク生成モジュール（CMGM）を実装する。
• 潜在的なデータ分布を捉えることで、クラスごとに複数のプロトタイプを学習するプロトタイプ生成モジュール（PGM）を設計し、特徴表現を強化する。
• マルチスケール特徴の統合を可能にし、特徴の豊かさを向上させるために、エンコーダー特徴統合モジュール（EFFM）を導入する。
• 学習されたプロトタイプとデコーダー特徴を統合し、最終的なセグメンテーションマスクを生成するプロトタイプマスク生成モジュール（PMGM）を採用する。
• 分布内およびOODデータの両方でセグメンテーション性能を最適化するために、クロスエントロピー損失とDice損失を用いて、エンドツーエンドのアーキテクチャを訓練する。

Figure 1 : The overall architecture of the proposed PrototypeLab. The input image is fed to the PVT-encoder to generate a coarse mask, which is combined with encoder features in the prototype generation module to generate various prototypes. Subsequently, the decoder is employed, which produces a fe

実験結果

リサーチクエスチョン

• RQ1 プロトタイプ学習は、異なる臨床施設の分布外（OOD）データセットにおけるポリープ分類モデルの一般化性能を向上させることができるか？
• RQ2 提案されたPrototypeLabは、多様な画像モダリティにおいて、微小、平坦、または部分的にしか見えない病変のような挑戦的なポリープタイプをどれほど効果的に処理できるか？
• RQ3 CMGM、PGM、PMGMの各モジュールが、ベースラインアーキテクチャと比較して、セグメンテーション精度とロバスト性をどの程度向上させるか？
• RQ4 マルチスケール特徴統合（EFFM）とプロトタイプベースの推論を統合することで、ノイズが多いまたはコントラストが低い内視鏡画像における性能が向上するか？
• RQ5 臨床現場への導入に適した、ほぼリアルタイムの推論速度を維持しながら、高い精度を維持できるか？

主な発見

• BKAI-IGH OODデータセットにおいて、PrototypeLabはDSC 0.9243、mIoU 0.8744を達成し、16の最先端（SOTA）手法を顕著に上回った。
• Kvasir-SEGデータセットでは、PrototypeLabはDSC 0.9086、mIoU 0.8544、HD 3.71を達成し、SSFormer-LをDSCで0.26%、mIoUで0.44%上回った。
• Kvasir-SEGで学習し、マルチセンターポリープデータセットPolypGenでテストした場合、PrototypeLabはDSC 0.7583、mIoU 0.6957、F2 0.7563を達成し、Polyp-PVTおよびSSFormer-Lを上回った。
• アブレーションスタディの結果、CMGM、PGM、PMGM、EFFM、LKDCブロックを含む完全なアーキテクチャは、ベースラインと比較してDSCが1.15%、mIoUが1.45%向上した。
• PrototypeLabは23.85 FPS（約30 FPS）で画像処理を実行し、高精度を維持したほぼリアルタイムの推論速度を実現した。
• 定性的な結果から、DuAT や Polyp-PVT といったSOTAベースラインと比較して、微小、平坦、ノイズの多いポリープにおける過剰および不十分なセグメンテーションが減少した。

Figure 2 : Qualitative results of models trained BKAI-IGH and tested on PolypGen. It can be observed that PrototypeLab produces a more accurate segmentation map in all the centers from C1 to C6.