wanghu178 / kunet

IJCAI2022

TODO

• Submit source code
• Submit implementation details
• Submit appendex
• Perfect link to the paper
• Citation
• Acknowledgment
• Getting Started
• Reference

KUNet $\rightarrow$ Paper Link

IJCAI2022 (long oral)

Overview

Comparison between the method of reversing the camera imaging pipeline and our method:

Overview of the method:

Implementation details

For the Head, Tail, D block of KIB module and KIC module, we use $3\times 3$ convolution with a step size of 1, and for the X and Y branches in the KIB module, we use $1 \times 1$ convolution. The activation function of all networks is ReLU function. Except for input and output layers, the channels of feature maps are all 64. The down-sampling operation uses a $3 \times 3$ convolution operation with a step size of 2. Up-sampling operation uses PixelShuffle[1]. We use ADAM optimizer[2] with the learning rate of $2e^{-4}$ decayed by a factor of 2 after every $20K$ iterations. The batch size is 12. All models are built on the PyTorch framework and trained with NVIDIA GeForce RTX 2080 SUPER. The total training time is about 6 days.

Getting Started

1. Dataset
2. Configure
3. Test
4. Train
5. Metrics

Dataset

We use two data sets to evaluate our method, i.e.,for image and video tasks. Both of these two data sets contain information about moving light sources, rich colors, highlights and bright. For the image task ,we use NTIRE 2021 dataset. They are 1416 paired training images and 78 test images. For the video task, we conduct experiment on HDRTV[3]. This dataset contains 1235 paired training pictures and 117 test pictures. Please refer to the paper for the details on the processing of this dataset. This dataset can be downloaded from Baidu Netdisk(accsee code: mpr6) .

Configure

h5py==3.6.0
lmdb==1.2.1
pandas==1.1.5
pillow==8.4.0
prefetch-generator==1.0.1
py-opencv==4.5.2
python==3.7.0
pytorch==1.9.0
scikit-image==0.18.3
scipy==1.7.0
tensorboard==2.8.0
torchvision==0.2.2
tqdm==4.61.2
yaml==0.2.5

How to test

we provide the pretrained models to test, which can be downloaded from the link provided by the dataset. Please put the obtained files into models according to the instructions.

• Modify the pictures and model paths in the *.yml file(cd code/options/test. )

• cd code/
  python /test.py -opt /options/test/KIB_mask.yml

How to train

• Prepare the data. Generate the sub-images with specific path size using ./scripts/extract_subimgs_single.py

• make sure that the paths and settings in ./options/train/train_KIB.yml are correct,then run

  cd codes/

• For the image task

  python train.py -opt option/train/train_KIB.yml

• For the video task

  python train.py -opt option/train/train_hdrtv.yml  

• All models and training states are stored in ./experiments.

Metrics

Two metrics are used to evaluate the quantitative performance of different methods on image tasks, including $PSNR$ and $\mu-psnr$ . Five metrics are used to evaluate the quantitative performance of different methods on video tasks, including PSNR, SSIM, SR_SIM, Delta EITP (ITU Rec.2124)and HDR-VDP3. we provide some reference codes in ./metrics for convenient usage.

Reference

1. W. Shi, J. Caballero, F. Husz´ar, J. Totz, A.P.

   Aitken, R. Bishop, D. Rueckert, and Z. Wang. Real-time single image and video super-resolution using an effificient sub-pixel convolutional neural network. In CVPR, pages 1874–1883, 2016.

2. H. Scharr. Optimal filters for extended optical flow. In IWCM, pages 14–29. Springer

3. X. Chen, Z. Zhang, J. Ren, L. Tian, Y. Qiao, and C. Dong. A new journey from SDRTV to HDRTV. In ICCV,

   pages 4500–4509, 2021.

Citation

If our work is helpful to you, please cite our paper:

@inproceedings{wang2022kunet,
  title={KUNet: Imaging Knowledge-Inspired Single HDR Image Reconstruction},
  author={Wang, Hu and Ye, Mao and ZHU, XIATIAN and Li, Shuai and Zhu, Ce and Li, Xue},
  booktitle={IJCAI-ECAI 2022}
}

Ackonwledgment

The code and format is inspired by HDRTV

kunet's Issues

About training dataset and get the .hdr image

Thanks for your work, and your reply for the last issue.
I have some questions about training dataset and geting the .hdr image.
Could I train it on 32-bit HDR image? I trained it on dataset "HDR-Real" and "HDR_Synth."
And could I just using cv2.imwrite('hdr_img_path.hdr', sr_img) to get the HDR image? It seems there are some problem when I using this function.
How can I get the .hdr image when I using this work?

Thank you~

AttributeError: 'GenerationModel' object has no attribute 'mask_pix'

I change network_G which_model_G of train_KIB.yml to KIB_DM_F_1x1_mask, but got this error:

22-06-30 11:35:58.139 - INFO: Model [GenerationModel] is created.
22-06-30 11:35:58.139 - INFO: Start training from epoch: 0, iter: 0
Traceback (most recent call last):
  File "F:/code/KUNet-main/codes/train.py", line 171, in <module>
    model.optimize_parameters(current_step)
  File "F:\code\KUNet-main\codes\models\Generation_condition.py", line 111, in optimize_parameters
    KIB1_loss = self.mask_pix(KIB1,self.real_H,mask)
AttributeError: 'GenerationModel' object has no attribute 'mask_pix'

How can I fix it, thanks!

where is the "./scripts/extract_subimgs_single.py" file?

Maybe a little miss

感谢您的工作，让我受益良多！
在原文中的KIC模块，将处理后的特征F波浪和余弦相似度的系数相乘。
但是在wanghu178/KUNet/tree/main/codes/models中，
KIC模块的实现：
def forward(self,x,guidance):
denos_x = self.denosing(x)
tmo_x_hat = self.tmo(denos_x)
guidance_map = self.cos_sim(tmo_x_hat,guidance).unsqueeze(dim=1)
out =self.refin(guidance_map*x)
return out

可见out =self.refin(guidance_map*x)，余弦相似度直接和输入的特征x相乘，而不是特征F波浪。请问这是一个小小的miss嘛？感谢你的回复！

Pretrained model

Hello,

I would appreciate it if you could assist me in running your code with your pretrained model. I encountered a difficulty while trying to access the model from the provided Baidu link, as it requires a Baidu account, which I currently do not possess. Creating a Baidu account is not a straightforward process for me. Hence, I kindly request you to consider uploading the checkpoint files to Google Drive or including them directly in your GitHub repository.

Thank you for your understanding.