GGDPolyp

Geometry-Guided Diffusion Improves AI-Based Polyp Detection, Classification, and Re-Identification

This repository provides training and inference code for generating temporally coherent colonoscopy polyp videos with geometry-guided diffusion.

Overview

The pipeline contains three main parts:

1. Training: train the video diffusion on 15-frame polyp video clips.
2. Short-video generation: generate a temporally coherent 15-frame polyp video from one input video and mask.
3. Long-video generation: recursively synthesize long videos with overlapping 15-frame windows.

Data Preparation

Prepare the processed LDPolypVideo dataset under the repository root:

LDPolypVideo/
├── videos/
│   ├── 47_0046.mp4
│   ├── ...
├── train.csv
├── Labelled/
│   └── Images/
└── Masks/

The processed videos should be placed in:

./LDPolypVideo/videos/

The training CSV file should be placed at:

./LDPolypVideo/train.csv

In train.csv, each .jpg entry corresponds to one processed video. The .jpg filename indicates the starting image/frame index of the corresponding video clip.

All processed videos used by this project are 15-frame clips.

Model Weights

Two model weight folders are required before inference.

1. Base inpainting model

Place the base inpainting model under:

./inpaint/

This folder is used as the pretrained inpainting backbone.

2. Trained video diffusion / motion module

Place the trained model weights under:

./out/unet/

Expected checkpoint example:

./out/unet/diffusion_pytorch_model.safetensors

Pretrained model weights and prepared resources can be downloaded from Google Drive.

Environment Setup

Create and activate the environment, then install dependencies:

pip install -r requirements.txt

The provided training script assumes a conda environment named GGDPolyp. If your environment name is different, update scripts/train.sh accordingly.

Training

The training script is:

./scripts/train.sh

Run training from the repository root:

bash scripts/train.sh

Video Synthesis / Testing

The video testing script is:

./inverse.py

Example usage:

python inverse.py \
  --video_path ./LDPolypVideo/videos/47_0046.mp4 \
  --mask_path ./saved_videos/user_mask.png

The script uses CoTracker to propagate the mask region, performs DDIM inversion, and generates a temporally coherent polyp video.

Long Video Generation

The long-video generation script is:

./generating_longvideo.py

Example usage:

python generating_longvideo.py \
  --image_dir ./LDPolypVideo/Labelled/Images/54 \
  --start_frame 1 \
  --end_frame 51 \
  --window_size 15 \
  --mask_path ./saved_videos/user_mask.png

This script recursively processes long image sequences with overlapping 15-frame windows.

The final generated video is saved to:

./saved_videos/long_video.mp4

Intermediate warped videos are saved to:

./saved_videos/warped_videos/