This repository is for "OoDHDR-Codec: Out-of-Distribution Generalization for HDR Image Compression"
(AAAI, 2022)
By Linfeng Cao, Aofan Jiang, Wei Li, Huaying Wu and Nanyang Ye
- Paper full supplementary is available here (Google Drive)
Overview of the proposed OoDHDR-codec framework.
DNN backbone in our project.
- Python (3.8.3)
- PyTorch (>=1.6.0)
(see setup.py
for the full list)
From source:
git clone https://github.com/caolinfeng/OoDHDR-codec OodHDR_codec
cd OodHDR_codec
pip install -U pip && pip install -e .
SDR training datasets can be downloaded from:
HDR training datasets (.hdr) can be downloaded from:
Evaluation datasets:
A training script with a regularization & fusion loss is provided in
examples/train_ood.py
. Please specify the corresponding HDR & SDR datasets in the file. The custom ImageFolder structure in our project only supports for the RGBE (.hdr) input currently.
You can run the script for a training pipeline:
python examples/train_ood.py --lambda 12 --epochs 250 --cuda --save --gpu 0,1,2,3 --batch-size 32 --rw 1e-5 --pw 2 --sdr_w 0.95
Note: 'rw, pw, sdr_w' are the hyper-parameters of the constructed loss. To achevie the optimal performance of a certain network, it is recommended to use the grid search. You can also modify other parameters to change the model and training strategy in the file or cmd.
To evaluate a trained model on HDR and SDR dataset, evaluation scripts (examples/test_hdr.py
, examples/test_sdr.py
) are provided. Please modify the testing dataset path in the corresponding file, and specify the trained model path in cmd:
python examples/test_hdr.py --pth /XXX.pth.tar
python examples/test_sdr.py --pth /XXX.pth.tar
- The PyTorch pre-trained models can be downloaded here (Google Drive)
To assess the compression performance on HDR images, the evaluation metrics of puPSNR, puSSIM and HDR-VDP (3.0.6) are used, the source codes (Matlab version) can be downloaded here.
If you find this code useful, please cite our paper:
@inproceedings{Cao2020OodHDR,
title = {OoDHDR-Codec: Out-of-Distribution Generalization for HDR Image Compression},
author = {Linfeng Cao, Aofan Jiang, Wei Li, Huaying Wu and Nanyang Ye},
booktitle = {Proceedings of the Thirty-Sixth AAAI Conference on Artificial Intelligence (AAAI)},
year = {2022}
}
- This project is developed based on CompressAI library