yolo 数据增强(旋转和拉伸),并计算更新后的label信息

yolo 数据增强(旋转和拉伸),并计算更新后的label信息

import cv2
import numpy as np

# 读取图像
img = cv2.imread(r'C:\Users\image\2.png')

# 原来的框
x, y, w, h = 500, 170, 200, 200

# 这里是需要转换的点
point1 = (x, y)
point2 = (x + w, y)
point3 = (x + w, y + w)
point4 = (x, y + w)

# 将四个点组合起来
l = [point1, point4, point2, point3]

angle = 200  # 旋转角度

# 拉伸参数(这里选则不拉伸,拉伸的话点的转化是正确的,但是框架是错误的)
scale_x = 1.33
scale_y = 0.8

# 旋转和拉伸图像
rows, cols = img.shape[:2]
M = cv2.getRotationMatrix2D((cols / 2, rows / 2), angle, 1)

img_rotated = cv2.warpAffine(img, M, (cols, rows))
img_rotated = cv2.resize(img_rotated, (int(img.shape[1] * scale_x), int(img.shape[0] * scale_y)))

nl = []  # 把转化后的点添加这里
for point in l:
    # 计算标记点point的新位置
    point_new = tuple(np.dot(M, np.array([point[0], point[1], 1]))[:2])
    point_new = (point_new[0] * scale_x, point_new[1] * scale_y)
    point_new = tuple(map(int, point_new))
    nl.append(point_new)

ZIP = list(zip(nl[0], nl[1], nl[2], nl[3]))

DRAW = False  # 把结果绘制出来,对比
if DRAW:
    # 绘制转化后的矩形框
    cv2.rectangle(img_rotated, (min(ZIP[0]), min(ZIP[1])), (max(ZIP[0]), max(ZIP[1])), (0, 0, 255), 2)

    # 绘制转化后的四个点
    cv2.circle(img_rotated, nl[0], 5, (255, 0, 0), -1)
    cv2.circle(img_rotated, nl[1], 5, (255, 0, 0), -1)
    cv2.circle(img_rotated, nl[2], 5, (255, 0, 0), -1)
    cv2.circle(img_rotated, nl[3], 5, (255, 0, 0), -1)

    cv2.rectangle(img, (x, y), (x + w, y + h), (0, 0, 255), 2)  # 绘制原来的框

    # 绘制转换前的四个点
    cv2.circle(img, (x, y), 5, (255, 0, 0), -1)
    cv2.circle(img, (x + w, y), 5, (255, 0, 0), -1)
    cv2.circle(img, (x, y + w), 5, (255, 0, 0), -1)
    cv2.circle(img, (x + w, y + w), 5, (255, 0, 0), -1)

    cv2.imshow('Original Image', img)
    cv2.imshow('Transformed Image', img_rotated)
    cv2.waitKey(0)
    cv2.destroyAllWindows()

save_txt = True  # 保存转换得到的图像和label信息 (保存的话需要把DRAW设置为false)
if save_txt:
    cv2.imwrite('trans_img.png', img_rotated)
    temp = [min(ZIP[0]), min(ZIP[1]), max(ZIP[0]) - min(ZIP[0]), max(ZIP[1]) - min(ZIP[1])]
    temp = ' '.join(list((map(str, temp))))
    with open('1.txt','w') as f:
        f.write('0 ')  # 类别信息
        f.write(temp)  # x, y, w, h

前后结果对比

  • 只拉伸不旋转
    在这里插入图片描述

  • 只旋转不拉伸
    在这里插入图片描述

  • 拉伸加旋转
    在这里插入图片描述

  • 大角度旋转加拉伸
    在这里插入图片描述

猜你喜欢

转载自blog.csdn.net/qq_45149610/article/details/130221825