利用Python将YUV格式文件批量转为BMP保存相对应文件夹下

利用python的tkinter模块选择要转化的yuv格式文件夹，遍历文件并转化格式，如果直接从yuv转rgb，cv2.COLOR_YUV2RGB_Y422会遇到通道不对的情况,所以在此先转NV12格式，再转RGB格式，此处定义三个函数，但这边用来yuv转nv12和nv12转rgb函数，完整代码（含调试注释）如下：

import numpy as np
import cv2
from PIL import Image
import os
import tkinter as tk
from tkinter import filedialog

root=tk.Tk()
root.withdraw()


# 指定输入YUV文件夹和输出BMP文件夹
input_folder = filedialog.askdirectory(title="input folder")
# output_folder = filedialog.askdirectory(title="output folder")

# 设置图像宽度和高度
width = 1920
height = 1536

# 定义YUV到RGB的转换函数
def yuv_to_rgb(yuv_frame, width, height):
    Y = yuv_frame[:width*height]
    U = yuv_frame[width*height:width*height + (width//2)*(height//2)]
    V = yuv_frame[width*height + (width//2)*(height//2):]

    Y = np.reshape(Y, (height, width))
    U = np.reshape(U, (height//2, width//2)).repeat(2, axis=0).repeat(2, axis=1)
    V = np.reshape(V, (height//2, width//2)).repeat(2, axis=0).repeat(2, axis=1)

    YUV = np.dstack((Y, U, V)).astype(np.uint8)
    # RGB = cv2.cvtColor(YUV, cv2.COLOR_YUV2RGB_I420)
    RGB = cv2.cvtColor(YUV, cv2.COLOR_YUV2RGB)
    return RGB

# 定义YUV到NV12的转换函数
def yuv_to_nv12(yuv_data,width, height):
    Y = yuv_data[:width*height]
    U = yuv_data[width*height:width*height + (width//2)*(height//2)]
    V = yuv_data[width*height + (width//2)*(height//2):]

    # 将Y、U、V分量重新排列以创建NV12数据
    UV = np.empty((height, width), dtype=np.uint8)
    UV[::2, ::2] = U.reshape((height//2, width//2))
    UV[1::2, ::2] = V.reshape((height//2, width//2))

    # 创建NV12数据
    nv12_data = np.dstack((Y.reshape((height, width)), UV))
    return nv12_data

# 定义NV12到BMP的转换函数
def nv12_to_bmp(nv12_data, width, height):
    # Convert NV12 to BGR
    yuv = np.frombuffer(nv12_data, dtype=np.uint8)
    bgr = cv2.cvtColor(yuv.reshape((int(1.5 * height), width)), cv2.COLOR_YUV2BGR_NV12)
    return bgr
    # cv2.imwrite(output_path, bgr)

# 遍历输入文件夹中的所有文件
for root, dirs, files in os.walk(input_folder):
    # print (root, dirs, files)
    # quit()
    for file in files:
        if file.endswith(".yuv"):
            # 构建输入文件的完整路径
            input_file_path = os.path.join(root, file)
            # print(input_file_path)
            # quit()
            # 构建输出文件的完整路径，将文件后缀改为.bmp
            output_file_path = os.path.join(root, os.path.splitext(file)[0] + ".bmp")
            # print(output_file_path)
            # quit()
            # 读取YUV文件
            with open(input_file_path, 'rb') as f:
                yuv_data = np.frombuffer(f.read(), dtype=np.uint8)
                # YUV转NV12
                nv12_frame = yuv_to_nv12(yuv_data, width, height)
                # print("nv12_frame:",nv12_frame)
                
                # NV12转BMP
                bmp_frame = nv12_to_bmp(yuv_data, width, height)
                # print("bmp_frame:",nv12_frame)

                # 将YUV转换为RGB
                # rgb_frame = yuv_to_rgb(yuv_data, width, height)
                # rgb_frame.save(output_file_path, 'bmp')

                # cv2.imwrite(output_file_path,rgb_frame)
                # print(f"Converted {input_file_path} to {output_file_path}")
                cv2.imwrite(output_file_path,bmp_frame)
                print(f"Converted {input_file_path} to {output_file_path}")
                

print("Conversion complete.")