diff --git a/av/sidedata/encparams.pyx b/av/sidedata/encparams.py similarity index 60% rename from av/sidedata/encparams.pyx rename to av/sidedata/encparams.py index 380dc53d4..399189a6b 100644 --- a/av/sidedata/encparams.pyx +++ b/av/sidedata/encparams.py @@ -1,19 +1,20 @@ -cimport libav as lib -from libc.stdint cimport int32_t, uint8_t - from enum import IntEnum -VideoEncParamsType = IntEnum( - "AVVideoEncParamsType", - { - "NONE": lib.AV_VIDEO_ENC_PARAMS_NONE, - "VP9": lib.AV_VIDEO_ENC_PARAMS_VP9, - "H264": lib.AV_VIDEO_ENC_PARAMS_H264, - "MPEG2": lib.AV_VIDEO_ENC_PARAMS_MPEG2, - }, -) - -cdef class VideoEncParams(SideData): +import cython +from cython.cimports import libav as lib +from cython.cimports.av.sidedata.sidedata import SideData +from cython.cimports.libc.stdint import uint8_t + + +class VideoEncParamsType(IntEnum): + NONE = lib.AV_VIDEO_ENC_PARAMS_NONE + VP9 = lib.AV_VIDEO_ENC_PARAMS_VP9 + H264 = lib.AV_VIDEO_ENC_PARAMS_H264 + MPEG2 = lib.AV_VIDEO_ENC_PARAMS_MPEG2 + + +@cython.cclass +class VideoEncParams(SideData): def __repr__(self): return f"" @@ -25,29 +26,37 @@ def nb_blocks(self): the values of blocks_offset / block_size are unspecified and should not be accessed. """ - return ( self.ptr.data).nb_blocks + return cython.cast( + cython.pointer[lib.AVVideoEncParams], self.ptr.data + ).nb_blocks @property def blocks_offset(self): """ Offset in bytes from the beginning of this structure at which the array of blocks starts. """ - return ( self.ptr.data).blocks_offset + return cython.cast( + cython.pointer[lib.AVVideoEncParams], self.ptr.data + ).blocks_offset @property def block_size(self): """ Size of each block in bytes. May not match sizeof(AVVideoBlockParams). """ - return ( self.ptr.data).block_size + return cython.cast( + cython.pointer[lib.AVVideoEncParams], self.ptr.data + ).block_size @property def codec_type(self): """ Type of the parameters (the codec they are used with). """ - cdef lib.AVVideoEncParamsType t = ( self.ptr.data).type - return VideoEncParamsType(t) + t: lib.AVVideoEncParamsType = cython.cast( + cython.pointer[lib.AVVideoEncParams], self.ptr.data + ).type + return VideoEncParamsType(cython.cast(cython.int, t)) @property def qp(self): @@ -57,15 +66,17 @@ def qp(self): combined with `delta_qp` and the per-block delta in a manner documented for each type. """ - return ( self.ptr.data).qp - + return cython.cast(cython.pointer[lib.AVVideoEncParams], self.ptr.data).qp + @property def delta_qp(self): """ Quantisation parameter offset from the base (per-frame) qp for a given plane (first index) and AC/DC coefficients (second index). """ - cdef lib.AVVideoEncParams *p = self.ptr.data + p: cython.pointer[lib.AVVideoEncParams] = cython.cast( + cython.pointer[lib.AVVideoEncParams], self.ptr.data + ) return [[p.delta_qp[i][j] for j in range(2)] for i in range(4)] def block_params(self, idx): @@ -85,25 +96,28 @@ def qp_map(self): """ import numpy as np - cdef int mb_h = (self.frame.ptr.height + 15) // 16 - cdef int mb_w = (self.frame.ptr.width + 15) // 16 - cdef int nb_mb = mb_h * mb_w - cdef int block_idx - cdef int y - cdef int x - cdef VideoBlockParams block + mb_h: cython.int = (self.frame.ptr.height + 15) // 16 + mb_w: cython.int = (self.frame.ptr.width + 15) // 16 + nb_mb: cython.int = mb_h * mb_w + block_idx, x, y = cython.declare(cython.int) + block: VideoBlockParams - # Validate number of blocks if self.nb_blocks != nb_mb: - raise RuntimeError("Expected frame size to match number of blocks in side data") - - # Validate type - cdef lib.AVVideoEncParamsType type = ( self.ptr.data).type - if type != lib.AVVideoEncParamsType.AV_VIDEO_ENC_PARAMS_MPEG2 and type != lib.AVVideoEncParamsType.AV_VIDEO_ENC_PARAMS_H264: + raise RuntimeError( + "Expected frame size to match number of blocks in side data" + ) + + type: lib.AVVideoEncParamsType = cython.cast( + cython.pointer[lib.AVVideoEncParams], self.ptr.data + ).type + if ( + type != lib.AVVideoEncParamsType.AV_VIDEO_ENC_PARAMS_MPEG2 + and type != lib.AVVideoEncParamsType.AV_VIDEO_ENC_PARAMS_H264 + ): raise ValueError("Expected MPEG2 or H264") # Create a 2-D map with the number of macroblocks - cdef int32_t[:, ::1] map = np.empty((mb_h, mb_w), dtype=np.int32) + map = np.empty((mb_h, mb_w), dtype=np.int32) # Fill map with quantization parameter per macroblock for block_idx in range(nb_mb): @@ -112,14 +126,19 @@ def qp_map(self): x = block.src_x // 16 map[y, x] = self.qp + block.delta_qp - return np.asarray(map) + return map -cdef class VideoBlockParams: - def __init__(self, VideoEncParams video_enc_params, int idx) -> None: - cdef uint8_t* base = video_enc_params.ptr.data - cdef Py_ssize_t offset = video_enc_params.blocks_offset + idx * video_enc_params.block_size - self.ptr = (base + offset) +@cython.cclass +class VideoBlockParams: + def __init__(self, video_enc_params: VideoEncParams, idx: cython.int) -> None: + base: cython.pointer[uint8_t] = cython.cast( + cython.pointer[uint8_t], video_enc_params.ptr.data + ) + offset: cython.Py_ssize_t = ( + video_enc_params.blocks_offset + idx * video_enc_params.block_size + ) + self.ptr = cython.cast(cython.pointer[lib.AVVideoBlockParams], base + offset) def __repr__(self): return f""