数据集获取

git clone https://github.com/CUHK-AIM-Group/Polyp-Gen.git

curl -L -o repo.zip https://github.com/CUHK-AIM-Group/Polyp-Gen/archive/refs/heads/main.zip
unzip repo.zip

数据集大小约 4.05 GB。

源站 README 摘录（使用方式）

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

📖 Paper ｜🤗 Model｜🤗 Dataset

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.
[02/26] We update the download link of the training and test dataset at HuggingFace link

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

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

<img src=“assets/matching.jpg” alt=“local_matching” width=“450”>

The third step is using the generated Mask to sample.

Acknowledgements

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

• [Diffusers](https://github.com/

数据加载示例（图像类）

from PIL import Image
import glob, os

files = (glob.glob(os.path.join(path, "**", "*.png"), recursive=True)
       + glob.glob(os.path.join(path, "**", "*.jpg"), recursive=True)
       + glob.glob(os.path.join(path, "**", "*.tif"), recursive=True))
print("图像文件数:", len(files))
img = Image.open(files[0]); print("尺寸/模式:", img.size, img.mode)

# torchvision Dataset 方式:
# from torchvision import datasets
# ds = datasets.ImageFolder(path)  # 要求 子目录=类别

目录组织与标注格式以源站说明和下载后实际文件为准。

完整仓库：github.com/CUHK-AIM-Group/Polyp-Gen