Skip to content

Commit dcdaeee

Browse files
DeepMC-DCVCDeepMC-DCVC
committed
DCVC
1 parent 9f984e4 commit dcdaeee

32 files changed

+4388
-0
lines changed

‎.gitignore

+6
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
*.pyc
2+
.vscode/
3+
*.bin
4+
*.png
5+
*.so
6+
build/

‎README.md

+94
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,94 @@
1+
# Introduction
2+
3+
Official Pytorch implementation for [Deep Contextual Video Compression](https://proceedings.neurips.cc/paper/2021/file/96b250a90d3cf0868c83f8c965142d2a-Paper.pdf), NeurIPS 2021
4+
5+
# Prerequisites
6+
* Python 3.8 and conda, get [Conda](https://www.anaconda.com/)
7+
* CUDA 11.0
8+
* Environment
9+
```
10+
conda create -n $YOUR_PY38_ENV_NAME python=3.8
11+
conda activate $YOUR_PY38_ENV_NAME
12+
13+
pip install torch==1.7.1+cu110 torchvision==0.8.2+cu110 torchaudio==0.7.2 -f https://download.pytorch.org/whl/torch_stable.html
14+
python -m pip install -r requirements.txt
15+
```
16+
17+
18+
19+
# Test dataset
20+
Currenlty the spatial resolution of video needs to be cropped into the integral times of 64.
21+
22+
The dataset format can be seen in dataset_config_example.json.
23+
24+
For example, one video of HEVC Class B can be prepared as:
25+
* Crop the original YUV via ffmpeg:
26+
```
27+
ffmpeg -pix_fmt yuv420p -s 1920x1080 -i BasketballDrive_1920x1080_50.yuv -vf crop=1920:1024:0:0 BasketballDrive_1920x1024_50.yuv
28+
```
29+
* Make the video path:
30+
```
31+
mkdir BasketballDrive_1920x1024_50
32+
```
33+
* Convert YUV to PNG:
34+
```
35+
ffmpeg -pix_fmt yuv420p -s 1920x1024 -i BasketballDrive_1920x1024_50.yuv -f image2 BasketballDrive_1920x1024_50/im%05d.png
36+
```
37+
At last, the folder structure of dataset is like:
38+
39+
/media/data/HEVC_B/
40+
* BQTerrace_1920x1024_60/
41+
- im00001.png
42+
- im00002.png
43+
- im00003.png
44+
- ...
45+
* BasketballDrive_1920x1024_50/
46+
- im00001.png
47+
- im00002.png
48+
- im00003.png
49+
- ...
50+
* ...
51+
/media/data/HEVC_D
52+
/media/data/HEVC_C/
53+
...
54+
55+
# Pretrained models
56+
57+
* Download CompressAI models
58+
```
59+
cd checkpoints/
60+
python download_compressai_models.py
61+
cd ..
62+
```
63+
64+
* Download [DCVC models](https://1drv.ms/u/s!AozfVVwtWWYoiS5mcGX320bFXI0k?e=iMeykH) and put them into /checkpoints folder.
65+
66+
# Test DCVC
67+
68+
Example of test the PSNR model:
69+
```bash
70+
python test_video.py --i_frame_model_name cheng2020-anchor --i_frame_model_path checkpoints/cheng2020-anchor-3-e49be189.pth.tar checkpoints/cheng2020-anchor-4-98b0b468.pth.tar checkpoints/cheng2020-anchor-5-23852949.pth.tar checkpoints/cheng2020-anchor-6-4c052b1a.pth.tar --test_config dataset_config_example.json --cuda true --cuda_device 0,1,2,3 --worker 4 --output_json_result_path DCVC_result_psnr.json --model_type psnr --recon_bin_path recon_bin_folder_psnr --model_path checkpoints/model_dcvc_quality_0_psnr.pth checkpoints/model_dcvc_quality_1_psnr.pth checkpoints/model_dcvc_quality_2_psnr.pth checkpoints/model_dcvc_quality_3_psnr.pth
71+
```
72+
73+
Example of test the MSSSIM model:
74+
```bash
75+
python test_video.py --i_frame_model_name bmshj2018-hyperprior --i_frame_model_path checkpoints/bmshj2018-hyperprior-ms-ssim-3-92dd7878.pth.tar checkpoints/bmshj2018-hyperprior-ms-ssim-4-4377354e.pth.tar checkpoints/bmshj2018-hyperprior-ms-ssim-5-c34afc8d.pth.tar checkpoints/bmshj2018-hyperprior-ms-ssim-6-3a6d8229.pth.tar --test_config dataset_config_example.json --cuda true --cuda_device 0,1,2,3 --worker 4 --output_json_result_path DCVC_result_msssim.json --model_type msssim --recon_bin_path recon_bin_folder_msssim --model_path checkpoints/model_dcvc_quality_0_msssim.pth checkpoints/model_dcvc_quality_1_msssim.pth checkpoints/model_dcvc_quality_2_msssim.pth checkpoints/model_dcvc_quality_3_msssim.pth
76+
```
77+
It is recommended that the ```--worker``` number is equal to your GPU number.
78+
79+
80+
# Acknowledgement
81+
The implementation is based on [CompressAI](https://github.com/InterDigitalInc/CompressAI) and [PyTorchVideoCompression](https://github.com/ZhihaoHu/PyTorchVideoCompression). The model weights of intra coding come from [CompressAI](https://github.com/InterDigitalInc/CompressAI).
82+
83+
# Citation
84+
If you find this work useful for your research, please cite:
85+
86+
```
87+
@article{li2021deep,
88+
title={Deep Contextual Video Compression},
89+
author={Li, Jiahao and Li, Bin and Lu, Yan},
90+
journal={Advances in Neural Information Processing Systems},
91+
volume={34},
92+
year={2021}
93+
}
94+
```

‎checkpoints/download_compressai_models.py

+19
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,19 @@
1+
import urllib.request
2+
3+
# The model weights of intra coding come from CompressAI.
4+
root_url = "https://compressai.s3.amazonaws.com/models/v1/"
5+
6+
model_names = [
7+
"bmshj2018-hyperprior-ms-ssim-3-92dd7878.pth.tar",
8+
"bmshj2018-hyperprior-ms-ssim-4-4377354e.pth.tar",
9+
"bmshj2018-hyperprior-ms-ssim-5-c34afc8d.pth.tar",
10+
"bmshj2018-hyperprior-ms-ssim-6-3a6d8229.pth.tar",
11+
"cheng2020-anchor-3-e49be189.pth.tar",
12+
"cheng2020-anchor-4-98b0b468.pth.tar",
13+
"cheng2020-anchor-5-23852949.pth.tar",
14+
"cheng2020-anchor-6-4c052b1a.pth.tar",
15+
]
16+
17+
for model in model_names:
18+
print(f"downloading {model}")
19+
urllib.request.urlretrieve(root_url+model, model)

‎dataset_config_example.json

+85
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,85 @@
1+
{
2+
"HEVC_B": {
3+
"base_path": "/media/data/HEVC_B",
4+
"sequences": {
5+
"BQTerrace_1920x1024_60": {"frames": 100, "gop": 10},
6+
"BasketballDrive_1920x1024_50": {"frames": 100, "gop": 10},
7+
"Cactus_1920x1024_50": {"frames": 100, "gop": 10},
8+
"Kimono1_1920x1024_24": {"frames": 100, "gop": 10},
9+
"ParkScene_1920x1024_24": {"frames": 100, "gop": 10}
10+
}
11+
},
12+
"HEVC_C": {
13+
"base_path": "/media/data/HEVC_C",
14+
"sequences": {
15+
"BQMall_832x448_60": {"frames": 100, "gop": 10},
16+
"BasketballDrill_832x448_50": {"frames": 100, "gop": 10},
17+
"PartyScene_832x448_50": {"frames": 100, "gop": 10},
18+
"RaceHorses_832x448_30": {"frames": 100, "gop": 10}
19+
}
20+
},
21+
"HEVC_D": {
22+
"base_path": "/media/data/HEVC_D",
23+
"sequences": {
24+
"BasketballPass_384x192_50": {"frames": 100, "gop": 10},
25+
"BlowingBubbles_384x192_50": {"frames": 100, "gop": 10},
26+
"BQSquare_384x192_50": {"frames": 100, "gop": 10},
27+
"RaceHorses_384x192_50": {"frames": 100, "gop": 10}
28+
}
29+
},
30+
"HEVC_E": {
31+
"base_path": "/media/data/HEVC_E",
32+
"sequences": {
33+
"FourPeople_1280x704_60": {"frames": 100, "gop": 10},
34+
"Johnny_1280x704_60": {"frames": 100, "gop": 10},
35+
"KristenAndSara_1280x704_60": {"frames": 100, "gop": 10}
36+
}
37+
},
38+
"UVG": {
39+
"base_path": "/media/data/UVGDataSet_crop",
40+
"sequences": {
41+
"Beauty_1920x1024_120fps_420_8bit_YUV": {"frames": 120, "gop": 12},
42+
"Bosphorus_1920x1024_120fps_420_8bit_YUV": {"frames": 120, "gop": 12},
43+
"HoneyBee_1920x1024_120fps_420_8bit_YUV": {"frames": 120, "gop": 12},
44+
"Jockey_1920x1024_120fps_420_8bit_YUV": {"frames": 120, "gop": 12},
45+
"ReadySteadyGo_1920x1024_120fps_420_8bit_YUV": {"frames": 120, "gop": 12},
46+
"ShakeNDry_1920x1024_120fps_420_8bit_YUV": {"frames": 120, "gop": 12},
47+
"YachtRide_1920x1024_120fps_420_8bit_YUV": {"frames": 120, "gop": 12}
48+
}
49+
},
50+
"MCL-JCV": {
51+
"base_path": "/media/data/MCL-JCV",
52+
"sequences": {
53+
"videoSRC01_1920x1024_30": {"frames": 120, "gop": 12},
54+
"videoSRC02_1920x1024_30": {"frames": 120, "gop": 12},
55+
"videoSRC03_1920x1024_30": {"frames": 120, "gop": 12},
56+
"videoSRC04_1920x1024_30": {"frames": 120, "gop": 12},
57+
"videoSRC05_1920x1024_25": {"frames": 120, "gop": 12},
58+
"videoSRC06_1920x1024_25": {"frames": 120, "gop": 12},
59+
"videoSRC07_1920x1024_25": {"frames": 120, "gop": 12},
60+
"videoSRC08_1920x1024_25": {"frames": 120, "gop": 12},
61+
"videoSRC09_1920x1024_25": {"frames": 120, "gop": 12},
62+
"videoSRC10_1920x1024_30": {"frames": 120, "gop": 12},
63+
"videoSRC11_1920x1024_30": {"frames": 120, "gop": 12},
64+
"videoSRC12_1920x1024_30": {"frames": 120, "gop": 12},
65+
"videoSRC13_1920x1024_30": {"frames": 120, "gop": 12},
66+
"videoSRC14_1920x1024_30": {"frames": 120, "gop": 12},
67+
"videoSRC15_1920x1024_30": {"frames": 120, "gop": 12},
68+
"videoSRC16_1920x1024_30": {"frames": 120, "gop": 12},
69+
"videoSRC17_1920x1024_24": {"frames": 120, "gop": 12},
70+
"videoSRC18_1920x1024_25": {"frames": 120, "gop": 12},
71+
"videoSRC19_1920x1024_30": {"frames": 120, "gop": 12},
72+
"videoSRC20_1920x1024_25": {"frames": 120, "gop": 12},
73+
"videoSRC21_1920x1024_24": {"frames": 120, "gop": 12},
74+
"videoSRC22_1920x1024_24": {"frames": 120, "gop": 12},
75+
"videoSRC23_1920x1024_24": {"frames": 120, "gop": 12},
76+
"videoSRC24_1920x1024_24": {"frames": 120, "gop": 12},
77+
"videoSRC25_1920x1024_24": {"frames": 120, "gop": 12},
78+
"videoSRC26_1920x1024_30": {"frames": 120, "gop": 12},
79+
"videoSRC27_1920x1024_30": {"frames": 120, "gop": 12},
80+
"videoSRC28_1920x1024_30": {"frames": 120, "gop": 12},
81+
"videoSRC29_1920x1024_24": {"frames": 120, "gop": 12},
82+
"videoSRC30_1920x1024_30": {"frames": 120, "gop": 12}
83+
}
84+
}
85+
}

‎requirements.txt

+6
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
numpy
2+
scipy
3+
matplotlib
4+
Pillow
5+
pytorch-msssim
6+
tqdm

‎src/cpp/3rdparty/CMakeLists.txt

+2
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,2 @@
1+
add_subdirectory(pybind11)
2+
add_subdirectory(ryg_rans)

‎src/cpp/3rdparty/pybind11/CMakeLists.txt

+22
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,22 @@
1+
# set(PYBIND11_PYTHON_VERSION 3.8 CACHE STRING "")
2+
configure_file(CMakeLists.txt.in pybind11-download/CMakeLists.txt)
3+
execute_process(COMMAND ${CMAKE_COMMAND} -G "${CMAKE_GENERATOR}" .
4+
RESULT_VARIABLE result
5+
WORKING_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR}/pybind11-download )
6+
if(result)
7+
message(FATAL_ERROR "CMake step for pybind11 failed: ${result}")
8+
endif()
9+
execute_process(COMMAND ${CMAKE_COMMAND} --build .
10+
RESULT_VARIABLE result
11+
WORKING_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR}/pybind11-download )
12+
if(result)
13+
message(FATAL_ERROR "Build step for pybind11 failed: ${result}")
14+
endif()
15+
16+
add_subdirectory(${CMAKE_CURRENT_BINARY_DIR}/pybind11-src/
17+
${CMAKE_CURRENT_BINARY_DIR}/pybind11-build/
18+
EXCLUDE_FROM_ALL)
19+
20+
set(PYBIND11_INCLUDE
21+
${CMAKE_CURRENT_BINARY_DIR}/pybind11-src/include/
22+
CACHE INTERNAL "")

‎src/cpp/3rdparty/pybind11/CMakeLists.txt.in

+33
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,33 @@
1+
cmake_minimum_required(VERSION 3.6.3)
2+
3+
project(pybind11-download NONE)
4+
5+
include(ExternalProject)
6+
if(IS_DIRECTORY "${PROJECT_BINARY_DIR}/3rdparty/pybind11/pybind11-src/include")
7+
ExternalProject_Add(pybind11
8+
GIT_REPOSITORY https://github.com/pybind/pybind11.git
9+
GIT_TAG v2.6.1
10+
GIT_SHALLOW 1
11+
SOURCE_DIR "${CMAKE_CURRENT_BINARY_DIR}/pybind11-src"
12+
BINARY_DIR "${CMAKE_CURRENT_BINARY_DIR}/pybind11-build"
13+
DOWNLOAD_COMMAND ""
14+
UPDATE_COMMAND ""
15+
CONFIGURE_COMMAND ""
16+
BUILD_COMMAND ""
17+
INSTALL_COMMAND ""
18+
TEST_COMMAND ""
19+
)
20+
else()
21+
ExternalProject_Add(pybind11
22+
GIT_REPOSITORY https://github.com/pybind/pybind11.git
23+
GIT_TAG v2.6.1
24+
GIT_SHALLOW 1
25+
SOURCE_DIR "${CMAKE_CURRENT_BINARY_DIR}/pybind11-src"
26+
BINARY_DIR "${CMAKE_CURRENT_BINARY_DIR}/pybind11-build"
27+
UPDATE_COMMAND ""
28+
CONFIGURE_COMMAND ""
29+
BUILD_COMMAND ""
30+
INSTALL_COMMAND ""
31+
TEST_COMMAND ""
32+
)
33+
endif()

‎src/cpp/3rdparty/ryg_rans/CMakeLists.txt

+21
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,21 @@
1+
configure_file(CMakeLists.txt.in ryg_rans-download/CMakeLists.txt)
2+
execute_process(COMMAND ${CMAKE_COMMAND} -G "${CMAKE_GENERATOR}" .
3+
RESULT_VARIABLE result
4+
WORKING_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR}/ryg_rans-download )
5+
if(result)
6+
message(FATAL_ERROR "CMake step for ryg_rans failed: ${result}")
7+
endif()
8+
execute_process(COMMAND ${CMAKE_COMMAND} --build .
9+
RESULT_VARIABLE result
10+
WORKING_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR}/ryg_rans-download )
11+
if(result)
12+
message(FATAL_ERROR "Build step for ryg_rans failed: ${result}")
13+
endif()
14+
15+
# add_subdirectory(${CMAKE_CURRENT_BINARY_DIR}/ryg_rans-src/
16+
# ${CMAKE_CURRENT_BINARY_DIR}/ryg_rans-build
17+
# EXCLUDE_FROM_ALL)
18+
19+
set(RYG_RANS_INCLUDE
20+
${CMAKE_CURRENT_BINARY_DIR}/ryg_rans-src/
21+
CACHE INTERNAL "")

‎src/cpp/3rdparty/ryg_rans/CMakeLists.txt.in

+33
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,33 @@
1+
cmake_minimum_required(VERSION 3.6.3)
2+
3+
project(ryg_rans-download NONE)
4+
5+
include(ExternalProject)
6+
if(EXISTS "${PROJECT_BINARY_DIR}/3rdparty/ryg_rans/ryg_rans-src/rans64.h")
7+
ExternalProject_Add(ryg_rans
8+
GIT_REPOSITORY https://github.com/rygorous/ryg_rans.git
9+
GIT_TAG c9d162d996fd600315af9ae8eb89d832576cb32d
10+
GIT_SHALLOW 1
11+
SOURCE_DIR "${CMAKE_CURRENT_BINARY_DIR}/ryg_rans-src"
12+
BINARY_DIR "${CMAKE_CURRENT_BINARY_DIR}/ryg_rans-build"
13+
DOWNLOAD_COMMAND ""
14+
UPDATE_COMMAND ""
15+
CONFIGURE_COMMAND ""
16+
BUILD_COMMAND ""
17+
INSTALL_COMMAND ""
18+
TEST_COMMAND ""
19+
)
20+
else()
21+
ExternalProject_Add(ryg_rans
22+
GIT_REPOSITORY https://github.com/rygorous/ryg_rans.git
23+
GIT_TAG c9d162d996fd600315af9ae8eb89d832576cb32d
24+
GIT_SHALLOW 1
25+
SOURCE_DIR "${CMAKE_CURRENT_BINARY_DIR}/ryg_rans-src"
26+
BINARY_DIR "${CMAKE_CURRENT_BINARY_DIR}/ryg_rans-build"
27+
UPDATE_COMMAND ""
28+
CONFIGURE_COMMAND ""
29+
BUILD_COMMAND ""
30+
INSTALL_COMMAND ""
31+
TEST_COMMAND ""
32+
)
33+
endif()

‎src/cpp/CMakeLists.txt

+20
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,20 @@
1+
cmake_minimum_required (VERSION 3.6.3)
2+
project (ErrorRecovery)
3+
4+
set(CMAKE_CONFIGURATION_TYPES "Debug;Release;RelWithDebInfo" CACHE STRING "" FORCE)
5+
6+
set(CMAKE_CXX_STANDARD 17)
7+
set(CMAKE_CXX_STANDARD_REQUIRED ON)
8+
set(CMAKE_CXX_EXTENSIONS OFF)
9+
10+
# treat warning as error
11+
if (MSVC)
12+
add_compile_options(/W4 /WX)
13+
else()
14+
add_compile_options(-Wall -Wextra -pedantic -Werror)
15+
endif()
16+
17+
# The sequence is tricky, put 3rd party first
18+
add_subdirectory(3rdparty)
19+
add_subdirectory (ops)
20+
add_subdirectory (rans)

‎src/cpp/ops/CMakeLists.txt

+25
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,25 @@
1+
cmake_minimum_required(VERSION 3.7)
2+
set(PROJECT_NAME MLCodec_CXX)
3+
project(${PROJECT_NAME})
4+
5+
set(cxx_source
6+
ops.cpp
7+
)
8+
9+
set(include_dirs
10+
${CMAKE_CURRENT_SOURCE_DIR}
11+
${PYBIND11_INCLUDE}
12+
)
13+
14+
pybind11_add_module(${PROJECT_NAME} ${cxx_source})
15+
16+
target_include_directories (${PROJECT_NAME} PUBLIC ${include_dirs})
17+
18+
# The post build argument is executed after make!
19+
add_custom_command(
20+
TARGET ${PROJECT_NAME} POST_BUILD
21+
COMMAND
22+
"${CMAKE_COMMAND}" -E copy
23+
"$<TARGET_FILE:${PROJECT_NAME}>"
24+
"${CMAKE_CURRENT_SOURCE_DIR}/../../entropy_models/"
25+
)

0 commit comments

Comments
 (0)