Authors: Xinzhen Ren, Wenju Zhou, Naitong Yuan, Fang Li, Yetian Ruan, and Huiyu Zhou.
Accurate judgment and identification of polyp size is crucial in endoscopic diagnosis. However, the indistinct boundaries of polyps lead to missegmentation and missed cancer diagnoses. In this paper, a prompt-based polyp segmentation method (PPSM) is proposed to assist in early-stage cancer diagnosis during endoscopy. It combines endoscopists' experience and artificial intelligence technology. Firstly, a prompt-based polyp segmentation network (PPSN) is presented, which contains the prompt encoding module (PEM), the feature extraction encoding module (FEEM), and the mask decoding module (MDM). The PEM encodes prompts to guide the FEEM for feature extracting and the MDM for mask generating. So that PPSN can segment polyps efficiently. Secondly, endoscopists' ocular attention data (gazes) are used as prompts, which can enhance PPSN's accuracy for segmenting polyps and obtain prompt data effectively in real-world. To reinforce the PPSN's stability, non-uniform dot matrix prompts are generated to compensate for frame loss during the eye-tracking. Moreover, a data augmentation method based on the segment anything model (SAM) is introduced to enrich the prompt dataset and improve the PPSN's adaptability. Experiments demonstrate the PPSM's superior accuracy and real-time capability. The results from cross-training and cross-testing on four datasets show the PPSM's excellent generalization.
The images and performance will be showed after the paper is accepted.
- Create and activate virtual environment:
python3 -m venv ~/PPSN-env
source ~/PPSN-env/bin/activate
- Clone the repository and navigate to new directory:
git clone https://github.com
cd ./PPSN
- Install the requirements:
pip install -r requirements.txt
You will need to download the dataset yourself
- Download and extract the Kvasir-SEG, the CVC-ClinicDB, the SUN-SEG and the PolyGen datasets.
The prompts can be generated in
drawPrompts.py
The SAM_Mask, SAM_Points, SAM_Edges, SAM_Polygons, and SAM_Scribbles generated by SAM. We generated the them base on ViT-B SAM model.
A dataset should have: original image, Ground Truth, and Prompts (e.g. Edges, Points....).
The types of prompts in the paper are: Edges, Points, Polygons, Scribbles and SAM_Mask. The number of input channels for PEM should change depending on the number of types of prompts.
├─Data
│ └─__pycache__
├─Metrics
│ └─__pycache__
├─Models
│ └─__pycache__
├─Others
│ ├─CFANet-main
│ │ ├─imgs
│ │ ├─lib
│ │ ├─mindspore
│ │ │ ├─lib
│ │ │ └─utils
│ │ └─utils
│ │ └─__pycache__
│ ├─Enhanced-U-Net-main
│ │ ├─datasets
│ │ │ └─__pycache__
│ │ ├─models
│ │ │ └─__pycache__
│ │ └─utils
│ │ └─__pycache__
│ ├─FAPNet-main
│ │ ├─imgs
│ │ ├─lib
│ │ │ └─__pycache__
│ │ └─utils
│ ├─FSAD-Net-master
│ ├─Graph_Region_Boudnary-main
│ │ ├─lib
│ │ │ └─__pycache__
│ │ ├─model
│ │ ├─res_weight
│ │ └─utils
│ │ └─__pycache__
│ ├─HarDNet-MSEG-master
│ │ ├─eval
│ │ ├─imgs
│ │ ├─lib
│ │ └─utils
│ │ └─__pycache__
│ ├─MSNet-M2SNet-main
│ │ ├─image
│ │ ├─model
│ │ │ └─__pycache__
│ │ ├─utils
│ │ └─weight
│ ├─PraNet-master
│ │ ├─eval
│ │ ├─imgs
│ │ ├─jittor
│ │ │ ├─lib
│ │ │ └─utils
│ │ ├─lib
│ │ └─utils
│ │ └─__pycache__
│ └─VPS-main
│ ├─assets
│ ├─docs
│ ├─eval
│ │ └─eval-result
│ │ ├─2015-MICCAI-UNet
│ │ ├─2018-TMI-UNet++
│ │ ├─2019-TPAMI-COSNet
│ │ ├─2020-AAAI-PCSA
│ │ ├─2020-MICCAI-23DCNN
│ │ ├─2020-MICCAI-ACSNet
│ │ ├─2020-MICCAI-PraNet
│ │ ├─2020-TIP-MATNet
│ │ ├─2021-ICCV-DCFNet
│ │ ├─2021-ICCV-FSNet
│ │ ├─2021-MICCAI-PNSNet
│ │ ├─2021-MICCAI-SANet
│ │ ├─2021-NIPS-AMD
│ │ └─2022-TMI-PNSPlus
│ ├─lib
│ │ ├─dataloader
│ │ ├─module
│ │ │ └─PNS
│ │ │ └─PNS_Module
│ │ └─utils
│ ├─scripts
│ ├─snapshot
│ ├─utils
│ └─数据集获取
└─Predictions
├─EUNet
│ ├─Trained on CVC
│ │ ├─test_on_CVC
│ │ ├─test_on_Kvasir
│ │ ├─test_on_polyGen
│ │ └─test_on_SUN
│ ├─Trained on Kvasir
│ │ ├─test_on_CVC
│ │ ├─test_on_Kvasir
│ │ ├─test_on_polyGen
│ │ └─test_on_SUN
│ ├─Trained on polyGen
│ │ ├─test_on_CVC
│ │ ├─test_on_Kvasir
│ │ ├─test_on_polyGen
│ │ └─test_on_SUN
│ └─Trained on SUN
│ ├─test_on_CVC
│ ├─test_on_Kvasir
│ ├─test_on_polyGen
│ └─test_on_SUN
├─FCBFormer
│ ├─Trained on CVC
│ │ ├─test_on_CVC
│ │ ├─test_on_Kvasir
│ │ ├─test_on_polyGen
│ │ └─test_on_SUN
│ ├─Trained on Kvasir
│ │ ├─test_on_CVC
│ │ ├─test_on_Kvasir
│ │ ├─test_on_polyGen
│ │ └─test_on_SUN
│ ├─Trained on polyGen
│ │ ├─test_on_CVC
│ │ ├─test_on_Kvasir
│ │ ├─test_on_polyGen
│ │ └─test_on_SUN
│ └─Trained on SUN
│ ├─test_on_CVC
│ ├─test_on_Kvasir
│ ├─test_on_polyGen
│ └─test_on_SUN
├─HarDMSEG
│ ├─Trained on CVC
│ │ ├─test_on_CVC
│ │ ├─test_on_Kvasir
│ │ ├─test_on_polyGen
│ │ └─test_on_SUN
│ ├─Trained on Kvasir
│ │ ├─test_on_CVC
│ │ ├─test_on_Kvasir
│ │ ├─test_on_polyGen
│ │ └─test_on_SUN
│ ├─Trained on polyGen
│ │ ├─test_on_CVC
│ │ ├─test_on_Kvasir
│ │ ├─test_on_polyGen
│ │ └─test_on_SUN
│ └─Trained on SUN
│ ├─test_on_CVC
│ ├─test_on_Kvasir
│ ├─test_on_polyGen
│ └─test_on_SUN
├─MSNet
│ ├─Trained on CVC
│ │ ├─test_on_CVC
│ │ ├─test_on_Kvasir
│ │ ├─test_on_polyGen
│ │ └─test_on_SUN
│ ├─Trained on Kvasir
│ │ ├─test_on_CVC
│ │ ├─test_on_Kvasir
│ │ ├─test_on_polyGen
│ │ └─test_on_SUN
│ ├─Trained on polyGen
│ │ ├─test_on_CVC
│ │ ├─test_on_Kvasir
│ │ ├─test_on_polyGen
│ │ └─test_on_SUN
│ └─Trained on SUN
│ ├─test_on_CVC
│ ├─test_on_Kvasir
│ ├─test_on_polyGen
│ └─test_on_SUN
├─ParNet
│ ├─Trained on CVC
│ │ ├─test_on_CVC
│ │ ├─test_on_Kvasir
│ │ ├─test_on_polyGen
│ │ └─test_on_SUN
│ ├─Trained on Kvasir
│ │ ├─test_on_CVC
│ │ ├─test_on_Kvasir
│ │ ├─test_on_polyGen
│ │ └─test_on_SUN
│ ├─Trained on polyGen
│ │ ├─test_on_CVC
│ │ ├─test_on_Kvasir
│ │ ├─test_on_polyGen
│ │ └─test_on_SUN
│ └─Trained on SUN
│ ├─test_on_CVC
│ ├─test_on_Kvasir
│ ├─test_on_polyGen
│ └─test_on_SUN
├─PPSN
│ ├─Trained on CVC
│ │ └─p1.0
│ │ ├─test_on_CVC
│ │ ├─test_on_Kvasir
│ │ ├─test_on_polyGen
│ │ └─test_on_SUN
│ ├─Trained on Kvasir
│ │ └─p1.0
│ │ ├─test_on_CVC
│ │ ├─test_on_Kvasir
│ │ ├─test_on_polyGen
│ │ └─test_on_SUN
│ ├─Trained on polyGen
│ │ └─p1.0
│ │ ├─test_on_CVC
│ │ ├─test_on_Kvasir
│ │ ├─test_on_polyGen
│ │ └─test_on_SUN
│ └─Trained on SUN
│ └─p1.0
│ ├─test_on_CVC
│ ├─test_on_Kvasir
│ ├─test_on_polyGen
│ └─test_on_SUN
├─U2Net
│ ├─Trained on CVC
│ │ ├─test_on_CVC
│ │ ├─test_on_Kvasir
│ │ ├─test_on_polyGen
│ │ └─test_on_SUN
│ ├─Trained on Kvasir
│ │ ├─test_on_CVC
│ │ ├─test_on_Kvasir
│ │ ├─test_on_polyGen
│ │ └─test_on_SUN
│ ├─Trained on polyGen
│ │ ├─test_on_CVC
│ │ ├─test_on_Kvasir
│ │ ├─test_on_polyGen
│ │ └─test_on_SUN
│ └─Trained on SUN
│ ├─test_on_CVC
│ ├─test_on_Kvasir
│ ├─test_on_polyGen
│ └─test_on_SUN
├─UNet
│ ├─Trained on CVC
│ │ ├─test_on_CVC
│ │ ├─test_on_Kvasir
│ │ ├─test_on_polyGen
│ │ └─test_on_SUN
│ ├─Trained on Kvasir
│ │ ├─test_on_CVC
│ │ ├─test_on_Kvasir
│ │ ├─test_on_polyGen
│ │ └─test_on_SUN
│ ├─Trained on polyGen
│ │ ├─test_on_CVC
│ │ ├─test_on_Kvasir
│ │ ├─test_on_polyGen
│ │ └─test_on_SUN
│ └─Trained on SUN
│ ├─test_on_CVC
│ ├─test_on_Kvasir
│ ├─test_on_polyGen
│ └─test_on_SUN
└─UNet++
├─test_on_CVC
└─test_on_Kvasir`
Train PPSN on the train split of a dataset:
python train.py --dataset=[train data] --data-root=[path]
- Replace
[train data]
with training dataset name (options:Kvasir
;CVC
). - Replace
[path]
with path to parent directory of/images
and/masks
directories (training on Kvasir-SEG); or parent directory of/Original
and/Ground Truth
directories (training on CVC-ClinicDB). - To train on multiple GPUs, include
--multi-gpu=true
.
Generate predictions from a trained model for a test split. Note, the test split can be from a different dataset to the train split:
python predict.py --train-dataset=[train data] --test-dataset=[test data] --data-root=[path]
- Replace
[train data]
with training dataset name (options:Kvasir
;CVC
). - Replace
[test data]
with testing dataset name (options:Kvasir
;CVC
). - Replace
[path]
with path to parent directory of/images
and/masks
directories (testing on Kvasir-SEG); or parent directory of/Original
and/Ground Truth
directories (testing on CVC-ClinicDB).
The intellectual property rights of the disposable electronic endoscopes, image processor, and eye tracker were transferred to a Shanghai Medical Device Co., Ltd. Some of the code is not publicly available. Please also do not use the exposed code for any commercial purposes, thanks.
No commercial use is permitted
This work was supported by the Shanghai Science and Technology Development Foundation (No. 21S31905100), the Shanghai Municipal Health Commission Foundation (No. 202150001), and the 111 Project (No. D18003).