Skip to content

computational-imaging/diffusion-in-the-dark

Repository files navigation

Diffusion in the Dark

Official repository for Diffusion in the Dark: A Diffuion Model for Low-Light Text Recognition
Cindy M. Nguyen, Eric R. Chan, Alexander W. Bergman, Gordon Wetzstein
WACV 2024

Project Page | Paper

arXiv

Updates

  • 11/06/23: 📣 Training code, inference code, and pretrained models are released.
  • 10/23/23: ✨ Diffusion in the Dark was accepted into WACV!

Setup

You can use the following to set up a Conda environment.
Environment and code was tested on Linux, NVIDIA Driver Version 545.23.06, CUDA Version 12.3. Code was run on a single NVIDIA RTX 3090.

conda env create -f environment.yml
conda activate did
pip install -r requirements.txt
pip install -e .

Download our pretrained model HERE. Put the pretrained model in a runs folder.

We put the datasets into a folder outside of the repo to create the following file structure

data
    LOL
    lowLightDataset
diffusion-in-the-dark

LOL (LOw-Light) dataset from here

File structure should be as follows:

LOL
    our485
        high
        low
    eval15
        high
        low

Low Light Dataset (based off of Seeing in the Dark)

From Xu, Ke, et al. "Learning to restore low-light images via decomposition-and-enhancement." This can be downloaded from here.
File structure should be as follows:

lowLightDataset
    gt
        test
        train
    input
        test
        train

Usage

To train on LowLightDataset or LOL dataset:

python train.py --outdir=runs --data=../data/lowLightDataset --batch 1 --curve linear --dataset lowlight --add_noise True --scale_norm True --use_lpips True

python train.py --outdir=runs --data=../data/LOL --batch 1 --curve linear --dataset lol --add_noise True --scale_norm True --use_lpips True

If you want to train on a custom dataset use data/raw_img_lol.ipynb as a base to sample 30 random images to find the mean and std of your data to perform data normalization with. Add the mean and std to training/constants.py

To test:

python inference.py --dataset lol

Reference

If you find this work useful, please consider citing us!

@inproceedings{nguyen2024diffusion,
  author    = {Nguyen, Cindy M and Chan, Eric R and Bergman, Alexander W and Wetzstein, Gordon},
  title     = {Diffusion in the Dark: A Diffusion Model for Low-Light Text Recognition},
  journal   = {Proceedings of the IEEE/CVF Winter Conference on Applications of Computer Vision},
  year      = {2024},
}

Acknowledgements

We thank the authors of EDM from which our repo is based off of.

About

Repository for Diffusion in the Dark (WACV 2024)

Resources

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published