工作小记-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)
1万+
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
