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
    No fields configured for issues without a type.

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions