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