python YUV 转RGB

一、将base64的图片数据转为RGB (速度很慢)

因为使用的Python内置的for循环遍历图片数据，所以速度会比较慢，转换一张200k左右的图片需要2s到3s左右。

#encoding:utf-8
import numpy as np
import cv2


video_cutoff = [0.0625, 0.5, 0.5]
full_cutoff = [0.0, 0.5, 0.5]


def rgb2yuv_cutoff(rgb_row, p):
    return [(a + p) for a in rgb_row]


def rgb2yuv_convert(rgb, format, cutoff):
    s = np.array([rgb2yuv_cutoff(np.array(np.dot(n / 255.0, format)), cutoff)
                  for n in rgb]);
    return np.clip(s, 0, 1.0)


yuv2rgb_601 = np.matrix([[1.164, 1.164, 1.164],
                         [0.0, -0.392, 2.017],
                         [1.596, -0.813, 0.0]])
yuv2rgb_709 = np.matrix([[1.164, 1.164, 1.164],
                         [0.0, -0.213, 2.112],
                         [1.787, -0.531, 0.0]])
yuv2rgb_601f = np.matrix([[1, 1, 1],
                          [0.0, -0.343, 1.765],
                          [1.4, -0.711, 0.0]])
rgb2yuv_601f = yuv2rgb_601f.I


def rgb2yuv(rgb):
    return (np.dot(rgb, rgb2yuv_601f) + full_cutoff).A1


# RGB Array编码为NV21字节序列，可直接写入文件
def rgbDataToNV21Bytes(rgbData):
    if type(rgbData) is not np.ndarray:
        raise NameError("bad type!")

        # width 和 height 必须 >= 2
    height = rgbData.shape[0]
    width = rgbData.shape[1]
    if rgbData.shape[2] != 3:
        raise NameError("only suporse RGB image!")

        # 写入缓冲区
    bytes = np.zeros(int(width * height * 1.5))

    yuvData = np.array(np.dot(rgbData.reshape((width * height, 3)) / 255.0, rgb2yuv_601f) + full_cutoff).reshape(
        (height, width, 3)) * 255

    offUV = width * height
    for y in range(height):
        for x in range(width):
            offY = y * width + x

            yuvb = yuvData[y][x]
            bytes[offY] = yuvb[0]

            if y % 2 == 0 and offY % 2 == 0:
                bytes[offUV] = yuvb[2]
                bytes[offUV + 1] = yuvb[1]
                offUV += 2
    bytes = np.clip(bytes, 0, 255)

    return bytes.astype('uint8').tostring()


# nv21 YVU 速度很慢
def nv21BytesToRGBdata(img_str, width, height):
    if type(img_str) is not str:
        raise NameError("bad type!")
    src = np.fromstring(img_str, np.uint8)

    numOfPixel = width * height
    positionOfV = numOfPixel
    yuvData = np.zeros((height, width, 3))

    bytes = np.array([src[0], src[positionOfV + 1], src[positionOfV]])
    yuv = bytes / 255.0 - [0, 0.5, 0.5]
    for i in range(height):
        startY = i * width
        step = i / 2 * width
        startV = positionOfV + step
        for j in range(width):
            Y = startY + j
            V = startV + j / 2 * 2
            U = V + 1
            index = Y * 3
            bytes = np.array([src[Y], src[U], src[V]])
            yuv = bytes / 255.0 - full_cutoff
            yuvData[i][j] = (np.dot(yuv, yuv2rgb_601f)) * 255

    yuvData = np.clip(yuvData, 0, 255)

    return yuvData.astype('uint8')

if __name__ == "__main__":

    import base64,time
    #文件中保存的是base64的图片数据
    with open("/root/image/text4.txt") as f:
        contents = ""
        lines = f.readlines()
        for line in lines:
            contents += line
    start_time = time.time()
    img_list = str(contents).split(",")
    img_data = base64.b64decode(img_list[1])
    img = nv21BytesToRGBdata(img_data,640,360)


 

 

二、将base64的图片数据转为RGB (速度很快)

主要利用numpy来实现的，所以速度上会比内置的for循环快很多，200k的图片大概只要零点几秒。

    import numpy as np
    import cv2


    def YUVtoRGB(byteArray,width,height):

        e = width * height
        Y = byteArray[0:e]
        Y = np.reshape(Y, (height, width))

        s = e
        V = byteArray[s::2]
        V = np.repeat(V, 2, 0)
        V = np.reshape(V, (height/2, width))
        V = np.repeat(V, 2, 0)

        U = byteArray[s + 1::2]
        U = np.repeat(U, 2, 0)
        U = np.reshape(U, (height/2, width))
        U = np.repeat(U, 2, 0)

        RGBMatrix = (np.dstack([Y, U, V])).astype(np.uint8)
        RGBMatrix = cv2.cvtColor(RGBMatrix, cv2.COLOR_YUV2RGB, 3)
        return RGBMatrix

    import base64,time
    #文件中保存的是base64的图片数据
    with open("/root/image/text4.txt") as f:
        contents = ""
        lines = f.readlines()
        for line in lines:
            contents += line
    start_time = time.time()
    img_list = str(contents).split(",")
    img_data = base64.b64decode(img_list[1])
    img = nv21BytesToRGBdata(img_data,640,360)
    nparr = np.fromstring(img_data, np.uint8)

    # img_np = cv2.imdecode(nparr,1)
    # print(img_np)
    img_np = YUVtoRGB(nparr,640,360)
    cv2.imwrite("/root/image/text4.jpg",img_np)
    print(time.time() - start_time)