Add hardware acceleration to video decoding

av/codec/context.pxd

@@ -63,10 +63,12 @@ cdef class CodecContext(object):
     cdef _send_packet_and_recv(self, Packet packet)
     cdef _recv_frame(self)
 
+    cdef _transfer_hwframe(self, Frame frame)
+
     # Implemented by children for the generic send/recv API, so we have the
     # correct subclass of Frame.
     cdef Frame _next_frame
     cdef Frame _alloc_next_frame(self)
 
 
-cdef CodecContext wrap_codec_context(lib.AVCodecContext*, const lib.AVCodec*, bint allocated)
+cdef CodecContext wrap_codec_context(lib.AVCodecContext*, const lib.AVCodec*, bint allocated, dict hwaccel)

av/codec/context.pyx

@@ -18,15 +18,15 @@ from av.utils cimport avdict_to_dict, avrational_to_fraction, to_avrational
 cdef object _cinit_sentinel = object()
 
 
-cdef CodecContext wrap_codec_context(lib.AVCodecContext *c_ctx, const lib.AVCodec *c_codec, bint allocated):
+cdef CodecContext wrap_codec_context(lib.AVCodecContext *c_ctx, const lib.AVCodec *c_codec, bint allocated, dict hwaccel):
     """Build an av.CodecContext for an existing AVCodecContext."""
 
     cdef CodecContext py_ctx
 
     # TODO: This.
     if c_ctx.codec_type == lib.AVMEDIA_TYPE_VIDEO:
         from av.video.codeccontext import VideoCodecContext
-        py_ctx = VideoCodecContext(_cinit_sentinel)
+        py_ctx = VideoCodecContext(_cinit_sentinel, hwaccel=hwaccel)
     elif c_ctx.codec_type == lib.AVMEDIA_TYPE_AUDIO:
         from av.audio.codeccontext import AudioCodecContext
         py_ctx = AudioCodecContext(_cinit_sentinel)
@@ -65,7 +65,7 @@ cdef class CodecContext(object):
     def create(codec, mode=None):
         cdef Codec cy_codec = codec if isinstance(codec, Codec) else Codec(codec, mode)
         cdef lib.AVCodecContext *c_ctx = lib.avcodec_alloc_context3(cy_codec.ptr)
-        return wrap_codec_context(c_ctx, cy_codec.ptr, True)
+        return wrap_codec_context(c_ctx, cy_codec.ptr, True, None)
 
     def __cinit__(self, sentinel=None, *args, **kwargs):
         if sentinel is not _cinit_sentinel:
@@ -283,10 +283,15 @@ cdef class CodecContext(object):
             return
         err_check(res)
 
+        frame = self._transfer_hwframe(frame)
+
         if not res:
             self._next_frame = None
             return frame
 
+    cdef _transfer_hwframe(self, Frame frame):
+        return frame
+
     cdef _recv_packet(self):
 
         cdef Packet packet = Packet()

av/container/core.pxd

@@ -35,6 +35,8 @@ cdef class Container(object):
     cdef readonly dict container_options
     cdef readonly list stream_options
 
+    cdef dict hwaccel
+
     cdef readonly StreamContainer streams
     cdef readonly dict metadata
 

av/container/core.pyx

@@ -35,7 +35,7 @@ cdef class Container(object):
     def __cinit__(self, sentinel, file_, format_name, options,
                   container_options, stream_options,
                   metadata_encoding, metadata_errors,
-                  buffer_size):
+                  buffer_size, hwaccel):
 
         if sentinel is not _cinit_sentinel:
             raise RuntimeError('cannot construct base Container')
@@ -59,6 +59,8 @@ cdef class Container(object):
         self.metadata_encoding = metadata_encoding
         self.metadata_errors = metadata_errors
 
+        self.hwaccel = hwaccel
+
         if format_name is not None:
             self.format = ContainerFormat(format_name)
 
@@ -188,7 +190,7 @@ cdef class Container(object):
 def open(file, mode=None, format=None, options=None,
          container_options=None, stream_options=None,
          metadata_encoding=None, metadata_errors='strict',
-         buffer_size=32768):
+         buffer_size=32768, hwaccel=None):
     """open(file, mode='r', format=None, options=None, metadata_encoding=None, metadata_errors='strict')
 
     Main entrypoint to opening files/streams.
@@ -206,6 +208,8 @@ def open(file, mode=None, format=None, options=None,
         ``str.encode`` parameter. Defaults to strict.
     :param int buffer_size: Size of buffer for Python input/output operations in bytes.
         Honored only when ``file`` is a file-like object. Defaults to 32768 (32k).
+    :param dict hwaccel: The desired device parameters to use for hardware acceleration
+        including device_type_name (e.x. cuda) and optional device (e.x. '/dev/dri/renderD128').
 
     For devices (via ``libavdevice``), pass the name of the device to ``format``,
     e.g.::
@@ -222,12 +226,15 @@ def open(file, mode=None, format=None, options=None,
     if mode is None:
         mode = 'r'
 
+    if hwaccel is not None:
+        hwaccel = dict(hwaccel)
+
     if mode.startswith('r'):
         return InputContainer(
             _cinit_sentinel, file, format, options,
             container_options, stream_options,
             metadata_encoding, metadata_errors,
-            buffer_size
+            buffer_size, hwaccel=hwaccel,
         )
     if mode.startswith('w'):
         if stream_options:
@@ -236,6 +243,6 @@ def open(file, mode=None, format=None, options=None,
             _cinit_sentinel, file, format, options,
             container_options, stream_options,
             metadata_encoding, metadata_errors,
-            buffer_size
+            buffer_size, hwaccel=hwaccel,
         )
     raise ValueError("mode must be 'r' or 'w'; got %r" % mode)

av/stream.pyx

@@ -85,7 +85,7 @@ cdef class Stream(object):
         else:
             self._codec = self._codec_context.codec
 
-        self.codec_context = wrap_codec_context(self._codec_context, self._codec, False)
+        self.codec_context = wrap_codec_context(self._codec_context, self._codec, False, container.hwaccel)
         self.codec_context.stream_index = stream.index
 
     def __repr__(self):

av/video/codeccontext.pxd

@@ -1,4 +1,6 @@
 
+cimport libav as lib
+
 from av.codec.context cimport CodecContext
 from av.video.format cimport VideoFormat
 from av.video.frame cimport VideoFrame
@@ -19,3 +21,9 @@ cdef class VideoCodecContext(CodecContext):
 
     # For decoding.
     cdef VideoFrame next_frame
+
+    # For hardware acceleration
+    cdef dict hwaccel
+    cdef lib.AVPixelFormat hw_pix_fmt
+    cdef lib.AVBufferRef* hw_device_ctx
+    cdef bint _setup_hw_decoder(self, lib.AVCodec *codec)

av/video/codeccontext.pyx

@@ -12,17 +12,82 @@ from av.video.frame cimport VideoFrame, alloc_video_frame
 from av.video.reformatter cimport VideoReformatter
 
 
+cdef lib.AVPixelFormat _get_hw_format(lib.AVCodecContext *ctx, lib.AVPixelFormat *pix_fmts):
+    i = 0
+    while pix_fmts[i] != -1:
+        if pix_fmts[i] == ctx.pix_fmt:
+            return pix_fmts[i]
+        i += 1
+
+    return lib.AV_PIX_FMT_NONE
+
+
 cdef class VideoCodecContext(CodecContext):
 
     def __cinit__(self, *args, **kwargs):
         self.last_w = 0
         self.last_h = 0
 
+        self.hw_pix_fmt = lib.AV_PIX_FMT_NONE
+        self.hw_device_ctx = NULL
+        self.hwaccel = kwargs.get("hwaccel", None)
+
     cdef _init(self, lib.AVCodecContext *ptr, const lib.AVCodec *codec):
         CodecContext._init(self, ptr, codec)  # TODO: Can this be `super`?
+
+        if self.hwaccel is not None:
+            self._setup_hw_decoder(codec)
+
         self._build_format()
         self.encoded_frame_count = 0
 
+    cdef bint _setup_hw_decoder(self, lib.AVCodec *codec):
+        # Get device type
+        device_type = lib.av_hwdevice_find_type_by_name(self.hwaccel["device_type_name"])
+        if device_type == lib.AV_HWDEVICE_TYPE_NONE:
+            raise ValueError("Device type {} is not supported.".format(self.hwaccel["device_type_name"]))
+
+        # Check that decoder is supported by this device
+        i = 0
+        while True:
+            config = lib.avcodec_get_hw_config(codec, i)
+
+            # Exhausted list
+            if not config:
+                break
+
+            if config.methods & lib.AV_CODEC_HW_CONFIG_METHOD_HW_DEVICE_CTX and config.device_type == device_type:
+                self.hw_pix_fmt = config.pix_fmt
+                break
+
+            i += 1
+
+        # Decoder is not supported by the desired device
+        if self.hw_pix_fmt == lib.AV_PIX_FMT_NONE:
+            return False
+
+        # Override the decoder context's get_format function
+        self.ptr.pix_fmt = self.hw_pix_fmt
+        self.ptr.get_format = _get_hw_format
+
+        # Create the hardware device context
+        cdef char* device = NULL
+        if "device" in self.hwaccel:
+            device_bytes = self.hwaccel["device"].encode()
+            device = device_bytes
+
+        err = lib.av_hwdevice_ctx_create(&self.hw_device_ctx, device_type, device, NULL, 0)
+        if err < 0:
+            raise RuntimeError("Failed to create specified HW device")
+
+        self.ptr.hw_device_ctx = lib.av_buffer_ref(self.hw_device_ctx)
+
+        return True
+
+    def __dealloc__(self):
+        if self.hw_device_ctx:
+            lib.av_buffer_unref(&self.hw_device_ctx)
+
     cdef _set_default_time_base(self):
         self.ptr.time_base.num = self.ptr.framerate.den or 1
         self.ptr.time_base.den = self.ptr.framerate.num or lib.AV_TIME_BASE
@@ -68,6 +133,24 @@ cdef class VideoCodecContext(CodecContext):
         cdef VideoFrame vframe = frame
         vframe._init_user_attributes()
 
+    cdef _transfer_hwframe(self, Frame frame):
+        cdef Frame frame_sw
+
+        if self.using_hwaccel and frame.ptr.format == self.hw_pix_fmt:
+            # retrieve data from GPU to CPU
+            frame_sw = self._alloc_next_frame()
+
+            ret = lib.av_hwframe_transfer_data(frame_sw.ptr, frame.ptr, 0)
+            if (ret < 0):
+                raise RuntimeError("Error transferring the data to system memory")
+
+            frame_sw.pts = frame.pts
+
+            return frame_sw
+
+        else:
+            return frame
+
     cdef _build_format(self):
         self._format = get_video_format(<lib.AVPixelFormat>self.ptr.pix_fmt, self.ptr.width, self.ptr.height)
 
@@ -162,3 +245,7 @@ cdef class VideoCodecContext(CodecContext):
     property coded_height:
         def __get__(self):
             return self.ptr.coded_height
+
+    property using_hwaccel:
+        def __get__(self):
+            return self.hw_device_ctx != NULL

include/libav.pxd

@@ -12,6 +12,8 @@ include "libavutil/dict.pxd"
 include "libavutil/error.pxd"
 include "libavutil/frame.pxd"
 include "libavutil/samplefmt.pxd"
+include "libavutil/buffer.pxd"
+include "libavutil/hwcontext.pxd"
 
 include "libavcodec/avcodec.pxd"
 include "libavdevice/avdevice.pxd"

include/libavcodec/avcodec.pxd

@@ -176,6 +176,10 @@ cdef extern from "libavcodec/avcodec.pyav.h" nogil:
         int get_buffer(AVCodecContext *ctx, AVFrame *frame)
         void release_buffer(AVCodecContext *ctx, AVFrame *frame)
 
+        # Hardware acceleration
+        AVBufferRef *hw_device_ctx
+        AVPixelFormat (*get_format)(AVCodecContext *s, const AVPixelFormat * fmt)
+
         # User Data
         void *opaque
 
@@ -185,6 +189,20 @@ cdef extern from "libavcodec/avcodec.pyav.h" nogil:
     cdef AVClass* avcodec_get_class()
     cdef int avcodec_copy_context(AVCodecContext *dst, const AVCodecContext *src)
 
+    # Hardware acceleration
+    enum:
+        AV_CODEC_HW_CONFIG_METHOD_HW_DEVICE_CTX = 0x01,
+        AV_CODEC_HW_CONFIG_METHOD_HW_FRAMES_CTX = 0x02,
+        AV_CODEC_HW_CONFIG_METHOD_INTERNAL      = 0x04,
+        AV_CODEC_HW_CONFIG_METHOD_AD_HOC        = 0x08,
+
+    cdef struct AVCodecHWConfig:
+        AVPixelFormat pix_fmt;
+        int methods;
+        AVHWDeviceType device_type;
+
+    cdef const AVCodecHWConfig* avcodec_get_hw_config(const AVCodec *codec, int index)
+
     cdef struct AVCodecDescriptor:
         AVCodecID id
         AVMediaType type

include/libavutil/buffer.pxd

@@ -0,0 +1,7 @@
+cdef extern from "libavutil/buffer.h" nogil:
+    cdef struct AVBuffer
+    cdef struct AVBufferRef
+
+    cdef AVBufferRef* av_buffer_ref(AVBufferRef *buf)
+    cdef void av_buffer_unref(AVBufferRef **buf)
+

include/libavutil/hwcontext.pxd

@@ -0,0 +1,22 @@
+cdef extern from "libavutil/hwcontext.h" nogil:
+
+    enum AVHWDeviceType:
+        AV_HWDEVICE_TYPE_NONE
+        AV_HWDEVICE_TYPE_VDPAU
+        AV_HWDEVICE_TYPE_CUDA
+        AV_HWDEVICE_TYPE_VAAPI
+        AV_HWDEVICE_TYPE_DXVA2
+        AV_HWDEVICE_TYPE_QSV
+        AV_HWDEVICE_TYPE_VIDEOTOOLBOX
+        AV_HWDEVICE_TYPE_D3D11VA
+        AV_HWDEVICE_TYPE_DRM
+        AV_HWDEVICE_TYPE_OPENCL
+        AV_HWDEVICE_TYPE_MEDIACODEC
+
+    cdef int av_hwdevice_ctx_create(AVBufferRef **device_ctx, AVHWDeviceType type, const char *device, AVDictionary *opts, int flags)
+
+    cdef AVHWDeviceType av_hwdevice_find_type_by_name(const char *name)
+    cdef const char *av_hwdevice_get_type_name(AVHWDeviceType type)
+
+    cdef int av_hwframe_transfer_data(AVFrame *dst, const AVFrame *src, int flags)
+

scripts/build-deps

@@ -25,6 +25,32 @@ fi
 
 mkdir -p "$PYAV_LIBRARY_ROOT"
 mkdir -p "$PYAV_LIBRARY_PREFIX"
+
+# Nvidia build
+CONFFLAGS_NVIDIA=""
+if [[ -e /usr/local/cuda ]]; then
+    # Get Nvidia headers for ffmpeg
+    cd $PYAV_LIBRARY_ROOT
+    if [[ ! -e "$PYAV_LIBRARY_ROOT/nv-codec-headers" ]]; then
+        git clone https://github.com/FFmpeg/nv-codec-headers.git
+        cd nv-codec-headers
+        make -j4
+        make PREFIX="$PYAV_LIBRARY_PREFIX" install
+    fi
+
+    PKG_CONFIG_PATH="/usr/local/lib/pkgconfig:$PKG_CONFIG_PATH"
+    CONFFLAGS_NVIDIA="--enable-cuda \
+                      --enable-cuvid \
+                      --enable-nvenc \
+                      --enable-nonfree \
+                      --enable-libnpp \
+                      --extra-cflags=-I/usr/local/cuda/include \
+                      --extra-ldflags=-L/usr/local/cuda/lib64"
+else
+    echo "WARNING: Did not find cuda libraries in /usr/local/cuda..."
+    echo "         Building without hardware acceleration support"
+fi
+
 cd "$PYAV_LIBRARY_ROOT"
 
 
@@ -52,6 +78,7 @@ fi
     --enable-debug=3 \
     --enable-gpl \
     --enable-libx264 \
+    $CONFFLAGS_NVIDIA \
     $CONFFLAGS \
     --prefix="$PYAV_LIBRARY_PREFIX" \
     || exit 2