PyNvVideoCodec解码使用

解码视频生成器例子，支持解码视频和视频流

try:
    import time
    from functools import wraps
    import cv2
    import numpy as np
    import pycuda.driver as cuda
    import PyNvVideoCodec as nvc
    import ctypes as C
    import os
except ImportError as e:
    pass
    
def cast_address_to_1d_bytearray(base_address, size):
    return np.ctypeslib.as_array(C.cast(base_address, C.POINTER(C.c_uint8)),
                                 shape=(size,))

def read_frames2(video_name, gpu_id=0, use_device_memory=False):
    nv_dmx = nvc.CreateDemuxer(filename=video_name)
    nv_dec = nvc.CreateDecoder(gpuid=gpu_id,
                               codec=nv_dmx.GetNvCodecId(),
                               cudacontext=0,
                               cudastream=0,
                               usedevicememory=use_device_memory)

    raw_frame = None

    seq_triggered = False
    # printing out FPS and pixel format of the stream for convenience
    print("FPS = ", nv_dmx.FrameRate())

    width = 0
    height = 0

    start = time.perf_counter()
    num_decoded_frames = 0

    try:
        # demuxer can be iterated, fetch the packet from demuxer
        for packet in nv_dmx:
            # Decode returns a list of packets, range of this list is from [0, size of (decode picture buffer)]
            # size of (decode picture buffer) depends on GPU, fur Turing series its 8
            for decoded_frame in nv_dec.Decode(packet):
                # 'decoded_frame' contains list of views implementing cuda array interface
                # for nv12, it would contain 2 views for each plane and two planes would be contiguous

                start_per = time.time()
                num_decoded_frames += 1

                if not seq_triggered:
                    decoded_frame_size = nv_dec.GetFrameSize()
                    print(f"decoded_frame_size={decoded_frame_size}")
                    width = nv_dec.GetWidth()
                    height = nv_dec.GetHeight()
                    print("width=", width)
                    print("height=", height)
                    raw_frame = np.ndarray(shape=decoded_frame_size, dtype=np.uint8)
                    print(f"raw_frame.shape={raw_frame.shape}")
                    seq_triggered = True

                luma_base_addr = decoded_frame.GetPtrToPlane(0)
                if use_device_memory:
                    cuda.memcpy_dtoh(raw_frame, luma_base_addr)
                    bits = bytearray(raw_frame)
                    # print(type(bits))
                    # np_arr = np.array(bits)
                    np_arr = np.frombuffer(bits, dtype=np.uint8)

                    bgr_frame = cv2.cvtColor(np_arr.reshape((height * 3 // 2, width)), cv2.COLOR_YUV2BGR_NV12)

                else:
                    new_array = cast_address_to_1d_bytearray(base_address=luma_base_addr,
                                                             size=decoded_frame.framesize())

                    bgr_frame = cv2.cvtColor(new_array.reshape((height * 3 // 2, width)), cv2.COLOR_YUV2BGR_NV12)

                yield bgr_frame

                elapsed_time = time.time() - start_per
                # print(f"NV解码一帧时间: 0.0599s")
    finally:
        final_elapsed = time.perf_counter() - start
        print(f"PyNvVideoCodec average per frame cost time: {final_elapsed / num_decoded_frames * 1000}ms")

def main():
    video_dir = "./video"
    vid = f"{video_dir}/2c1080p.mp4"
    for i, frame in enumerate(read_frames2(vid)):
        pass
       
if __name__ == '__main__':
    main()