前言

  对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