TNRD-pytorch
Trainable Nonlinear Reaction Diffusion: A Flexible Framework for Fast and Effective Image Restoration - unofficial pytorch implementation
This folder is based on https://github.com/kligvasser/xUnit/tree/master/denoising and aim to bring back to life the use of TNRD based architechture. Note that TNRD model has very few parameters while achieving high PSNR on several tasks.
Dependecies and installation
The code is based on pytorch 1.0+. Please install depencise using:
python -m pip install -r requirements.txt
Repredue the results
For denosing please use:
python main.py --root path-to-data --g-model g_tnrd --d-model d_tnrd --model-config "{'gen_blocks':3, 'dis_blocks':4, 'in_channels':1}" --reconstruction-weight 1.0 --perceptual-weight 0 --adversarial-weight 0 --tnrd-energy-weight 0 --crop-size 100 --gray-scale --noise-sigma 50 --epochs 600 --step-size 150 --batch-size 2 --eval-every 10 --print-every 100 --lr 1e-3
For super resulotion please use:
python3 main.py --root path-to-data --g-model g_tnrd --d-model d_tnrd --model-config "{'scale':4, 'gen_blocks':3, 'dis_blocks':1}" --scale 4 --reconstruction-weight 1.0 --perceptual-weight 0 --adversarial-weight 0 --crop-size 40 --device-ids 0 --batch-size 1 --eval-every 100 --lr 1e-3 --step-size 200