This code is intended for training Super Resolution (SR) algorithms in Pytorch.
Typologies are available for training at this point:
- Single image super resolution network based on SRResNet architecture ("Photo-Realistic Single Image Super-Resolution Using a Generative Adversarial Network") but with reduced number of channels and depthwise convolution in the decoder
- Attention-Based single image super resolution network with reduced number of channels and changes in the network architecture
- Tiny model to upsampling scanned text images
The PSNR values were calculated with Y channel from YCrCb image.
Model | Set5, PSNRx3, dB | Set5, PSNRx4, dB |
---|---|---|
SmallModel | 33.15 | 31.16 |
- Ubuntu* 16.04 or higher
- OpenVINO™ 2019 R1 or higher
- Python* 3
-
Create virtual environment:
virtualenv venv -p python3 --prompt="(sr)"
-
Activate virtual environment and setup OpenVINO™ variables:
. venv/bin/activate . /opt/intel/openvino/bin/setupvars.sh
TIP: Good practice is adding
. /opt/intel/openvino/bin/setupvars.sh
to the end of thevenv/bin/activate
.echo ". /opt/intel/openvino/bin/setupvars.sh" >> venv/bin/activate
-
Install the module
pip3 install -e .
Create two directories for train and test images. Train images may have any resolution higher than path_size
.
Validation images should have the resolution like the path_size
.
./data
├── train
│ ├── 000000.png
│ ...
└── val
├── 000000.png
...
Use the tools/train.py
script to start the training process:
python3 tools/train.py --config configs/smallmodel_scale4.yaml
To start from a pretrained checkpoint, set init_checkpoint
in config.
Use the tools/test.py
script to evaluate the trained model.
python3 tools/test.py --test_data_path PATH_TO_TEST_DATA \
--models_path PATH_TO_MODELS_PATH \
--exp_name EXPERIMENT_NAME
python3 tools/export.py --models_path PATH_TO_MODELS_PATH \
--exp_name EXPERIMENT_NAME \
--input_size 200 200 \
--data_type FP32
python3 tools/infer.py --model PATH_TO_CHECKPOINT \
--scale 4 \
image_path
python3 tools/infer_ie.py --model <PATH_TO_IR_XML> \
image_path
A network cannot be reshaped after conversation to the IR. Set the input_size
when running tools/export.py
.