工作杂记-YUV的dump和read

工作杂记-YUV的dump和read

工作中涉及到模型验证相关的工作，这里是三个模型的共同作用，在感知模型读取图片的时候，把输入替换成自己给定的输入，保证一致性，再来看输出。
知道了模型输入的宽高和步长之后，有如下的dump_yuv和read_yuv函数。其中bpu_addr和bpu_addr_uv是y分量和uv分量的地址。

bool ReadYUVImage(const std::string &img_path, int width, int height,
                  uint64_t bpu_addr, uint64_t bpu_addr_uv) {
  std::ifstream input_img_file(img_path, std::ios::binary);
  input_img_file.seekg(0, std::ios::end);
  int len = input_img_file.tellg();
  input_img_file.seekg(0, std::ios::beg);
  std::vector<uint8_t> img_data(width * height * 3 / 2);
  if (len < width * height * 3 / 2) {
    HSLOG_E << "file length is not right " << len << " img_file:" << img_path;
    return false;
  } else {
    input_img_file.read(reinterpret_cast<char *>(img_data.data()),
                        width * height * 3 / 2);
    memcpy(reinterpret_cast<void *>(bpu_addr), img_data.data(), width * height);
    memcpy(reinterpret_cast<void *>(bpu_addr_uv),
           img_data.data() + width * height, width * height / 2);
  }
  input_img_file.close();
  return true;
}

void DumpYUVImage(const std::string &out_file, int width, int height,
                  uint64_t bpu_addr, uint64_t bpu_addr_uv) {
  std::ofstream fout(out_file, std::ios::out | std::ios::binary);
  std::vector<uint8_t> src_input_image(width * height * 3 / 2);
  memcpy(src_input_image.data(), reinterpret_cast<void *>(bpu_addr),
         width * height);
  memcpy(src_input_image.data() + width * height,
         reinterpret_cast<void *>(bpu_addr_uv), width * height / 2);
  fout.write(reinterpret_cast<char *>(src_input_image.data()),
             width * height * 3 / 2);
  fout.close();
}

利用dump生成图片 yuv2img

这里需要了解yuv相关的排布知识。
YUV格式有两大类：planar和packed。
对于planar的YUV格式，先连续存储所有像素点的Y，紧接着存储所有像素点的U，随后是所有像素点的V。
对于packed的YUV格式，每个像素点的Y,U,V是连续交*存储的。

YUV，分为三个分量，“Y”表示明亮度（Luminance或Luma），也就是灰度值；而“U”和“V” 表示的则是色度（Chrominance或Chroma），作用是描述影像色彩及饱和度，用于指定像素的颜色。

planar的YUV格式分为YUV420P和YUV420SP，YUV420P包含I420和YV12。I420格式和YV12格式的不同处在U平面和V平面的位置不同。在I420格式中，U平面紧跟在Y平面之后，然后才是V平面（即：YUV）；但YV12则是相反（即：YVU）。
YUV420SP, Y分量平面格式，UV打包格式, 即NV12。 NV12与NV21类似，U 和 V 交错排列,不同在于UV顺序。
I420: YYYYYYYY UU VV =>YUV420P
YV12: YYYYYYYY VV UU =>YUV420P
NV12: YYYYYYYY UVUV =>YUV420SP
NV21: YYYYYYYY VUVU =>YUV420SP

注意转化时要知道文件的uv排列先后顺序，实测u，v颠倒后的图片色差会有变化。
正确的结果

错误的结果

核心代码，这里有一个opencv的坑，在merge时 会报错Invalid number of channels in input image: ‘VScn::contains(scn)’ ‘scn’，原因时Y U V的shape不一致，解决方案时把UV插值成和Y一样的，在merge，这样shape就成了(h，w，3)这种形式。再调用cv2.COLOR_YUV2BGR去做转化。

    # 分离UV通道为U和V
    U = UV[:, ::2]
    V = UV[:, 1::2]

    # 扩展U和V通道为与Y通道一致的大小
    U = cv2.resize(U, (width, height), interpolation=cv2.INTER_LINEAR)
    V = cv2.resize(V, (width, height), interpolation=cv2.INTER_LINEAR)

    # 合并Y、U、V通道数据
    YUV420p_image = cv2.merge((Y, U, V))
    # 转换为RGB格式
    image_rgb = cv2.cvtColor(YUV420p_image, cv2.COLOR_YUV2BGR)

yuv2img代码

如下是根据dump的yuv文件生成图片，其中三个函数，分别为灰度(only y 分量)，420sp直接转RGB，420sp 提取 U、V分量后，再插值和Y分量组成YUV444，再转成RGB。

Y = np.frombuffer(buffer[:Y_size], dtype=np.uint8).reshape(height, width)
UV = np.frombuffer(buffer[Y_size:], dtype=np.uint8).reshape(height // 2, width)

# 分离UV通道为U和V
U = UV[:, ::2] # 实际上是选择 UV 数组的所有行，但只选择每一行中的偶数列元素，即第0列、第2列、第4列等
V = UV[:, 1::2] 

# 扩展U和V通道为与Y通道一致的大小
U = cv2.resize(U, (width, height), interpolation=cv2.INTER_LINEAR)
V = cv2.resize(V, (width, height), interpolation=cv2.INTER_LINEAR)

# 合并Y、U、V通道数据
YUV420p_image = cv2.merge((Y, U, V))

import cv2
import numpy as np
import re
import os

def convert_YUV420sp_RGB(yuv_file):
        # 从文件名中提取宽度和高度
    match = re.search(r'_w_(\d+)_h_(\d+)', yuv_file)
    if match:
        width = int(match.group(1))
        height = int(match.group(2))
    else:
        print("无法提取分辨率信息:", yuv_file)
        return
    print(width, height)
    with open(yuv_file, "rb") as f:
        buffer = f.read()
    image = np.frombuffer(buffer, np.uint8).reshape(height*3//2, width)
    print(image.size, image.shape)
    image_rgb = cv2.cvtColor(image, cv2.COLOR_YUV420SP2RGB)
    
    yuv_file = os.path.basename(yuv_file)
    output_file = yuv_file.replace(".yuv", ".jpg")
    output_file = os.path.join(output_folder_path, f"{output_file}.jpg")
    cv2.imwrite(output_file, image_rgb)
    # print(output_file)

def convert_YUV420sp_YUV444_RGB(yuv_file):
    # 从文件名中提取宽度和高度
    match = re.search(r'_w_(\d+)_h_(\d+)', yuv_file)
    if match:
        width = int(match.group(1))
        height = int(match.group(2))
    else:
        print("无法提取分辨率信息:", yuv_file)
        return

    # 读取YUV420sp图像数据
    with open(yuv_file, "rb") as f:
        buffer = f.read()

    # 解析Y、UV通道数据
    Y_size = width * height
    UV_size = width * height // 2

    Y = np.frombuffer(buffer[:Y_size], dtype=np.uint8).reshape(height, width)
    UV = np.frombuffer(buffer[Y_size:], dtype=np.uint8).reshape(height // 2, width)

    # 分离UV通道为U和V
    U = UV[:, ::2]
    V = UV[:, 1::2]

    # 扩展U和V通道为与Y通道一致的大小
    U = cv2.resize(U, (width, height), interpolation=cv2.INTER_LINEAR)
    V = cv2.resize(V, (width, height), interpolation=cv2.INTER_LINEAR)

    # 合并Y、U、V通道数据
    YUV420p_image = cv2.merge((Y, U, V))
    # 转换为RGB格式
    image_rgb = cv2.cvtColor(YUV420p_image, cv2.COLOR_YUV2BGR)

    yuv_file = os.path.basename(yuv_file)
    output_file = yuv_file.replace(".yuv", ".jpg")
    output_file = os.path.join(output_folder_path, f"{output_file}.jpg")
    cv2.imwrite(output_file, image_rgb)
    # print(output_file)

def convert_YUV420sp_GRAY(yuv_file):
    file_name = os.path.splitext(yuv_file)[0]
    match = re.search(r'_w_(\d+)_h_(\d+)', yuv_file)
    if match:
        width = int(match.group(1))
        height = int(match.group(2))
    else:
        print("无法提取分辨率信息:", yuv_file)
        return

    # 打开YUV文件并读取数据
    
    with open(yuv_file, 'rb') as file:
        buffer = file.read()

    Y_size = width * height
    Y_data = buffer[:Y_size]

    # 将Y数据转换为NumPy数组
    Y = np.frombuffer(Y_data, dtype=np.uint8).reshape((height, width))
    gray_image = cv2.cvtColor(Y, cv2.COLOR_GRAY2BGR)

    yuv_file = os.path.basename(yuv_file)
    output_file = yuv_file.replace(".yuv", ".jpg")
    output_file = os.path.join(output_folder_path, f"{output_file}.jpg")
    cv2.imwrite(output_file, gray_image)
    return

# 定义输入YUV文件夹路径
input_folder_path = "./dump_yuv"  # 替换为包含YUV文件的文件夹路径

# 定义输出JPEG文件夹路径
output_folder_path = input_folder_path + "_output"  # 保存JPEG文件的文件夹路径
if not os.path.exists(output_folder_path):
    os.makedirs(output_folder_path)

# 获取输入文件夹中的所有YUV文件
yuv_files = [file for file in os.listdir(input_folder_path) if file.endswith(".yuv")]

for yuv_file in yuv_files:
    print(yuv_file)
    # convert_YUV420sp_YUV444_RGB(input_folder_path + "/" + yuv_file)
    # convert_YUV420sp_RGB(input_folder_path + "/" + yuv_file)
    convert_YUV420sp_GRAY(input_folder_path + "/" + yuv_file)

model_id:0 percepts_[0] rect num:1
model_id:0 rect:388.95 1034.04 520.791 1163.19 7.03816 0
model_id:0 percepts_[5] rect num:8
model_id:0 rect:2164.7 1034.81 2261.26 1110.25 8.7623 5
model_id:0 rect:2464.56 996.751 2519.25 1059.25 7.19807 5
model_id:0 rect:2654.99 593.22 2744.83 696.733 8.5591 5
model_id:0 rect:2914.5 570.069 3007.19 666.134 8.76954 5
model_id:0 rect:3044.18 553.215 3152.04 664.448 8.68193 5
model_id:0 rect:3185.64 543.125 3295.19 652.673 7.15572 5
model_id:0 rect:3750.75 609.548 3836 781.291 8.73659 5
model_id:0 rect:3761.16 792.687 3836 1073.72 9.44854 5
model_id:0 percepts_[11] rect num:3
model_id:0 rect:247.464 1250.77 334.153 1311.08 7.29449 11
model_id:0 rect:1664.7 1409.73 2243.79 1526.45 7.30121 11
model_id:0 rect:3057.62 1149.15 3528.76 1186.85 7.29843 11
model_id:0 percepts_[13] rect num:3
model_id:0 rect:1390.61 1212.43 1441.98 1423.03 8.79426 13
model_id:0 rect:1992.3 1189.52 2025.68 1335.91 7.29169 13
model_id:0 rect:2816.59 1182.99 2865.39 1378.19 7.29921 13

再额外附一段画框的代码

import cv2
import re

# 读取图片
image_path = './dump_yuv_v3_output/model_50_frame_1694608049106_w_3840_h_2160.jpg.jpg'
image = cv2.imread(image_path)

# 打开包含矩形框坐标的文件
with open('rect.txt', 'r') as file:
    lines = file.readlines()

# 正则表达式模式，用于匹配包含model_id和percepts的行
model_pattern = re.compile(r'model_id:\d+ percepts_\[(\d+)\] rect num:(\d+)')

# 正则表达式模式，用于匹配包含model_id和rect的行
rect_pattern = re.compile(r'model_id:\d+ rect:(\d+(?:\.\d+)?) (\d+(?:\.\d+)?) (\d+(?:\.\d+)?) (\d+(?:\.\d+)?)')

# 初始化一个字典来存储percepts_值和颜色的映射
percepts_color_mapping = {
    "0": (0, 0, 255),   # 红色
    "5": (0, 255, 0),   # 绿色
    "11": (255, 0, 0),   # 蓝色
    "13": (0, 255, 255), # 黄色
}


# 遍历文件中的每一行
for line in lines:
    # 尝试匹配包含矩形坐标的行
    model_match = model_pattern.match(line)
    if model_match:
        percept_index = model_match.group(1)
        num_rects = model_match.group(2)
        print(f"Percept Index: {percept_index}")
        print(f"Number of Rectangles: {num_rects}")

    # 匹配对应的框
    rect_match = rect_pattern.search(line)
    if rect_match:
        x1, y1, x2, y2 = map(float, rect_match.groups())
        
        # 获取颜色
        color = percepts_color_mapping.get(percept_index)  # 默认为红色
        
        # 绘制矩形框
        x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
        print(x1, y1, x2, y2)
        cv2.rectangle(image, (x1, y1), (x2, y2), color, 2)

        font = cv2.FONT_HERSHEY_SIMPLEX
        cv2.putText(image, str(percept_index), (x1, y1 - 10), font, 0.5, color, 1, cv2.LINE_AA)

# 保存包含绘制框的图像
output_image_path = 'output_image.png'
cv2.imwrite(output_image_path, image)

# 保存包含绘制框的图像
output_image_path = 'output_image.png'
cv2.imwrite(output_image_path, image)