Skip to content

VideoPlane.__dlpack__() when device is cpu #2217

@Yves33

Description

@Yves33

Hi,
Many thanks for the amazing work, specially for implementing HWaccel and __dlpack__() support.

Actual behavior: VideoPlane.__dlpack__() only works when hwaccel=HWAccel(device_type="cuda").
Expected behavior: __dlpack__() should return a valid pycapsule for any av.VideoPlane when data is on cpu (whatever te number of video planes (rgb, y+uv, y+u+v).

minimal code:

import numpy as np
from ctypes import c_byte,POINTER,cast
import av
from av.codec.hwaccel import HWAccel

## pick you flavor!
hwaccel = HWAccel(device_type="cuda",is_hw_owned=True)  ## works
hwaccel = HWAccel(device_type="cuda",is_hw_owned=False) ## works
hwaccel=None ## fails - requires numpy intermediate array

outcontainer=av.open("test_cpu_dlpack.mp4", "w")
outstream=outcontainer.add_stream("hevc_nvenc",rate=av.utils.Fraction(30000,1001))
# decode using cuvid
with av.open("./video4.mkv", "r", hwaccel=hwaccel) as c:
    for inframe in c.decode(video=0):
        y=inframe.planes[0]
        uv=inframe.planes[1]
        if not hwaccel:
            ## there is no reason why we could not get a __dlpack__() from video planes on the cpu!
            ## presently fails at File "av/video/plane.py", line 178, in av.video.plane.VideoPlane.__dlpack__
            ##    raise NotImplementedError("unsupported sw_format for DLPack export")
            ## the next lines are a workaround,(creates numpy array from ptr) but it's not really readable!
            y=np.ctypeslib.as_array(cast(y.buffer_ptr,POINTER(c_byte)),
                                    shape=(y.height, y.width)).astype(np.uint8)
            uv=np.ctypeslib.as_array(cast(uv.buffer_ptr,POINTER(c_byte)),
                                     shape=(uv.height, uv.width*2)).astype(np.uint8)

        outframe=av.VideoFrame.from_dlpack((y,uv), format='nv12')
        for packet in outstream.encode(outframe):
            outcontainer.mux(packet)
for packet in outstream.encode(None):
    outcontainer.mux(packet)

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions