工作小记-YUV的dump和read
- 工作杂记-YUV的dump和read
- 利用dump生成图片 yuv2img
- yuv2img代码
工作杂记-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和VU = 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 osdef 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)returnprint(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 * heightUV_size = width * height // 2Y = 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和VU = 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 * heightY_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_SIMPLEXcv2.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)