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:

• DIV2K
• Flickr2K

HDR training datasets (.hdr) can be downloaded from:

• HDRIHEVEN
• pfstools (HDR Image Gallery)
• HDRLabs
• Funt and Shi 2010

Evaluation datasets:

• HDR (Google Drive)
• SDR-Kodak

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