在训练一个自己的yolo模型目标检测模型时,使用公共数据集时,通常要将图片缩放处理,而此时图片对应的标签文件中目标的坐标也应进行同等的变换,这样才能保证模型的正确训练。(当然,如果是自建的数据集,则将图片进行缩放后,使用Labelimg标注工具进行手动标注即可)
已知文件夹a中存放着批量图片,文件夹b中存放着文件夹a中每张图片对应的标签文件.txt,现要求将文件夹a中的图片缩放为256像素*256像素,并且对文件夹b中的标签文件中的目标的坐标进行同等变换,然后将缩放后的图片另存到文件夹c中,同等变换后的标签文件另存到文件夹d中。
import os
import cv2
# 输入文件夹路径
input_img_folder = r'E:\系统默认\桌面\a'
input_label_folder = r'E:\系统默认\桌面\b'
# 输出文件夹路径
output_img_folder = r'E:\系统默认\桌面\c'
output_label_folder = r'E:\系统默认\桌面\d'
# 目标尺寸
target_size = (256, 256)
# 确保输出文件夹存在
if not os.path.exists(output_img_folder):
os.makedirs(output_img_folder)
if not os.path.exists(output_label_folder):
os.makedirs(output_label_folder)
# 遍历文件夹a中的每张图片
for file_name in os.listdir(input_img_folder):
if file_name.endswith('.png'):
# 读取图片
img_path = os.path.join(input_img_folder, file_name)
img = cv2.imread(img_path)
# 缩放图片
resized_img = cv2.resize(img, target_size)
# 保存缩放后的图片到文件夹c
output_img_path = os.path.join(output_img_folder, file_name)
cv2.imwrite(output_img_path, resized_img)
# 处理对应的标签文件
label_file_name = file_name.replace('.png', '.txt')
label_file_path = os.path.join(input_label_folder, label_file_name)
output_label_file_path = os.path.join(output_label_folder, label_file_name)
# 读取原始标签文件内容
with open(label_file_path, 'r') as f:
lines = f.readlines()
# 对每个目标的坐标进行同等变换
transformed_lines = []
for line in lines:
# 解析坐标
class_id, x, y, width, height = map(float, line.strip().split())
# 计算变换后的坐标
x *= target_size[0] / img.shape[1]
y *= target_size[1] / img.shape[0]
width *= target_size[0] / img.shape[1]
height *= target_size[1] / img.shape[0]
# 将变换后的坐标保存为字符串
transformed_line = f'{class_id} {x} {y} {width} {height}\n'
transformed_lines.append(transformed_line)
# 保存变换后的标签文件到文件夹d
with open(output_label_file_path, 'w') as f:
f.writelines(transformed_lines)
上述代码同时实现了将图片进行缩放和其对应标签文件的同等变换。?
这段代码假设标签文件格式为一行一个目标,每行包括类别ID、目标在图像中的归一化Bounding Box(以图片宽度和高度为单位),用空格分隔。如果标签文件格式不同,需要进行相应的修改。(这也是yolo官方模型的标签格式,建议首先确保标签格式与此保持一致)