Skip to content

Polyp-Gen: Realistic and Diverse Polyp Image Generation for Endoscopic Dataset Expansion

Notifications You must be signed in to change notification settings

CUHK-AIM-Group/Polyp-Gen

Folders and files

NameName
Last commit message
Last commit date

Latest commit

579d772 · Sep 16, 2024

History

4 Commits
Sep 16, 2024
Sep 6, 2024
Sep 6, 2024
Sep 16, 2024
Sep 6, 2024
Sep 6, 2024
Sep 6, 2024
Sep 6, 2024
Sep 16, 2024
Sep 6, 2024
Sep 16, 2024
Sep 6, 2024
Sep 6, 2024
Sep 6, 2024
Sep 6, 2024

Repository files navigation

Polyp-Gen: Realistic and Diverse Polyp Image Generation for Endoscopic Dataset Expansion

introduction

Setup

git clone https://github.com/Saint-lsy/Polyp-Gen.git
cd Polyp-Gen
conda create -n PolypGen python=3.10
conda activate PolypGen
pip install -r requirements.txt

Data Preparation

This model was trained by LDPolypVideo dataset.

We filtered out some low-quality images with blurry, reflective, and ghosting effects, and finally select 55,883 samples including 29,640 polyp frames and 26,243 non-polyp frames.

Our dataset can be downloaded at here.

Training

The pre-trained model is Stable Diffusion Inpainting-2, availble on Huggingface

You can train your own model using the script:

bash scripts/train.sh

Sampling

sampling

Checkpoint

You can download the chekpoints of our Polyp_Gen from here.

Sampling with Specified Mask

python sample_one_image.py

Sampling with Mask Proposer

The weight of pretrained DINOv2 can be found here.

The first step is building database and Global Retrieval.

python GlobalRetrieval.py --data_path /path/of/non-polyp/images --database_path /path/to/build/database --image_path /path/of/query/image/

The second step is Local Matching for query image.

python LocalMatching.py --ref_image /path/ref/image --ref_mask /path/ref/mask --query_image /path/query/image --mask_proposal /path/to/save/mask

One Demo of LocalMatching

python LocalMatching.py --ref_image demos/img_1513_neg.jpg --ref_mask  demos/mask_1513.jpg --query_image  demos/img_1592_neg.jpg --mask_proposal gen_mask.jpg

local_matching

The third step is using the generated Mask to sample.

Acknowledgements

The code is based on the following projects. Greatly thanks to these authors!

About

Polyp-Gen: Realistic and Diverse Polyp Image Generation for Endoscopic Dataset Expansion

Resources

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published