This is the official implementation of Multi-stage Spatio-Temporal Fusion Network for Fast and Accurate Video Bit-depth Enhancement (IEEE Transactions on Multimedia 2023) in Pytorch.
Copyright(c) 2023 Jing Liu
If you use this code, please cite the following publication:
J. Liu, Z. Fan, Z. Yang, Y. Su, and X. Yang, "Multi-stage Spatio-Temporal Fusion Network for Fast and Accurate Video Bit-depth Enhancement", to appear in IEEE Transactions on Multimedia.
You can run conda env create -f mstfn.yml
to create the environment
- python=3.6
- ubuntu 20.04
- pytorch=1.6.0
- torchvision=0.7
- opencv-python==4.6
We provide folders "./dataset/test_frames_16bit" and "./dataset/test_frames_12bit" to realize 4-bit to 16-bit and 8-bit to 12-bit BDE tasks respectively. When testing, prepare the testing dataset, and modify the dataset path and other related content in the code. We provide a test sequence(16bit and 12bit) for testing. You can directly test on the testing frames by running
$ python test.py
- We provide recovery results in the folder "result". When testing, the predicted results are saved in the folder "test" .
- The files "./metrics/SSIM_PSNR.m" is used to calculate PSNR and SSIM.