yolo旋转图片和标签坐标,数据增强
前言
对yolo做数据增强,旋转图片后,标签的坐标也要跟着旋转,网上看了一大堆不是旋转的图片有问题就是旋转后的标签有问题,自己写了一个,欢迎指点评论。
直接上代码,有注释,看不懂的欢迎评论区和私信
import cv2
import numpy as np
import os
from math import cos,sin
jiaodu=30 #旋转的角度
zq_id="zq1" #旋转后的名字(在原有名字的后面加上zq1)
img_folder="images" #存储需要处理的图片的文件夹路径
img_save_folder="images" #旋转后需要保存图像的文件夹路径
txt_folder="labelTxt" #对应图片标签的文件夹路径(标签文件名需要和对应的图片文件名一样)
txt_save_folder="labelTxt" #旋转后需要保存标签的文件夹路径
Add_label_information=False #是否在旋转后的标签文件中加入标签信息
#坐标旋转计算函数
def rotate_xy(x, y, angle, cx, cy): #参数需要旋转点的x、y坐标,旋转角度、旋转中心点的x、y坐标
"""
点(x,y) 绕(cx,cy)点旋转
"""
# print(cx,cy)
angle = angle * np.pi / 180 #旋转角度需要转换一下格式
x_new = (x - cx) * cos(angle) - (y - cy) * sin(angle) + cx #计算旋转后的x坐标
y_new = (x - cx) * sin(angle) + (y - cy) * cos(angle) + cy #计算旋转后的y坐标
return [x_new, y_new] #返回旋转后的[x,y]坐标
#图像和坐标旋转增强
for img_path in os.listdir(img_folder):
#图像旋转
img=cv2.imread(img_folder+"/"+img_path) #读图
height,width=img.shape[:2] #获取宽和高
cx,cy=height/2,width/2 #计算旋转中心点的坐标
M=cv2.getRotationMatrix2D((cx,cy),jiaodu,1) #根据旋转角度计算映射矩阵
dst=cv2.warpAffine(img,M,(height,width),borderValue=(255,255,255)) #根据映射矩阵旋转图像
cv2.imwrite(img_save_folder+"/"+os.path.splitext(img_path)[0]+zq_id+'.jpg',dst) #保存图像
#标签旋转
txt=open(txt_folder+"/"+os.path.splitext(img_path)[0]+".txt").read() #读标签
points_list=[] #用于保存原来的坐标点[[x1,y1],[x2,y2],[x3,y3],[x4,y4]]
for i in range(8)[::2]:
points_list.append([txt.split(' ')[i],txt.split(' ')[i+1]]) #将原始坐标组合为【x1,y1】的形式存储到列表中
new_newpoints=[] #格式和points_list一样,用于保存旋转后的坐标
for i in points_list:
new_newpoints.append(rotate_xy(float(i[0]),float(i[1]),-jiaodu,cx,cy))#将每一个点进行旋转
#加入标签信息
if Add_label_information: #如果需要加入标签信息请修改标签信息那一行
sava_txt=''
for i in new_newpoints:
for j in i:
sava_txt+=str(j)+' '
sava_txt+=txt.split(' ')[-2]+' '+txt.split(' ')[-1] #最后加入标签信息,这里我最后有两个标签信息,
#一般情况只有一个标签信息,删掉前面的txt.split(' ')[-2]+' '+就行了
dir_name=txt_save_folder+"/"+os.path.splitext(img_path)[0]+zq_id+".txt"
else: #如果不加入标签信息只保存旋转后的坐标那么最开头的Add_label_information设置为False
sava_txt=''
for i in new_newpoints:
for j in i:
sava_txt+=str(j)+' ' #将坐标由列表转化为yolo的格式,中间分割由','变为' '
dir_name=txt_save_folder+"/"+os.path.splitext(img_path)[0]+zq_id+".txt"
fb = open(dir_name ,mode='w',encoding='utf-8') #创建一个dir_name.txt的文本,
fb.write(str(sava_txt)) # 写入变量sava_txt
放一下我的目录结构
images和labelTxt分别存放需要增强的图片和标签
这里我增强后的路径也写的这两个,增强后的图片和标签也保存在这两个文件夹里面
我的label文件是这样的,前面八个是四个点的坐标都一样
最后面我是两个标签,如果只有一个标签在代码里面改一下加入标签信息那里,注释有写
总结
人懒此处省略,有问题欢迎评论区,阿巴阿巴。。。。。。。。。。。。。