A Pytorch implemented Deep SR-ITM (ICCV2019 oral)
Soo Ye Kim, Jihyong Oh, Munchurl Kim. Deep SR-ITM: Joint Learning of Super-Resolution and Inverse Tone-Mapping for 4K UHD HDR Applications. IEEE International Conference on Computer Vision, 2019.
Note: The code is completely adapted from https://github.com/sooyekim/Deep-SR-ITM
but rewritten in pytorch format. This repository is NOT aimed to improve the baseline, but to retain the original settings in a different implementation. If you have any questions for the details of the implementations, please refer to the original repo.
- Ubuntu 16.04 LTS
- python 3.7.5
- pytorch 1.3.1
- torchvision 0.4.2
- CUDA 10.1
- opencv 3.4.2
- numpy 1.17.4
- Download training and testing data from
https://github.com/sooyekim/Deep-SR-ITM
- Use Matlab to transform the data with extension,
.mat
, into'.png'
form (No matter SDR or HDR images) - Prepare the data as following...
${DATA_ROOT}
├── trainset_SDR
│ ├── 000001.png
│ ├── 000002.png
│ ├── ...
│ └── 039840.png
├── trainset_HDR
│ ├── 000001.png
│ ├── 000002.png
│ ├── ...
│ └── 039840.png
├── testset_SDR
│ ├── 000001.png
│ ├── 000002.png
│ ├── ...
│ └── 000028.png
└── testset_SDR
├── 000001.png
├── 000002.png
├── ...
└── 000028.png
# Prepare CUDA Installation
...
# git clone repository
git clone https://github.com/greatwallet/Pytorch-Implemented-Deep-SR-ITM.git
cd Pytorch-Implemented-Deep-SR-ITM
# create conda environment
conda create --n env-sr-itm python=3.7 -y
conda activate env-sr-itm
conda install -c pytorch pytorch -y
conda install -c pytorch torchvision -y
conda install -c conda-forge opencv -y
conda install numpy -y
# set soft link to data path
ln -s ${DATA_ROOT} ./
The default parameters in the scripts is set strictly according to the original repo. However, please modify the parameters in the script if you would like.
python train_base_net.py
python train_full_net.py
Download pretrained checkpoints from here.
Please specify the path of the testset and other settings in the test.py
.
python test.py
Note: The difference between the val
and test
phase in YouTubeDataset.__init__
would be that:
val
: the SDR and HDR images must be both be provided, and the size of SDR image must be IDENTICAL with HDR images, and the YouTubeDataset will resize the SDR images for the net later on.test
: the HDR images may or may not be provided to the dataset. If provided, the size of SDR image should bek
times smaller than HDR images (assumingk
is the parameterscale
of networks)
SSIM and MS-SSIM functions are borrowed from https://github.com/VainF/pytorch-msssim
Please contact me via email (cxt_tsinghua@126.com) for any problems regarding the released code.