[Paper Review] 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 DetectionMedicine3 citations

TL;DR

This paper proposes PrototypeLab, a prototype learning-based segmentation model that enhances robustness for out-of-distribution (OOD) polyp segmentation across diverse colonoscopy imaging modalities (WLI, BLI, LCI, FICE). By leveraging a Coarse Mask Generation Module (CMGM), Prototype Generation Module (PGM), and Prototype Mask Generation Module (PMGM), the model learns class-specific prototypes capturing shape, texture, and color features, achieving a DSC ≥ 90% and mIoU ≥ 85% on OOD datasets with near real-time inference.

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.

Motivation & Objective

• Address the poor generalization of existing polyp segmentation models on out-of-distribution (OOD) datasets from different clinical centers.
• Overcome challenges posed by small, flat, partially visible, or camouflaged polyps with indistinct boundaries and variable imaging conditions.
• Develop a robust segmentation model that generalizes well across diverse imaging modalities (WLI, BLI, LCI, FICE) and multi-center datasets.
• Integrate prototype learning to capture intrinsic data distribution features, improving model generalization and reliability in clinical settings.
• Achieve high accuracy and near real-time inference speed for practical deployment in computer-aided diagnosis (CADx) systems.

Proposed method

• Utilize a Pyramid Vision Transformer (PVT) encoder to extract multi-scale features from input colonoscopy images.
• Implement a Coarse Mask Generation Module (CMGM) to generate initial coarse masks used to guide prototype learning.
• Design a Prototype Generation Module (PGM) that learns multiple prototypes per class by capturing underlying data distribution, enhancing feature representation.
• Introduce an Encoder Feature Fusion Module (EFFM) to enable multi-scale feature integration and improve feature richness.
• Employ a Prototype Mask Generation Module (PMGM) that fuses learned prototypes with decoder features to generate final segmentation masks.
• Train the end-to-end architecture using cross-entropy and Dice loss to optimize segmentation performance on in-distribution and OOD data.

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

Experimental results

Research questions

• RQ1 Can prototype learning improve generalization of polyp segmentation models on out-of-distribution (OOD) datasets from different clinical centers?
• RQ2 How effective is the proposed PrototypeLab in handling challenging polyp types such as diminutive, flat, or partially visible lesions across diverse imaging modalities?
• RQ3 To what extent do the CMGM, PGM, and PMGM components contribute to improved segmentation accuracy and robustness compared to baseline architectures?
• RQ4 Does the integration of multi-scale feature fusion (EFFM) and prototype-based reasoning enhance performance on noisy or low-contrast endoscopic images?
• RQ5 Can the model maintain high accuracy while achieving near real-time inference speed suitable for clinical deployment?

Key findings

• PrototypeLab achieves a DSC of 0.9243 and mIoU of 0.8744 on the BKAI-IGH OOD dataset, significantly outperforming 16 SOTA methods.
• On the Kvasir-SEG dataset, PrototypeLab achieves a DSC of 0.9086, mIoU of 0.8544, and HD of 3.71, surpassing SSFormer-L by 0.26% in DSC and 0.44% in mIoU.
• When trained on Kvasir-SEG and tested on the multi-center PolypGen dataset, PrototypeLab achieves a DSC of 0.7583, mIoU of 0.6957, and F2 of 0.7563, outperforming both Polyp-PVT and SSFormer-L.
• Abllation studies confirm that the full architecture (with CMGM, PGM, PMGM, EFFM, and LKDC blocks) improves DSC by 1.15% and mIoU by 1.45% over the baseline.
• PrototypeLab processes images at 23.85 FPS (≈30 FPS), demonstrating near real-time inference speed while maintaining high accuracy.
• Qualitative results show that PrototypeLab reduces over- and under-segmentation on diminutive, flat, and noisy polyps compared to SOTA baselines like DuAT and Polyp-PVT.

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.