try to convert 8/10 bits nv12 to 8bit jpg file

try to convert 8/10 bits nv12 to 8 bit jpg file

import cv2
import numpy as np

def yuv2bgr(filename,height,width,startfrm,bit_depth):
    if bit_depth=="8bits":
        bpp=1; dtype='uint8'
    else:
        bpp=2; dtype='uint16'

    with open(filename,'rb') as fp:
        frame_len=height*width*bpp*3//2

        fp.seek(0,2)
        ps=fp.tell()
        numfrm=ps//frame_len

        fp.seek(0,0)
        for i in range(numfrm):
            raw = fp.read(frame_len)
            yuv = np.frombuffer(raw, dtype=dtype)
            yuv = yuv.reshape((height*3//2, width))

            if bpp==1:
                bgr_img = cv2.cvtColor(yuv, cv2.COLOR_YUV2BGR_NV12)
            else:
                yuv_tmp = yuv * 255.0 / 1023
                yuv_uint8=yuv_tmp.astype(np.uint8)
                bgr_img = cv2.cvtColor(yuv_uint8, cv2.COLOR_YUV2BGR_NV12)

            cv2.imwrite(r'%s_yuv2bgr_%d.jpg' % (bit_depth, i + 1), bgr_img)
            print("Extract frame %d" % (i+1))


if __name__=='__main__':
    path=r"sample_1.nv12"
    width=640;height=480
    bit_depth = "8bits"

    # path=r"sample_2.nv12"
    # width = 640;height = 480
    # bit_depth="10bits"

    startfrm=0
    _=yuv2bgr(filename=path,height=height,width=width,startfrm=0,bit_depth=bit_depth)