Implementation of a VAE-based image restoration technique for degraded photos that learns from unpaired clean and noisy images. Improves on the original method (Wan et al. 2022) by replacing transposed convolution layers with resize-convolution blocks to eliminate checkerboard artifacts.
Based on the paper:
WAN, Ziyu, ZHANG, Bo, CHEN, Dongdong, et al. "Old photo restoration via deep latent space translation." IEEE Transactions on Pattern Analysis and Machine Intelligence, 2022. arXiv:2009.07047
src/train.py- Script for training different modules/modelssrc/inference.py- Script for denoising images using a trained modelsrc/evaluate.py- Functions for model evaluation (can be modified for custom evaluation)src/models.py- Implementation of model classes (VAE1, VAE2, and Mapping)src/networks.py- Shared network structuressrc/dataset.py- Dataset builders and dataloaderssrc/utils.py- Image processing and evaluation utilitiessrc/config.py- Configuration parameters for training and evaluation
The code has been tested with Python 3.8.0. Install dependencies with:
pip install -r requirements.txtThe model consists of three modules that must be trained separately:
- VAE1 and VAE2 can be trained independently
- The mapping module requires pre-trained VAE1 and VAE2 models
Training each module:
# Train VAE1
python src/train.py --cfg-path configs/vae.yaml --stage vae1 --output-path checkpoints/vae1.ckpt
# Train VAE2
python src/train.py --cfg-path configs/vae.yaml --stage vae2 --output-path checkpoints/vae2.ckpt
# Train mapping (requires trained VAE1 and VAE2)
python src/train.py --cfg-path configs/mapping.yaml --stage mapping --output-path checkpoints/full.ckptMonitor training progress with TensorBoard:
tensorboard --logdir lightning_logsRun inference on images using a trained model:
python src/inference.py -i <input_folder> -o <output_folder> -m <fullmodel_path>Note: The model is not optimized for processing large images and may result in out-of-memory errors with high-resolution inputs.