Keep frame on CUDA device

[jmspereira]

Hey everyone,

I am trying to use the HWAccel feature of pyav, and for that, I am doing the following:

hwaccel = HWAccel(device_type="cuda", allow_software_fallback=False)
codec = av.codec.CodecContext.create("h264", mode="r", hwaccel=hwaccel)

for packet in packets:
    for frame in codec.decode(packet):
           ....

However, I would like to keep the frame in GPU, but even if I do not call the to_ndarray method, the output frame is already in CPU.
I am observing the increase in memory when I put the frames on a list.

Is there any way of keeping the frame on the GPU?
Best,

[jmspereira]

@WyattBlue, from what I am seeing, this is not possible with the current PyAv version.
However, I think that this can be simply implemented by having an option on the HWAccel object (e.g., keep_output_on_device) and here if that option is set, the output should be kept on the CUDA device, right?

[jmspereira]

I was wrong, the change is not enough because then you are not able to access the planes to create a cupy array or a tensor.
I tried to access the pointers of the planes by accessing directly the self.ptr.data[0] and self.ptr.data[1] of the frame (since it is an nv12 frame). However, after creating a cupy array that points to those planes, the frame is not correct, despite not being totally garbage since some parts of the image are correct (the same code works when I revert the change on the _transfer_hwframe method and I use the pointers to the CPU frame).

Do you have any idea of the required changes to be possible to have this working? Or is it even possible?

[materight]

How did you do the conversion from nv12? The planes on the GPU might have a different "pitch" memory layout than the ones on CPU, you can check it using self.ptr.linesize[0] and self.ptr.linesize[1]. I was able to do something similar some time ago, using the NPP library for the conversion (here), but it would be nice to be able to do it natively in pyav.

[jmspereira]

I used a CUDA kernel in cupy to handle that conversion. I could be misunderstanding, but I looked into your library, and you are receiving the frame that AV returns, which is a frame that, at least for the most recent AV version, is already in CPU, no?

To integrate this directly into AV maybe we should use a format like dlpack, which does not depend on any specific library (like pytorch or cupy, however, they can create tensors from dlpacks). Unfortunately, I do not understand enough of cython or cuda memory management to open a PR with this 😞

[materight]

My library was written for an old version of pyav, when hw acceleration was not supported yet, so the frames are never transferred to cpu. Yeah dlpack would be interesting to try, I might take a look when I have some time.

[jmspereira]

Thank you! Please let me know if you need any help