Super-resolution using deep neural networks (U-Net / RUNet)

Python 3.9.5 | Conda package manager

Implementation of U-Net and RUNet architecture for super-resolution task with a focus on comics images.

Quick start

Prepare enviroment:

conda create -n BP-py395 python=3.9.5
conda activate BP-py395

Install dependencies:

pip install -r requirements.txt

Usage:

1. Train -> ./src/train.py --config "config_path"
2. Validation -> ./src/valid.py --config "config_path"
3. Use model -> ./src/main.py -h

Traning

Script used to train deep neural network on given dataset.

> python ./src/train.py -h 

usage: train.py [-h] --config CONFIG

optional arguments:
  -h, --help       show this help message and exit
  --config CONFIG  config file path

Structure of config file for training

{
	"model":"UNET",
	"loss":"MSE",
	"loss_layers": [0,1,2,3],
	
	"epochs":300,
	"upscale_factor":2,
	"crop_size":128,
	"batch_size":16,
	"lr":0.001,
	
	"dataset_train":"./datasets/comics_train",
	"dataset_valid":"./datasets/comics_valid",
	
	"save_path":"./",
	"save_name":"unet"
}

Requried:

• model -> model used for traning -> available values: UNET, RUNET
• loss -> loss function used for traning -> available values: MSE, PERCEPTUAL (Mean Square Error or Perceptual Loss fucntion)
• epochs -> number of train epochs
• upscale_factor -> upscale factor to train with
• crop_size -> crop size of dataset images feeded to model
• batch_size -> size of batch
• lr -> starting learning rate for model
• dataset_train -> path to folder with dataset train images
• dataset_valid -> path to folder with dataset validation images
• save_path -> path where model should be saved
• save_name -> name as model should be saved

Optional:

• loss_layers -> used to determine number of extracted blocks from loss network used for perceptual loss
• debug -> turn on debug mode. At the end of the training, validation and training loss is saved to the cwd directory in numpy format

Validation

Script used to validate trained model.

> python ./src/valid.py -h

usage: valid.py [-h] --config CONFIG

optional arguments:
  -h, --help       show this help message and exit
  --config CONFIG   onfig file path

Structure of config file for validation

{
	"model":"unet_mse.pt",
	"upscale_factor":2,
	"crop_size":1024,
	"dataset_valid":"./datasets/comics_valid",
}

Requried:

• model -> path to file of saved model
• upscale_factor -> upscale factor to valid with
• crop_size -> crop size of dataset images feeded to model
• dataset_valid -> path to folder with dataset validation images

Optional:

• loss_layers -> used to determine number of extracted blocks from loss network used for perceptual loss
• debug -> turn on debug mode. Save high-resolutioin,bilinear,output images durning the traning to cwd.

Use model

Script to use trained model.

> python ./src/main.py -h

usage: main.py [-h] --image IMAGE --model MODEL [--upscale UPSCALE] [--device DEVICE]

optional arguments:
  -h, --help         show this help message and exit
  --image IMAGE      path to image file
  --model MODEL      path to model file
  --upscale UPSCALE  upscale factor
  --device DEVICE    device CPU or CUDA

Examples: -> upscale by factor 2 default value ./src/main.py --image "./images/image1.jpg" --model "model2x.pt"

-> upscale by factor 4 ./src/main.py --image "./images/image2.jpg" --model "model4x.pt" --upscale 4

-> upscale by factor 4 using cuda ./src/main.py --image "./images/image2.jpg" --model "model4x.pt" --upscale 4 --device "cuda"