数据集扩充增广/格式转换
适用于yolov3,v5,v7,v8,v9,v10,v11,v12目标检测数据集,支持同步更新标签(txt)
小工具,Windows10以上系统可以直接打开使用,无需安装,无需配置环境
翻转、旋转、缩放、拉伸、剪裁、平移、亮度、对比度、饱和度、噪声等十多种增广方式
可选择增广倍数,生成的图像和标签重命名不会覆盖原数据
新增格式转换(XML→TXT,TXT→XML)和划分数据集功能
数据集增强工具,核心在于图像变换的同时,能够同步、准确地更新 YOLO 格式的.txt标注文件。
功能(随机旋转和翻转),我们需要处理图像像素的变化,同时根据几何变换规则计算边界框坐标的新位置。
以下是构建该系统的核心 Python 代码。你可以将其保存为main.py,并配合PyQt5运行。
核心依赖
你需要安装以下库:
pipinstallopencv-python numpy PyQt5核心代码实现
核心逻辑:随机旋转(含标签同步更新)和水平翻转。
importcv2importnumpyasnpimportosimportrandomclassDataAugmentor:def__init__(self):passdefload_yolo_labels(self,label_path,img_width,img_height):"""读取YOLO格式标签并转换为像素坐标"""boxes=[]ifnotos.path.exists(label_path):returnboxeswithopen(label_path,'r')asf:forlineinf.readlines():parts=line.strip().split()class_id=int(parts[0])# 归一化中心点x, y, 宽w, 高hx_center=float(parts[1])*img_width y_center=float(parts[2])*img_height w=float(parts[3])*img_width h=float(parts[4])*img_height# 转换为左上角和右下角坐标 (x1, y1, x2, y2)x1=x_center-w/2y1=y_center-h/2x2=x_center+w/2y2=y_center+h/2boxes.append([class_id,x1,y1,x2,y2])returnboxesdefsave_yolo_labels(self,boxes,label_path,img_width,img_height):"""将像素坐标转换回YOLO格式并保存"""withopen(label_path,'w')asf:forboxinboxes:class_id=box[0]x1,y1,x2,y2=box[1],box[2],box[3],box[4]# 计算中心点和宽高x_center=(x1+x2)/2y_center=(y1+y2)/2w=x2-x1 h=y2-y1# 归一化x_center/=img_width y_center/=img_height w/=img_width h/=img_height# 防止坐标溢出(0-1之间)x_center=max(0,min(1,x_center))y_center=max(0,min(1,y_center))w=max(0,min(1,w))h=max(0,min(1,h))f.write(f"{class_id}{x_center}{y_center}{w}{h}\n")defaugment_random_rotation(self,img,boxes,angle_range=(-90,90)):""" 随机旋转增强 :param img: OpenCV图像对象 :param boxes: [[class_id, x1, y1, x2, y2], ...] :param angle_range: 旋转角度范围 """h,w=img.shape[:2]angle=random.uniform(angle_range[0],angle_range[1])# 获取旋转矩阵center=(w/2,h/2)M=cv2.getRotationMatrix2D(center,angle,scale=1.0)cos=np.abs(M[0,0])sin=np.abs(M[0,1])# 计算新的图像尺寸以避免裁剪new_w=int((h*sin)+(w*cos))new_h=int((h*cos)+(w*sin))# 调整旋转矩阵的平移部分M[0,2]+=(new_w/2)-center[0]M[1,2]+=(new_h/2)-center[1]# 旋转图像rotated_img=cv2.warpAffine(img,M,(new_w,new_h))# 旋转边界框new_boxes=[]forboxinboxes:cls_id=box[0]# 获取矩形的四个顶点pts=np.array([[box[1],box[2]],# x1, y1[box[3],box[2]],# x2, y1[box[3],box[4]],# x2, y2[box[1],box[4]]# x1, y2],dtype=np.float32)# 对顶点应用旋转ones=np.ones(shape=(len(pts),1))pts_ones=np.hstack([pts,ones])transformed_pts=M.dot(pts_ones.T).T# 获取新的外接矩形x_coords=transformed_pts[:,0]y_coords=transformed_pts[:,1]x1,y1=np.min(x_coords),np.min(y_coords)x2,y2=np.max(x_coords),np.max(y_coords)# 简单过滤:如果框完全在图像外,则丢弃(这里简化处理,未做复杂裁剪逻辑)ifx2>0andy2>0andx1<new_wandy1<new_h:# 限制在图像范围内x1,y1=max(0,x1),max(0,y1)x2,y2=min(new_w,x2),min(new_h,y2)new_boxes.append([cls_id,x1,y1,x2,y2])returnrotated_img,new_boxesdefaugment_horizontal_flip(self,img,boxes):"""水平翻转"""h,w=img.shape[:2]flipped_img=cv2.flip(img,1)# 1表示水平翻转new_boxes=[]forboxinboxes:cls_id,x1,y1,x2,y2=box# 翻转逻辑:新的x = 图像宽 - 旧的xnew_x1=w-x2 new_x2=w-x1 new_boxes.append([cls_id,new_x1,y1,new_x2,y2])returnflipped_img,new_boxes如何集成到 UI
截图中的界面使用了 PyQt5。你可以创建一个按钮槽函数来调用上面的类:
# 伪代码示例:如何在 PyQt5 中调用defon_start_augment(self):aug=DataAugmentor()# 1. 获取路径img_path=self.ui.image_path.text()txt_path=self.ui.txt_path.text()# 2. 读取img=cv2.imread(img_path)h,w=img.shape[:2]boxes=aug.load_yolo_labels(txt_path,w,h)# 3. 执行增强 (例如随机旋转)# 注意:这里需要把界面上输入的角度范围传进来,例如 (-90, 90)new_img,new_boxes=aug.augment_random_rotation(img,boxes,angle_range=(-90,90))# 4. 保存save_img_path="output/aug_001.jpg"save_txt_path="output/aug_001.txt"cv2.imwrite(save_img_path,new_img)aug.save_yolo_labels(new_boxes,save_txt_path,new_img.shape[1],new_img.shape[0])print("增强完成!")关键点解析
- 坐标变换逻辑:YOLO 使用归一化坐标
(x_center, y_center, w, h)。在进行几何变换(如旋转)时,必须先将其转换为像素坐标(x1, y1, x2, y2),变换后再转回归一化坐标保存。 - 旋转后的尺寸变化:图片旋转后,为了防止图像四个角被裁剪,通常画布尺寸会变大(代码中通过计算
new_w和new_h实现)。 - 多格式支持:如果要实现 XML 转 TXT,可以使用 Python 内置的
xml.etree.ElementTree解析 Pascal VOC 格式的 XML 文件,提取<bndbox>数据,然后套用上面的save_yolo_labels逻辑即可。
以上代码逻辑完整覆盖了图像增强中坐标同步更新的核心需求,可直接用于构建数据集处理工具。