yuv420sp转jpg

代码功能：将yuv420sp（NV12）数据转换成jpg图片
需要注意：文件名需包含图片宽高信息，如 ‘00001_1920_1080.yuv420sp’

import os
import cv2
import numpy as np
from tqdm import tqdm


def yuv420sp2bgr(path, w, h):
    yuv = np.zeros((h, w, 3), dtype=np.uint8)
    u = np.zeros((int(h / 2), int(w / 2)), dtype=np.uint8)
    v = np.zeros((int(h / 2), int(w / 2)), dtype=np.uint8)
    with open(path, 'rb') as f:
        for row in range(h):
            for col in range(w):
                yuv[row, col, 0] = int(f.read(1).hex(), 16)
        for row in range(int(h / 2)):
            for col in range(int(w / 2)):
                u[row, col] = int(f.read(1).hex(), 16)
                v[row, col] = int(f.read(1).hex(), 16)

    yuv[:, :, 1] = cv2.resize(u, (w, h))
    yuv[:, :, 2] = cv2.resize(v, (w, h))
    bgr = cv2.cvtColor(yuv, cv2.COLOR_YUV2BGR)
    return bgr


# yuv420sp文件名需包含图片宽高信息，如 '00001_1920_1080.yuv420sp'
# yuv420sp文件是yuv420图像存储格式的数据位: yyyyyy...uvuv..
path_yuv420sp = r'/media/fang/3441-913C/tsr/img'  # yuv420sp文件路径
path_output = path_yuv420sp  # jpg图片输出路径，不要有中文

for root, dirs, files in os.walk(path_yuv420sp):
    for file in tqdm(files):
        if file.split(".")[-1] == 'yuv420sp':
            w, h = int(file.split(".")[0].split('_')[-2]), int(file.split(".")[0].split('_')[-1])
            bgr = yuv420sp2bgr(root + '/' + file, w, h)
            cv2.imwrite(path_output + '/' + file.replace('yuv420sp', 'jpg'), bgr)