Skip to content

Commit d92cd21

Browse files
NicolasHugfmassa
authored andcommitted
[fbsync] Add video GPU decoder (#5019)
Summary: * [WIP] Add video GPU decoder * Expose use_dev_frame to python class and handle it internally * Fixed invalid argument CUDA error * Fixed empty and missing frames * Free remaining frames in the queue * Added nv12 to yuv420 conversion support for host frames * Added unit test and cleaned up code * Use CUDA_HOME inside if * Undo commented out code * Add Readme * Remove output_format and use_device_frame optional arguments from the VideoReader API * Cleaned up init() * Fix warnings * Fix python linter errors * Fix linter issues in setup.py * clang-format * Make reformat private * Member function naming * Add comments * Variable renaming * Code cleanup * Make return type of decode() void * Replace printing errors with throwing runtime_error * Replaced runtime_error with TORCH_CHECK in demuxer.h * Use CUDAGuard instead of cudaSetDevice * Remove printf * Use Tensor instead of uint8* and remove cuMemAlloc/cuMemFree * Use TORCH_CHECK instead of runtime_error * Use TORCHVISION_INCLUDE and TORCHVISION_LIBRARY to pass video codec location * Include ffmpeg_include_dir * Remove space * Removed use of runtime_error * Update Readme * Check for bsf.h * Change struct initialisation style * Clean-up get_operating_point * Make variable naming convention uniform * Move checking for bsf.h around * Fix linter error Reviewed By: datumbox, prabhat00155 Differential Revision: D33405358 fbshipit-source-id: 0e6251389508309a23c7afd843f298208dcd67e8 Co-authored-by: Francisco Massa <[email protected]>
1 parent c2b6655 commit d92cd21

File tree

10 files changed

+1020
-1
lines changed

10 files changed

+1020
-1
lines changed

setup.py

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -427,6 +427,59 @@ def get_extensions():
427427
)
428428
)
429429

430+
# Locating video codec
431+
# CUDA_HOME should be set to the cuda root directory.
432+
# TORCHVISION_INCLUDE and TORCHVISION_LIBRARY should include the location to
433+
# video codec header files and libraries respectively.
434+
video_codec_found = (
435+
extension is CUDAExtension
436+
and CUDA_HOME is not None
437+
and any([os.path.exists(os.path.join(folder, "cuviddec.h")) for folder in vision_include])
438+
and any([os.path.exists(os.path.join(folder, "nvcuvid.h")) for folder in vision_include])
439+
and any([os.path.exists(os.path.join(folder, "libnvcuvid.so")) for folder in library_dirs])
440+
)
441+
442+
print(f"video codec found: {video_codec_found}")
443+
444+
if (
445+
video_codec_found
446+
and has_ffmpeg
447+
and any([os.path.exists(os.path.join(folder, "libavcodec", "bsf.h")) for folder in ffmpeg_include_dir])
448+
):
449+
gpu_decoder_path = os.path.join(extensions_dir, "io", "decoder", "gpu")
450+
gpu_decoder_src = glob.glob(os.path.join(gpu_decoder_path, "*.cpp"))
451+
cuda_libs = os.path.join(CUDA_HOME, "lib64")
452+
cuda_inc = os.path.join(CUDA_HOME, "include")
453+
454+
ext_modules.append(
455+
extension(
456+
"torchvision.Decoder",
457+
gpu_decoder_src,
458+
include_dirs=include_dirs + [gpu_decoder_path] + [cuda_inc] + ffmpeg_include_dir,
459+
library_dirs=ffmpeg_library_dir + library_dirs + [cuda_libs],
460+
libraries=[
461+
"avcodec",
462+
"avformat",
463+
"avutil",
464+
"swresample",
465+
"swscale",
466+
"nvcuvid",
467+
"cuda",
468+
"cudart",
469+
"z",
470+
"pthread",
471+
"dl",
472+
],
473+
extra_compile_args=extra_compile_args,
474+
)
475+
)
476+
else:
477+
print(
478+
"The installed version of ffmpeg is missing the header file 'bsf.h' which is "
479+
"required for GPU video decoding. Please install the latest ffmpeg from conda-forge channel:"
480+
" `conda install -c conda-forge ffmpeg`."
481+
)
482+
430483
return ext_modules
431484

432485

test/test_video_gpu_decoder.py

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,41 @@
1+
import os
2+
3+
import pytest
4+
import torch
5+
from torchvision.io import _HAS_VIDEO_DECODER, VideoReader
6+
7+
try:
8+
import av
9+
except ImportError:
10+
av = None
11+
12+
VIDEO_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "assets", "videos")
13+
14+
test_videos = [
15+
"RATRACE_wave_f_nm_np1_fr_goo_37.avi",
16+
"TrumanShow_wave_f_nm_np1_fr_med_26.avi",
17+
"v_SoccerJuggling_g23_c01.avi",
18+
"v_SoccerJuggling_g24_c01.avi",
19+
"R6llTwEh07w.mp4",
20+
"SOX5yA1l24A.mp4",
21+
"WUzgd7C1pWA.mp4",
22+
]
23+
24+
25+
@pytest.mark.skipif(_HAS_VIDEO_DECODER is False, reason="Didn't compile with support for gpu decoder")
26+
class TestVideoGPUDecoder:
27+
@pytest.mark.skipif(av is None, reason="PyAV unavailable")
28+
def test_frame_reading(self):
29+
for test_video in test_videos:
30+
full_path = os.path.join(VIDEO_DIR, test_video)
31+
decoder = VideoReader(full_path, device="cuda:0")
32+
with av.open(full_path) as container:
33+
for av_frame in container.decode(container.streams.video[0]):
34+
av_frames = torch.tensor(av_frame.to_ndarray().flatten())
35+
vision_frames = next(decoder)["data"]
36+
mean_delta = torch.mean(torch.abs(av_frames.float() - decoder._reformat(vision_frames).float()))
37+
assert mean_delta < 0.1
38+
39+
40+
if __name__ == "__main__":
41+
pytest.main([__file__])
Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,21 @@
1+
GPU Decoder
2+
===========
3+
4+
GPU decoder depends on ffmpeg for demuxing, uses NVDECODE APIs from the nvidia-video-codec sdk and uses cuda for processing on gpu. In order to use this, please follow the following steps:
5+
6+
* Download the latest `nvidia-video-codec-sdk <https://developer.nvidia.com/nvidia-video-codec-sdk/download>`_
7+
* Extract the zipped file.
8+
* Set TORCHVISION_INCLUDE environment variable to the location of the video codec headers(`nvcuvid.h` and `cuviddec.h`), which would be under `Interface` directory.
9+
* Set TORCHVISION_LIBRARY environment variable to the location of the video codec library(`libnvcuvid.so`), which would be under `Lib/linux/stubs/x86_64` directory.
10+
* Install the latest ffmpeg from `conda-forge` channel.
11+
12+
.. code:: bash
13+
14+
conda install -c conda-forge ffmpeg
15+
16+
* Set CUDA_HOME environment variable to the cuda root directory.
17+
* Build torchvision from source:
18+
19+
.. code:: bash
20+
21+
python setup.py install

0 commit comments

Comments
 (0)