Enhanced EDSR

Enhanced EDSR is one of the AI models included in the Intel® Library for Video Super Resolution. It is designed to upscale images and video to delive a high quality visual output. The model is trained using Basic Super Restoration (BasicSR) toolbox, a pytorch based framework that streamlines the development and training of the image/video resoration models. This repository provides guidance on how to train the model and include tools to convert the trained pytorch model to OpenVINO™ using OpenVINO inference API, enabling superior throughput and reduced latency. The training is supported on both CPU and Nvidia GPU.

For integration into media workflows, please refer to the iVSR implementation https://github.com/OpenVisualCloud/iVSR, which offers an FFmpeg plugin integrated with the dnn_processing filter.

Pre-trained model files are available upon request. Please contact us.

Environment Setup

Before beginning the process of training EDSR, first time users must install the required dependencies. The provided prepare_env.sh script automates the environment setup and ensures that all necessary software packages are properly installed including

• python 3.12
• BasicSR python dependencies
• FFmpeg (optional for dataset preparation)

A python virtual environment by the name of venv_edsr will be created in your project folder. It is recommended to activate this virtual environment before proceeding with training Enhanced EDSR.

bash ./prepare_env.sh
source ./venv_edsr/bin/activate

Dataset Preparation

• Put the HR(high resolution) data under dataset/your_dataset/HR like:

dataset
└── your_dataset
    ├── HR
    │   ├── 89-00001.png
    │   ├── 90-00001.png
    │   ├── 91-00001.png
    │   └── ......
    └── meta_info_your_dataset.txt

You can use make_dataset.sh to extract HR frames from videos. The script asks for a single folder containing all videos you want converted into images. Specifying the # of images generated per video and specifying the output folder where the images are saved are optional arguments

Usage: bash ./make_dataset.sh input_video_folder [output_image_folder num_frames_per_video]

• The EDSR (Enhanced Deep Super-Resolution) model does not require manually generating low-resolution (LR) images before training. Instead, the LR data is created on the fly by the training pipeline using a defined degradation algorithm. The key to controlling the final video quality produced by the model lies in the design of two components: The modified Ground-Truth (GT) generation algorithm, which defines how the high-resolution (HR) reference images , and The degradation model, which specifies how LR images are derived from those GT frames during training.

• If using BasicSR for training as suggested, it is required to generate a text file summarizing basic information about your HR folder. Anytime you update or change your HR images, please regenerate this metadata.

Before using BasicSR, please add ./BasicSR path to PYTHONPATH export PYTHONPATH="./BasicSR"

cd BasicSR
python3 scripts/data_preparation/generate_meta_info.py /path_to/HR
cd ..

EDSR model training

This repository includes both the PyTorch implementation from https://github.com/sanghyun-son/EDSR-PyTorch and the BasicSR toolbox for training the model. However, we recommend using the BasicSR toolbox, and the steps outlined below are based on that approach.

Train EDSR model based on BasicSR toolbox

This project relies on BasicSR for training Enhanced EDSR. BasicSR toolbox support EDSR model training and implement high order degradation process used by Real-ESRGAN model to simulate real wold scenarios, enabled for the Enhanced EDSR model training.

Dependencies

pip3 install -r ./BasicSR/requirements.txt

How to train EDSR model

add ./BasicSR path to PYTHONPATH export PYTHONPATH="./BasicSR"

BasicSR toolbox uses yaml files for configuration of training or test. The document Config.md introduce the detailed.

The configuration file template for Enhanced EDSR training is located in ./BasicSR/options/train/EDSR/. Please set the dataset and metainfo path, along with any other parameters for training configuration.

Here is the commandline for training training EDSR with 3 ResidualBlock and 64 features number for better quality

python3 ./BasicSR/basicsr/train.py -opt ./BasicSR/options/train/EDSR/train_EDSR_Lx2_r3_f64_pixabay.yml

training EDSR with 3 ResidualBlock and 32 features number for higher speed

python3 ./BasicSR/basicsr/train.py -opt ./BasicSR/options/train/EDSR/train_EDSR_Lx2_r3_f32_pixabay.yml

How to test EDSR model

Test EDSR model without metrics

The inference_edsr.py file is the inference sample of Enhanced EDSR pytorch model, can use the sample to do inference only. Here is the usage of the sample.

python3 ./BasicSR/inference/inference_edsr.py --model_path ./BasicSR/experiments/EDSR_Lx2_r3_f32/models/net_g_latest.pth --block_num 3 --feat_num 32 --input xxx_LR_folder --output xxx_output_folder

Test EDSR model with metrics

add ./BasicSR path to PYTHONPATH export PYTHONPATH="./BasicSR"

BasicSR toolbox support psnr, ssim, lpips and niqe metrics calculation. The configuration file template for Enhanced EDSR test is located in ./BasicSR/options/test/EDSR/. Please set correct parameters values in test configuration file before start testing.

Here is the commandline for test test EDSR with 3 ResidualBlock and 64 features number

python3 ./BasicSR/basicsr/test.py -opt ./BasicSR/options/test/EDSR/test_EDSR_Lx2_r3_f64.yml

test EDSR with 3 ResidualBlock and 32 features number

python3 ./BasicSR/basicsr/test.py -opt ./BasicSR/options/test/EDSR/test_EDSR_Lx2_r3_f32.yml

Convert PyTorch model to OpenVINO IR

Install ONNX and OpenVINO via pip3 install onnx openvino then execute the following command to convert PyTorch->ONNX->OpenVINO . This step is optional to convert the model from Pytorch to the OpenVINO IR format for faster speed on x86 architecture and Intel GPUs.

python3 ./BasicSR/scripts/convert_pytorch_to_onnx_ov.py --py_model_path ./BasicSR/experiments/EDSR_Lx2_r3_f32/models/net_g_latest.pth  --block_num 3 --feat_num 32  --output_model_path ./openvino_IR/ --bit [8 | 10 | 16]