305 lines
10 KiB
Python
305 lines
10 KiB
Python
import cython
|
|
import cython.cimports.libav as lib
|
|
from cython.cimports.av.error import err_check
|
|
from cython.cimports.av.video.format import (
|
|
VideoFormatComponent,
|
|
get_pix_fmt,
|
|
get_video_format,
|
|
)
|
|
from cython.cimports.av.video.frame import VideoFrame
|
|
from cython.cimports.cpython import PyBUF_WRITABLE, PyBuffer_FillInfo
|
|
from cython.cimports.cpython.buffer import Py_buffer
|
|
from cython.cimports.cpython.pycapsule import (
|
|
PyCapsule_GetPointer,
|
|
PyCapsule_IsValid,
|
|
PyCapsule_New,
|
|
)
|
|
from cython.cimports.libc.stdlib import free, malloc
|
|
|
|
|
|
@cython.final
|
|
@cython.cclass
|
|
class VideoPlane(Plane):
|
|
def __cinit__(self, frame: VideoFrame, index: cython.int):
|
|
# The palette plane has no associated component or linesize; set fields manually
|
|
fmt = frame.format
|
|
if frame.ptr.hw_frames_ctx:
|
|
frames_ctx: cython.pointer[lib.AVHWFramesContext] = cython.cast(
|
|
cython.pointer[lib.AVHWFramesContext], frame.ptr.hw_frames_ctx.data
|
|
)
|
|
fmt = get_video_format(
|
|
frames_ctx.sw_format, frame.ptr.width, frame.ptr.height
|
|
)
|
|
|
|
if index == 1 and fmt.name == "pal8":
|
|
self.width = 256
|
|
self.height = 1
|
|
self.buffer_size = 256 * 4
|
|
return
|
|
|
|
for i in range(fmt.ptr.nb_components):
|
|
if fmt.ptr.comp[i].plane == index:
|
|
component = VideoFormatComponent(fmt, i)
|
|
self.width = component.width
|
|
self.height = component.height
|
|
break
|
|
else:
|
|
raise RuntimeError(f"could not find plane {index} of {fmt!r}")
|
|
|
|
# Sometimes, linesize is negative (and that is meaningful). We are only
|
|
# insisting that the buffer size be based on the extent of linesize, and
|
|
# ignore it's direction.
|
|
self.buffer_size = abs(self.frame.ptr.linesize[self.index]) * self.height
|
|
|
|
@cython.cfunc
|
|
def _buffer_size(self) -> cython.size_t:
|
|
return self.buffer_size
|
|
|
|
@property
|
|
def line_size(self):
|
|
"""
|
|
Bytes per horizontal line in this plane.
|
|
|
|
:type: int
|
|
"""
|
|
return self.frame.ptr.linesize[self.index]
|
|
|
|
@cython.cfunc
|
|
def _buffer_writable(self) -> cython.bint:
|
|
if self.frame.ptr.hw_frames_ctx:
|
|
return False
|
|
return True
|
|
|
|
def __getbuffer__(self, view: cython.pointer[Py_buffer], flags: cython.int):
|
|
if self.frame.ptr.hw_frames_ctx:
|
|
raise TypeError(
|
|
"Hardware frame planes do not support the Python buffer protocol. "
|
|
"Use DLPack (__dlpack__) or download to a software frame."
|
|
)
|
|
if flags & PyBUF_WRITABLE and not self._buffer_writable():
|
|
raise ValueError("buffer is not writable")
|
|
|
|
ptr: cython.p_void = self._buffer_ptr()
|
|
line_size: cython.int = self.frame.ptr.linesize[self.index]
|
|
if line_size < 0:
|
|
height: cython.int = self.height
|
|
ptr = cython.cast(cython.p_char, ptr) + (height - 1) * line_size
|
|
|
|
PyBuffer_FillInfo(view, self, ptr, self._buffer_size(), 0, flags)
|
|
|
|
def __dlpack_device__(self):
|
|
if self.frame.ptr.hw_frames_ctx:
|
|
if cython.cast(lib.AVPixelFormat, self.frame.ptr.format) != get_pix_fmt(
|
|
b"cuda"
|
|
):
|
|
raise NotImplementedError(
|
|
"DLPack export is only implemented for CUDA hw frames"
|
|
)
|
|
return (kCuda, self.frame._device_id)
|
|
return (kCPU, 0)
|
|
|
|
def __dlpack__(self, *, stream: int | None = None):
|
|
if self.frame.ptr.buf[0] == cython.NULL:
|
|
raise TypeError(
|
|
"DLPack export requires a refcounted AVFrame (frame.buf[0] is NULL)"
|
|
)
|
|
|
|
sw_fmt: lib.AVPixelFormat
|
|
device_type: cython.int
|
|
device_id: cython.int
|
|
|
|
if self.frame.ptr.hw_frames_ctx:
|
|
if cython.cast(lib.AVPixelFormat, self.frame.ptr.format) != get_pix_fmt(
|
|
b"cuda"
|
|
):
|
|
raise NotImplementedError(
|
|
"DLPack export is only implemented for CUDA hw frames"
|
|
)
|
|
if stream is not None:
|
|
raise NotImplementedError(
|
|
"CUDA stream synchronization is not supported. "
|
|
"Pass stream=None and synchronize before calling __dlpack__."
|
|
)
|
|
|
|
frames_ctx: cython.pointer[lib.AVHWFramesContext] = cython.cast(
|
|
cython.pointer[lib.AVHWFramesContext], self.frame.ptr.hw_frames_ctx.data
|
|
)
|
|
sw_fmt = frames_ctx.sw_format
|
|
device_type = kCuda
|
|
device_id = self.frame._device_id
|
|
else:
|
|
sw_fmt = cython.cast(lib.AVPixelFormat, self.frame.ptr.format)
|
|
device_type = kCPU
|
|
device_id = 0
|
|
|
|
line_size = self.line_size
|
|
if line_size < 0:
|
|
raise NotImplementedError(
|
|
"negative linesize is not supported for DLPack export"
|
|
)
|
|
|
|
nv12 = get_pix_fmt(b"nv12")
|
|
p010le = get_pix_fmt(b"p010le")
|
|
p016le = get_pix_fmt(b"p016le")
|
|
|
|
ndim, bits, itemsize = cython.declare(cython.int)
|
|
s0, s1, s2 = cython.declare(int64_t)
|
|
st0, st1, st2 = cython.declare(int64_t)
|
|
|
|
if sw_fmt == nv12:
|
|
itemsize = 1
|
|
bits = 8
|
|
if self.index == 0:
|
|
ndim = 2
|
|
s0 = self.frame.ptr.height
|
|
s1 = self.frame.ptr.width
|
|
st0 = line_size
|
|
st1 = 1
|
|
elif self.index == 1:
|
|
ndim = 3
|
|
s0 = self.frame.ptr.height // 2
|
|
s1 = self.frame.ptr.width // 2
|
|
s2 = 2
|
|
st0 = line_size
|
|
st1 = 2
|
|
st2 = 1
|
|
else:
|
|
raise ValueError("invalid plane index for NV12")
|
|
elif sw_fmt == p010le or sw_fmt == p016le:
|
|
itemsize = 2
|
|
bits = 16
|
|
if line_size % itemsize:
|
|
raise ValueError("linesize is not aligned to dtype")
|
|
if self.index == 0:
|
|
ndim = 2
|
|
s0 = self.frame.ptr.height
|
|
s1 = self.frame.ptr.width
|
|
st0 = line_size // itemsize
|
|
st1 = 1
|
|
elif self.index == 1:
|
|
ndim = 3
|
|
s0 = self.frame.ptr.height // 2
|
|
s1 = self.frame.ptr.width // 2
|
|
s2 = 2
|
|
st0 = line_size // itemsize
|
|
st1 = 2
|
|
st2 = 1
|
|
else:
|
|
raise ValueError("invalid plane index for P010/P016")
|
|
else:
|
|
raise NotImplementedError("unsupported sw_format for DLPack export")
|
|
|
|
frame_ref: cython.pointer[lib.AVFrame] = lib.av_frame_alloc()
|
|
if frame_ref == cython.NULL:
|
|
raise MemoryError("av_frame_alloc() failed")
|
|
err_check(lib.av_frame_ref(frame_ref, self.frame.ptr))
|
|
|
|
shape = cython.cast(
|
|
cython.pointer[int64_t], malloc(ndim * cython.sizeof(int64_t))
|
|
)
|
|
strides = cython.cast(
|
|
cython.pointer[int64_t], malloc(ndim * cython.sizeof(int64_t))
|
|
)
|
|
if shape == cython.NULL or strides == cython.NULL:
|
|
if shape != cython.NULL:
|
|
free(shape)
|
|
if strides != cython.NULL:
|
|
free(strides)
|
|
lib.av_frame_free(cython.address(frame_ref))
|
|
raise MemoryError("malloc() failed")
|
|
|
|
if ndim == 2:
|
|
shape[0] = s0
|
|
shape[1] = s1
|
|
strides[0] = st0
|
|
strides[1] = st1
|
|
else:
|
|
shape[0] = s0
|
|
shape[1] = s1
|
|
shape[2] = s2
|
|
strides[0] = st0
|
|
strides[1] = st1
|
|
strides[2] = st2
|
|
|
|
ctx = cython.cast(
|
|
cython.pointer[cython.p_void], malloc(3 * cython.sizeof(cython.p_void))
|
|
)
|
|
if ctx == cython.NULL:
|
|
free(shape)
|
|
free(strides)
|
|
lib.av_frame_free(cython.address(frame_ref))
|
|
raise MemoryError("malloc() failed")
|
|
|
|
ctx[0] = cython.cast(cython.p_void, frame_ref)
|
|
ctx[1] = cython.cast(cython.p_void, shape)
|
|
ctx[2] = cython.cast(cython.p_void, strides)
|
|
|
|
managed = cython.cast(
|
|
cython.pointer[DLManagedTensor], malloc(cython.sizeof(DLManagedTensor))
|
|
)
|
|
if managed == cython.NULL:
|
|
free(ctx)
|
|
free(shape)
|
|
free(strides)
|
|
lib.av_frame_free(cython.address(frame_ref))
|
|
raise MemoryError("malloc() failed")
|
|
|
|
managed.dl_tensor.data = cython.cast(cython.p_void, frame_ref.data[self.index])
|
|
managed.dl_tensor.device_type = device_type
|
|
managed.dl_tensor.device_id = device_id
|
|
managed.dl_tensor.ndim = ndim
|
|
managed.dl_tensor.dtype = DLDataType(code=1, bits=bits, lanes=1)
|
|
managed.dl_tensor.shape = shape
|
|
managed.dl_tensor.strides = strides
|
|
managed.dl_tensor.byte_offset = 0
|
|
managed.manager_ctx = cython.cast(cython.p_void, ctx)
|
|
managed.deleter = _dlpack_managed_tensor_deleter
|
|
|
|
try:
|
|
capsule = PyCapsule_New(
|
|
cython.cast(cython.p_void, managed),
|
|
b"dltensor",
|
|
_dlpack_capsule_destructor,
|
|
)
|
|
except Exception:
|
|
_dlpack_managed_tensor_deleter(managed)
|
|
raise
|
|
|
|
return capsule
|
|
|
|
|
|
@cython.cfunc
|
|
@cython.nogil
|
|
@cython.exceptval(check=False)
|
|
def _dlpack_managed_tensor_deleter(
|
|
managed: cython.pointer[DLManagedTensor],
|
|
) -> cython.void:
|
|
if managed == cython.NULL:
|
|
return
|
|
ctx = cython.cast(cython.pointer[cython.p_void], managed.manager_ctx)
|
|
if ctx != cython.NULL:
|
|
frame_ref = cython.cast(cython.pointer[lib.AVFrame], ctx[0])
|
|
shape = cython.cast(cython.pointer[int64_t], ctx[1])
|
|
strides = cython.cast(cython.pointer[int64_t], ctx[2])
|
|
|
|
if frame_ref != cython.NULL:
|
|
lib.av_frame_free(cython.address(frame_ref))
|
|
if shape != cython.NULL:
|
|
free(shape)
|
|
if strides != cython.NULL:
|
|
free(strides)
|
|
free(ctx)
|
|
|
|
free(managed)
|
|
|
|
|
|
@cython.cfunc
|
|
@cython.exceptval(check=False)
|
|
def _dlpack_capsule_destructor(capsule: object) -> cython.void:
|
|
if PyCapsule_IsValid(capsule, b"dltensor"):
|
|
managed = cython.cast(
|
|
cython.pointer[DLManagedTensor],
|
|
PyCapsule_GetPointer(capsule, b"dltensor"),
|
|
)
|
|
if managed != cython.NULL:
|
|
managed.deleter(managed)
|