Docker demo support

.dockerignore

@@ -0,0 +1,2 @@
+data
+data.orig

.gitignore

@@ -1,4 +1,5 @@
 .ipynb_checkpoints
 .div2k
 .ckpt
+data
 

README.md

@@ -169,7 +169,7 @@ The following training examples use the [training and validation datasets](#div2
 training API is designed around *steps* (= minibatch updates) rather than *epochs* to better match the descriptions in the 
 papers.
 
-## EDSR
+### EDSR
 
 ```python
 from model.edsr import edsr
@@ -204,7 +204,7 @@ trainer.model.save_weights('weights/edsr-16-x4/weights.h5')
 Interrupting training and restarting it again resumes from the latest saved checkpoint. The trained Keras model can be
 accessed with `trainer.model`.
 
-## WDSR
+### WDSR
 
 ```python
 from model.wdsr import wdsr_b
@@ -236,7 +236,7 @@ print(f'PSNR = {psnr.numpy():3f}')
 trainer.model.save_weights('weights/wdsr-b-32-x4/weights.h5')
 ```
 
-## SRGAN
+### SRGAN
 
 ### Generator pre-training
 
@@ -275,7 +275,7 @@ gan_trainer.generator.save_weights('weights/srgan/gan_generator.h5')
 gan_trainer.discriminator.save_weights('weights/srgan/gan_discriminator.h5')
 ```
 
-## SRGAN for fine-tuning EDSR and WDSR models
+### SRGAN for fine-tuning EDSR and WDSR models
 
 It is also possible to fine-tune EDSR and WDSR x4 models with SRGAN. They can be used as drop-in replacement for the
 original SRGAN generator. More details in [this article](https://krasserm.github.io/2019/09/04/super-resolution/).
@@ -299,3 +299,33 @@ generator.load_weights('weights/wdsr-b-16-32/weights.h5')
 gan_trainer = SrganTrainer(generator=generator, discriminator=discriminator())
 gan_trainer.train(train_ds, steps=200000)
 ```
+
+## Docker
+
+You may try out models in docker using the provided Dockerfile.
+
+```bash
+docker build -t super-resolution -f docker/Dockerfile .
+docker run -ti --rm -v "${PWD}":/working super-resolution edsr demo/0829x4-crop.png
+docker run -ti --rm -v "${PWD}":/working super-resolution wdsr demo/0829x4-crop.png
+docker run -ti --rm -v "${PWD}":/working super-resolution srgan demo/0829x4-crop.png
+```
+
+If you wish to use nvidia CUDA in passthrough with nvidia-docker, you may do that as well.
+
+```bash
+docker run -ti --rm -v "${PWD}":/working --gpus all super-resolution edsr demo/0829x4-crop.png
+```
+
+After running, look in the demo folder for the scaled images to compare.
+
+### Training with DIV2K in Docker
+
+You can also use the Docker container to train your own models, and then leverage them for image processing.
+
+```bash
+mkdir -p data/images
+cp demo/*.png data/images
+docker build -t super-resolution -f docker/Dockerfile .
+docker run -ti --rm -v "${PWD}":/working -v "${PWD}/data":/data --gpus all super-resolution div2k_train
+```

docker/Dockerfile

@@ -0,0 +1,45 @@
+FROM ubuntu:20.04
+
+WORKDIR /app
+
+COPY environment.yml /app
+
+RUN \
+    apt-get update &&\
+    DEBIAN_FRONTEND='noninteractive' apt-get install -y python wget nvidia-cuda-dev=10.1.243-3 nvidia-cuda-toolkit=10.1.243-3 &&\
+    wget https://martin-krasser.de/sisr/weights-edsr-16-x4.tar.gz &&\
+    wget https://martin-krasser.de/sisr/weights-wdsr-b-32-x4.tar.gz &&\
+    wget https://martin-krasser.de/sisr/weights-srgan.tar.gz &&\
+    tar xvfz weights-edsr-16-x4.tar.gz &&\
+    tar xvfz weights-wdsr-b-32-x4.tar.gz &&\
+    tar xvfz weights-srgan.tar.gz &&\
+    rm -f weights-edsr-16-x4.tar.gz &&\
+    rm -f weights-wdsr-b-32-x4.tar.gz &&\
+    rm -f weights-srgan.tar.gz &&\
+    wget https://developer.download.nvidia.com/compute/machine-learning/repos/ubuntu1804/x86_64/libcudnn7_7.6.5.32-1+cuda10.1_amd64.deb &&\
+    wget https://developer.download.nvidia.com/compute/machine-learning/repos/ubuntu1804/x86_64/libcudnn7-dev_7.6.5.32-1+cuda10.1_amd64.deb &&\
+    apt-get install -y ./libcudnn7_7.6.5.32-1+cuda10.1_amd64.deb &&\
+    apt-get install -y ./libcudnn7-dev_7.6.5.32-1+cuda10.1_amd64.deb &&\
+    rm -f ./libcudnn7_7.6.5.32-1+cuda10.1_amd64.deb &&\
+    rm -f ./libcudnn7-dev_7.6.5.32-1+cuda10.1_amd64.deb &&\
+    cd /root &&\
+    wget https://repo.anaconda.com/miniconda/Miniconda3-py39_4.9.2-Linux-x86_64.sh &&\
+    bash Miniconda3-py39_4.9.2-Linux-x86_64.sh -b &&\
+    rm -f Miniconda3-py39_4.9.2-Linux-x86_64.sh &&\
+    miniconda3/condabin/conda init bash &&\
+    bash -c " \
+        export PS1='$ ' &&\
+        . /root/.bashrc &&\
+        cd /app &&\
+        conda env create -f environment.yml \
+    " &&\
+    apt-get clean
+
+COPY model/ /app/model/
+COPY *.py /app/
+COPY LICENSE /app
+COPY docker/*.py /app/
+COPY docker/run.sh /app
+
+ENTRYPOINT ["./run.sh"]
+

docker/div2k_train.py

@@ -0,0 +1,61 @@
+import os
+import matplotlib.pyplot as plt
+
+from data import DIV2K
+from model.srgan import generator, discriminator
+from train import SrganTrainer, EdsrTrainer
+
+# %matplotlib inline
+
+# Location of model weights (needed for demo)
+weights_dir = '/data/weights/div2k'
+os.makedirs(weights_dir, exist_ok=True)
+weights_file = lambda filename: os.path.join(weights_dir, filename)
+
+# Set up images
+images_dir = '/data/images'
+os.makedirs(images_dir, exist_ok=True)
+caches_dir = '/data/caches'
+os.makedirs(caches_dir, exist_ok=True)
+div2k_train = DIV2K(
+                    scale=4,
+                    subset='train',
+                    downgrade='bicubic',
+                    images_dir=images_dir,
+                    caches_dir=caches_dir,
+)
+div2k_valid = DIV2K(
+                    scale=4,
+                    subset='valid',
+                    downgrade='bicubic',
+                    images_dir=images_dir,
+                    caches_dir=caches_dir,
+)
+
+# Do pre-training
+check_dir = f'/data/ckpt/pre_generator'
+os.makedirs(check_dir, exist_ok=True)
+train_ds = div2k_train.dataset(batch_size=16, random_transform=True)
+valid_ds = div2k_valid.dataset(batch_size=16, random_transform=True, repeat_count=1)
+pre_trainer = EdsrTrainer(model=generator(), checkpoint_dir=check_dir)
+pre_trainer.train(
+                    train_ds,
+                    valid_ds.take(1),
+                    #steps=1000000, 
+                    steps=1000, 
+                    evaluate_every=1000, 
+                    save_best_only=False
+)
+pre_trainer.model.save_weights(weights_file('pre_generator.h5'))
+
+# Do gan-training
+gan_generator = generator()
+gan_generator.load_weights(weights_file('pre_generator.h5'))
+gan_trainer = SrganTrainer(generator=gan_generator, discriminator=discriminator())
+gan_trainer.train(
+                    train_ds,
+                    #steps=200000
+                    steps=100
+)
+gan_trainer.generator.save_weights(weights_file('gan_generator.h5'))
+gan_trainer.discriminator.save_weights(weights_file('gan_discriminator.h5'))

docker/edsr.py

@@ -0,0 +1,33 @@
+import os
+import sys
+import matplotlib.pyplot as plt
+
+from data import DIV2K
+from model.edsr import edsr
+
+from utils import load_image
+from model import resolve_single
+
+lr_image_path = sys.argv[1]
+
+# Number of residual blocks
+depth = 16
+
+# Super-resolution factor
+scale = 4
+
+weights_dir = f'weights/edsr-{depth}-x{scale}'
+weights_file = os.path.join(weights_dir, 'weights.h5')
+
+model = edsr(scale=scale, num_res_blocks=depth)
+model.load_weights(weights_file)
+
+lr = load_image(lr_image_path)
+sr = resolve_single(model, lr)
+
+fig = plt.figure()
+ax = plt.Axes(fig, [0., 0., 1., 1.])
+ax.set_axis_off()
+fig.add_axes(ax)
+plt.imshow(sr)
+plt.savefig(sys.argv[1] + '_edsr.png', dpi=300, bbox_inches='tight', pad_inches=0)

docker/run.sh

@@ -0,0 +1,10 @@
+#!/bin/bash
+
+export PS1='$ '
+. /root/.bashrc
+
+export PATH=/usr/local/cuda-10.1/bin${PATH:+:${PATH}}
+export LD_LIBRARY_PATH=/usr/local/cuda-10.1/lib64${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}
+
+conda activate sisr
+python ${1}.py "/working/${2}"

docker/srgan.py

@@ -0,0 +1,26 @@
+import os
+import sys
+import matplotlib.pyplot as plt
+
+from data import DIV2K
+from model.srgan import generator
+
+from utils import load_image
+from model import resolve_single
+
+lr_image_path = sys.argv[1]
+
+weights_file = 'weights/srgan/gan_generator.h5'
+
+model = generator()
+model.load_weights(weights_file)
+
+lr = load_image(lr_image_path)
+sr = resolve_single(model, lr)
+
+fig = plt.figure()
+ax = plt.Axes(fig, [0., 0., 1., 1.])
+ax.set_axis_off()
+fig.add_axes(ax)
+plt.imshow(sr)
+plt.savefig(sys.argv[1] + '_srgan.png', dpi=300, bbox_inches='tight', pad_inches=0)

docker/wdsr.py

@@ -0,0 +1,33 @@
+import os
+import sys
+import matplotlib.pyplot as plt
+
+from data import DIV2K
+from model.wdsr import wdsr_b
+
+from utils import load_image
+from model import resolve_single
+
+lr_image_path = sys.argv[1]
+
+# Number of residual blocks
+depth = 32
+
+# Super-resolution factor
+scale = 4
+
+weights_dir = f'weights/wdsr-b-{depth}-x{scale}'
+weights_file = os.path.join(weights_dir, 'weights.h5')
+
+model = wdsr_b(scale=scale, num_res_blocks=depth)
+model.load_weights(weights_file)
+
+lr = load_image(lr_image_path)
+sr = resolve_single(model, lr)
+
+fig = plt.figure()
+ax = plt.Axes(fig, [0., 0., 1., 1.])
+ax.set_axis_off()
+fig.add_axes(ax)
+plt.imshow(sr)
+plt.savefig(sys.argv[1] + '_wdsr.png', dpi=300, bbox_inches='tight', pad_inches=0)