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
前后结果对比
-
只拉伸不旋转
-
只旋转不拉伸
-
拉伸加旋转
-
大角度旋转加拉伸