前言
在用模型推理后，经常需要看推理效果，一个一个图核对，非常麻烦，编写下面的脚本用于显示标注和推理结果在一张图，看效果。
1.程序解释
draw_bounding_boxes(images_path, labels_path, pre_labels_path, output_path, label_index=None):
- images_path: 包含图像文件的文件夹路径。
- labels_path: 包含真实标签文件的文件夹路径。
- pre_labels_path: 包含预测标签文件的文件夹路径。
- output_path: 输出绘制了边界框后图像的保存路径。
- label_index=None: 要绘制的标签的类别索引，如果为None，则绘制所有类别的标签。
check_label_counts(images_path, labels_path, pre_labels_path, output_path, label_index=None):
- images_path: 包含图像文件的文件夹路径。
- labels_path: 包含真实标签文件的文件夹路径。
- pre_labels_path: 包含预测标签文件的文件夹路径。
- output_path: 输出错误图像的保存路径。
l- abel_index=None: 要检查的标签的类别索引，如果为None，则检查所有类别的标签。
这些函数都是用来处理目标检测模型的输出结果的。第一个函数draw_bounding_boxes用于绘制边界框，第二个函数check_label_counts用于检查标签数量是否一致。这些函数的参数主要是文件夹路径和标签索引，用于指定要处理的图像和标签。
2.完整代码
代码如下（示例）：
import cv2
import os
from tqdm import tqdm
def draw_bounding_boxes(images_path, labels_path, pre_labels_path, output_path, label_index=None):
# 获取文件夹中所有图片文件的列表
image_files = os.listdir(images_path)
label_files = os.listdir(labels_path)
pre_label_files = os.listdir(pre_labels_path)
print(f'图片数量： {len(image_files)}')
print(f'标签文件数量： {len(label_files)}')
print(f'预测文件数量： {len(pre_label_files)}')
# 遍历images文件夹下的图片
for image_file in tqdm(image_files, desc='draw_bounding_boxes'):
# 获取文件名（不带扩展名）
filename = os.path.splitext(image_file)[0]
# 检查是否存在相同文件名的标签文件和预测文件
if f'{filename}.txt' in label_files and f'{filename}.txt' in pre_label_files:
# 读取图像
img = cv2.imread(os.path.join(images_path, image_file))
# 获取图像的宽度和高度
img_h, img_w, _ = img.shape
# 读取YOLO格式的标签文件（第一个标签文件）
with open(os.path.join(labels_path, f'{filename}.txt'), 'r') as label_file:
lines = label_file.readlines()
# 遍历第一个标签文件中的每个检测框信息
for line in lines:
class_id, x_center, y_center, width, height = map(float, line.split())
if label_index is None or int(class_id) == label_index:
# 将YOLO格式的标注信息转换为实际图像上的坐标
x_center *= img_w
y_center *= img_h
width *= img_w
height *= img_h
# 计算边框的左上角和右下角坐标
x1 = int(x_center - width / 2)
y1 = int(y_center - height / 2)
x2 = int(x_center + width / 2)
y2 = int(y_center + height / 2)
color = (0, 255, 0) # 设置边界框颜色为绿色
alpha = 0.5 # 设置透明度为 0.7（更高的透明度）
overlay = img.copy()
cv2.rectangle(overlay, (x1, y1), (x2, y2), color, -1) # 画矩形填充到overlay上
cv2.addWeighted(overlay, alpha, img, 1 - alpha, 0, img) # 将overlay融合到img中
# 读取另一个YOLO格式的标签文件（第二个标签文件）
with open(os.path.join(pre_labels_path, f'{filename}.txt'), 'r') as pre_label_file:
lines = pre_label_file.readlines()
# 遍历第二个标签文件中的每个检测框信息
for line in lines:
class_id, x_center, y_center, width, height = map(float, line.split())
if label_index is None or int(class_id) == label_index:
# 将YOLO格式的标注信息转换为实际图像上的坐标
x_center *= img_w
y_center *= img_h
width *= img_w
height *= img_h
# 计算边框的左上角和右下角坐标
x1 = int(x_center - width / 2)
y1 = int(y_center - height / 2)
x2 = int(x_center + width / 2)
y2 = int(y_center + height / 2)
color = (255, 0, 0) # 设置边界框颜色为蓝色
thickness = 2 # 设置边框线条粗细
cv2.rectangle(img, (x1, y1), (x2, y2), color, thickness) # 画矩形框线
# 保存绘制了边界框的图像到error文件夹下
output_file = os.path.join(output_path, f'{filename}_bbox.bmp')
cv2.imwrite(output_file, img)
def check_label_counts(images_path, labels_path, pre_labels_path, output_path, label_index=None):
# 创建保存错误图像的文件夹
error_img_folder = os.path.join(output_path, 'error_img')
os.makedirs(error_img_folder, exist_ok=True)
# 获取文件夹中所有图片文件的列表
image_files = os.listdir(images_path)
# 遍历images文件夹下的图片
for image_file in tqdm(image_files, desc='check_label_counts'):
# 获取文件名（不带扩展名）
filename = os.path.splitext(image_file)[0]
# 读取标签文件内容
with open(os.path.join(labels_path, f'{filename}.txt'), 'r') as label_file:
labels_data = label_file.readlines()
with open(os.path.join(pre_labels_path, f'{filename}.txt'), 'r') as pre_label_file:
pre_labels_data = pre_label_file.readlines()
# 检查标签数量是否相同
if len(labels_data) != len(pre_labels_data):
# 读取图像
img = cv2.imread(os.path.join(images_path, image_file))
img_h, img_w, _ = img.shape
img_labels = img.copy()
img_pre_labels = img.copy()
img_copy_labels = img.copy()
# 在图像上显示不匹配的标签
color1 = (255, 0, 0) # 设置红色框
color2 = (0, 0, 255) # 设置蓝色框
thickness = 2 # 设置边框线条粗细
for label in labels_data:
class_id, x_center, y_center, width, height = map(float, label.split())
if label_index is None or int(class_id) == label_index:
# 将YOLO格式的标注信息转换为实际图像上的坐标
x_center *= img_w
y_center *= img_h
width *= img_w
height *= img_h
# 计算边框的左上角和右下角坐标
x1 = int(x_center - width / 2)
y1 = int(y_center - height / 2)
x2 = int(x_center + width / 2)
y2 = int(y_center + height / 2)
cv2.rectangle(img_labels, (x1, y1), (x2, y2), color1, thickness) # 绘制红色框
color = (0, 255, 0) # 设置边界框颜色为绿色
alpha = 0.5 # 设置透明度为 0.7（更高的透明度）
overlay = img_copy_labels.copy()
cv2.rectangle(overlay, (x1, y1), (x2, y2), color, -1) # 画矩形填充到overlay上
cv2.addWeighted(overlay, alpha, img_copy_labels, 1 - alpha, 0, img_copy_labels) # 将overlay融合到img中
label_txt_path = os.path.join(error_img_folder, filename, f'{filename}_label.txt')
if not os.path.exists(label_txt_path):
os.makedirs(os.path.dirname(label_txt_path), exist_ok=True)
with open(label_txt_path, 'w') as label_txt:
label_txt.writelines(labels_data)
for label in pre_labels_data:
class_id, x_center, y_center, width, height = map(float, label.split())
if label_index is None or int(class_id) == label_index:
# 将YOLO格式的标注信息转换为实际图像上的坐标
x_center *= img_w
y_center *= img_h
width *= img_w
height *= img_h
# 计算边框的左上角和右下角坐标
x1 = int(x_center - width / 2)
y1 = int(y_center - height / 2)
x2 = int(x_center + width / 2)
y2 = int(y_center + height / 2)
cv2.rectangle(img_pre_labels, (x1, y1), (x2, y2), color2, thickness) # 绘制蓝色框
cv2.rectangle(img_copy_labels, (x1, y1), (x2, y2), color1, thickness)
with open(os.path.join(error_img_folder, filename, f'{filename}_pre_labels.txt'), 'w') as pre_labels_txt:
pre_labels_txt.writelines(pre_labels_data)
# 保存图像和标签图到以图片名为名的子文件夹中
image_folder_path = os.path.join(error_img_folder, filename)
os.makedirs(image_folder_path, exist_ok=True)
cv2.imwrite(os.path.join(image_folder_path, f'{filename}.bmp'), img) # 保存原始图像
cv2.imwrite(os.path.join(image_folder_path, f'{filename}_labels.bmp'), img_labels) # 保存labels标签图
cv2.imwrite(os.path.join(image_folder_path, f'{filename}_pre_labels.bmp'), img_pre_labels) # 保存pre_labels标签图
cv2.imwrite(os.path.join(image_folder_path, f'{filename}_match_labels.bmp'), img_copy_labels) # 保存pre_labels标签图
# 输入参数
images_path = '/codeyard/yolov5_6.0_tensorrt_TL/datasetall/images'
labels_path = '/codeyard/yolov5_6.0_tensorrt_TL/datasetall/labels'
pre_labels_path = '/codeyard/yolov5_6.0_tensorrt_TL/runs/detect/exp35/labels'
output_path = '/codeyard/yolov5_6.0_tensorrt_TL/save_error'
draw_label_index = None # 若要显示所有标签，则将此处改为 None
check_label_index = None
# 调用函数，传入标签序号进行测试
draw_bounding_boxes(images_path, labels_path, pre_labels_path, output_path, draw_label_index)
check_label_counts(images_path, labels_path, pre_labels_path, output_path,check_label_index)
2.读入数据
代码如下（示例）：
data = pd.read_csv(
'https://labfile.oss.aliyuncs.com/courses/1283/adult.data.csv')
print(data.head())
该处使用的url网络请求的数据。
总结
提示：这里对文章进行总结：
例如：以上就是今天要讲的内容，本文仅仅简单介绍了pandas的使用，而pandas提供了大量能使我们快速便捷地处理数据的函数和方法。
扫一扫
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
