Text-based-denoising

This project implements diffusion based transformer model to denoise the images with text prompt guidance.

Dataset Information

Note: Our project utilizes a custom dataset, distinct from the one used in the original WindowSeat project.

• Dataset Access: Download via Google Drive
• Text Embeddings and trained checkpoints Access: Download via Google Drive
• Overview: To enhance model performance, we reorganized the data with an optimized structure.
• Composition: The dataset combines selected public data with 3,000 synthesized images (including rainstreaks, raindrops, and reflections) to improve robustness.

Dataset Preparation (download + synthesis, pre-split)

1. Download the dataset, unzip it, and place it in the root directory.
2. (Optional) If need to synthesize more data:
   • Check image_synthesizer/readme.txt.
3. Run DatasetManager.py to generate documentation files.
4. Check the dataset_metadata/ folder for:
   • dataset_info and dataset_structure.
   • train.json, valid.json, and test.json for the DataLoader.

Dataset Structure

The image dataset is organzied as follows:

dataset_full
├── blur
│   └── real
│       ├── blended
│       └── clean
├── raindrop
│   ├── real
│   │   ├── blended
│   │   └── clean
│   └── syn_zka
│       ├── blended
│       └── clean
├── rainstreak
│   ├── real
│   │   ├── blended
│   │   └── clean
│   └── syn_zka
│       ├── blended
│       └── clean
├── rainstreak_raindrop
│   └── real
│       ├── blended
│       └── clean
└── reflection
    ├── syn
    │   ├── blended
    │   └── clean
    └── syn_zka
        ├── blended
        └── clean

These images are denoised by DIT, so which is text guided network that takes text embeddings as input, so we create text embeddings from the texttual data and store it, these embeddings are generated for each task, for instance we have text prompts for blur taks that look like the following:

    "De-blur this image to restore sharp details.",
    "Remove the motion blur from the photo.",
    "Sharpen the blurry areas and recover clarity.",
    "Fix the out-of-focus blur in this picture.",
    "Clean the hazy and blurred regions.",
    "Make the blurry image crisp and clear.",
    "Recover the lost details from the blurred scene.",
    "Eliminate the camera shake blur.",
    "Restore the sharp edges of the blurred objects.",
    "Please deblur the foreground and background.",
    "Clear the motion artifacts from the image.",
    "Perform high-quality deblurring on this scene.",
    "Refocus this blurry photograph.",
    "Improve the resolution by removing blur.",
    "Sharp image restoration from blurred input.",
    "Remove defocus blur for a clearer view.",
    "Enhance the clarity by fixing the blur.",
    "Deblur the scene while maintaining natural textures.",
    "Get rid of the blurriness in this frame.",
    "Turn this blurry shot into a sharp one."

And the embeddings are stored before we run the training:

├── blur
│   ├── 0.pt
│   ├── 1.pt
│   ├── 2.pt
│   └── ...
├── raindrop
│   ├── 0.pt
│   ├── 1.pt
│   ├── 2.pt
│   └── ...
├── rainstreak
│   ├── 0.pt
│   ├── 1.pt
│   ├── 2.pt
│   └── ...
├── rainstreak_raindrop
│   ├── 0.pt
│   ├── 1.pt
│   ├── 2.pt
│   └── ...
└── reflection
    ├── 0.pt
    ├── 1.pt
    ├── 2.pt
    └── ...

Installation

Anconda environment recommended.

git clone https://github.com/mfaizan-ai/text-based-denoising
cd text-based-denoising

create a virtual environment in Anaconda and activate it.

conda create -n windowseat_test_env python=3.12.11 -y 
conda activate windowseat

Now install all the dependencies

pip install --upgrade pip
pip install -r requirements.txt

Usage

To train the model, subject a slurm job on the HPC or submit to GPU. Make sure to update paths to train the training successufly.

sbatch run_training.sh

Note: for better reproducibility and viewing on the command line, its NOT recommended to run the slurm script for either test or training. The command that starts with sbatch, they are written for longer training times and test.

to run with the command line, make sure you activate the conda environment and on hte compute node, to get a compute node:

source /lustre/disk/home/users/mfaizan/bash_scripts/slurm_jobs_submission.sh
srun --jobid=$SLURM_JOBID --pty bash

and activate cuda for getting the gpu:

source /lustre/disk/home/users/mfaizan/bash_scripts/setup_everything.sh

after doing this you must cd into the project directory, depending on the project location where you cloned, in my case:

cd /lustre/disk/home/users/mfaizan/windowseat-reflection-removal/text-based-denoising

now actiavate the conda environment for the denoising diffusion, depending on your conda name:

conda activate windowseat_test_env

for accessing the dataset and the text embeddings, set hte project path so we can use that in training and test

PROJECT=/lustre/disk/home/users/mfaizan/windowseat-reflection-removal/text-based-denoising

then you run training on command line:

python train_windowseat.py \
    --data-root    $PROJECT/dataset \
    --meta-dir     $PROJECT/dataset_metadata \
    --embed-dir    $PROJECT/text/text_embeddings \
    --output-dir   runs/combined_test_training \
    --use-task-aware-loss \
     --use-multitask-lora \
    --total-steps  11000 \
    --batch-size   2 \
    --grad-accum   1 \
    --lr-start     1e-5 \
    --lr-peak      1e-4 \
    --lr-decay     5e-6 \
    --warmup-steps 100 \
    --resolution   608 \
    --lambda-psnr  0.1 \
    --lambda-ssim  20.0 \
    --lora-rank    128 \
    --lora-alpha   128 \
    --lora-dropout 0.0 \
    --log-interval  50 \
    --val-interval  500 \
    --save-interval 1000 \
    --num-workers  8 \
    --wandb-project combined_test_training \
    --run-name     combined_test_training_11k

similarly the test can be run with slurm

sbatch run_test.sh

And also a test script can be used on terminal:

python test.py \
      --checkpoint  $PROJECT/runs/baseline_full/checkpoint_best.pt \
      --data-root   $PROJECT/dataset \
      --meta-dir    $PROJECT/dataset_metadata \
      --embed-dir   $PROJECT/text/text_embeddings \
      --output-dir  runs/test_run
 

this will load the check point and performance inference on the test set, the model return the metrics on the test set and the inference over a few images from each category to see how model denoise those categories.

Running demo with Gradio

To launch gradio to run the demo for restoration, its recommended that you clone the repo locally so you can use those for testing with the app: if you can't find then in that computer you can find at

/Users/faizan/Documents/PhD/advance_ai/project/demo_dataset

# Terminal 1 — on laptop, SSH to cluster then gpu01
ssh maguire01

cd /lustre/disk/home/users/mfaizan/windowseat-reflection-removal/text-based-denoising

source ../../bash_scripts/slurm_jobs_submission.sh

srun --jobid=$SLURM_JOBID --pty bash

source ../../bash_scripts/setup_everything.sh

cd /lustre/disk/home/users/mfaizan/windowseat-reflection-removal/text-based-denoising

conda activate windowseat

python app.py \
    --checkpoint runs/baseline_full/checkpoint_best.pt \
    --embed-dir  text/text_embeddings \
    --host 0.0.0.0 \
    --port 7862

# if gradio can't be found please install wiht 
pip install gradio


# Terminal 2 — fresh terminal on laptop (not SSH'd)
ssh -L 7862:gpu01:7862 maguire01

# Browser
http://localhost:7862