ffmpeg-python icon indicating copy to clipboard operation
ffmpeg-python copied to clipboard

How to read video frames into a numpy array, using GPU ?

Open uhtmilkshake opened this issue 3 years ago • 1 comments

I'm using this example here to grab frames from a video and put them into a numpy array, but in order to use GPU, one has to specify the nvenc decoder and that changes the output buffer size, which makes it impossible to reshape it properly. For example for a small video, buffer goes from 1272499200 in size to 35578084 in size Here is what I have currently, which does not work, I get a "cannot reshape array of size Xinto shape (height,width,3)" error. (which makes sense, math doesn't result in an integer anymore) Looking for any advice on how to deal with or avoid this problem. Even a completely different approach would be fine. Thanks!

input_args = {
    "hwaccel": "nvdec",
    "vcodec": "h264_cuvid",
    "c:v": "h264_cuvid"
}

output_args = {
    "vcodec": "hevc_nvenc",
    "c:v": "hevc_nvenc",
    "preset": "fast",
}

probe = ffmpeg.probe("sample2.mp4")
video_info = next(s for s in probe['streams'] if s['codec_type'] == 'video')
width = int(video_info['width'])
height = int(video_info['height'])

out, err = (
    ffmpeg
    .input('sample2.mp4',**input_args)
    .output('pipe:', format='rawvideo', pix_fmt='rgb24',**output_args)
    .run(capture_stdout=True)
)
video = np.frombuffer(out, np.uint8).reshape([-1, height, width, 3])

uhtmilkshake avatar Sep 19 '22 14:09 uhtmilkshake

Hey @sensitivecereal, I was actually trying to figure out how to use the nvenc decoder while reading the raw data for my project. I think I've got it working. I believe you just need to remove your output_args. Since you're wanting the raw data I don't think you should specify those args since they're telling ffmpeg how to encode the data.

Also, I'd recommend taking a look at this example in the repo - I found it really helpful. It processes one frame at a time so you won't use up all your memory. 😄

lawrencefoley avatar Oct 13 '22 05:10 lawrencefoley