利用opencv对图像和检测框做任意角度的旋转

一.钢筋比赛中的数据扩充

#coding:utf-8
#数据集扩增
import cv2
import math
import numpy as np
import xml.etree.ElementTree as ET
import osdef rotate_image(src, angle, scale=1):w = src.shape[1]h = src.shape[0]# 角度变弧度rangle = np.deg2rad(angle)  # angle in radians# now calculate new image width and heightnw = (abs(np.sin(rangle) * h) + abs(np.cos(rangle) * w)) * scalenh = (abs(np.cos(rangle) * h) + abs(np.sin(rangle) * w)) * scale# ask OpenCV for the rotation matrixrot_mat = cv2.getRotationMatrix2D((nw * 0.5, nh * 0.5), angle, scale)# calculate the move from the old center to the new center combined# with the rotationrot_move = np.dot(rot_mat, np.array([(nw - w) * 0.5, (nh - h) * 0.5, 0]))# the move only affects the translation, so update the translation# part of the transformrot_mat[0, 2] += rot_move[0]rot_mat[1, 2] += rot_move[1]dst = cv2.warpAffine(src, rot_mat, (int(math.ceil(nw)), int(math.ceil(nh))), flags=cv2.INTER_LANCZOS4)# 仿射变换return dst# 对应修改xml文件
def rotate_xml(src, xmin, ymin, xmax, ymax, angle, scale=1.):w = src.shape[1]h = src.shape[0]rangle = np.deg2rad(angle)  # angle in radians# now calculate new image width and height# 获取旋转后图像的长和宽nw = (abs(np.sin(rangle)*h) + abs(np.cos(rangle)*w))*scalenh = (abs(np.cos(rangle)*h) + abs(np.sin(rangle)*w))*scale# ask OpenCV for the rotation matrixrot_mat = cv2.getRotationMatrix2D((nw*0.5, nh*0.5), angle, scale)# calculate the move from the old center to the new center combined# with the rotationrot_move = np.dot(rot_mat, np.array([(nw-w)*0.5, (nh-h)*0.5,0]))# the move only affects the translation, so update the translation# part of the transformrot_mat[0, 2] += rot_move[0]rot_mat[1, 2] += rot_move[1]# print('rot_mat=', rot_mat)# rot_mat是最终的旋转矩阵# point1 = np.dot(rot_mat, np.array([xmin, ymin, 1]))          #这种新画出的框大一圈# point2 = np.dot(rot_mat, np.array([xmax, ymin, 1]))# point3 = np.dot(rot_mat, np.array([xmax, ymax, 1]))# point4 = np.dot(rot_mat, np.array([xmin, ymax, 1]))point1 = np.dot(rot_mat, np.array([(xmin+xmax)/2, ymin, 1]))   # 获取原始矩形的四个中点，然后将这四个点转换到旋转后的坐标系下# print('point1=',point1)point2 = np.dot(rot_mat, np.array([xmax, (ymin+ymax)/2, 1]))# print('point2=', point2)point3 = np.dot(rot_mat, np.array([(xmin+xmax)/2, ymax, 1]))# print('point3=', point3)point4 = np.dot(rot_mat, np.array([xmin, (ymin+ymax)/2, 1]))# print('point4=', point4)concat = np.vstack((point1, point2, point3, point4))            # 合并np.array# print('concat=', concat)# 改变array类型concat = concat.astype(np.int32)rx, ry, rw, rh = cv2.boundingRect(concat)                        #rx,ry,为新的外接框左上角坐标，rw为框宽度，rh为高度，新的xmax=rx+rw,新的ymax=ry+rhreturn rx, ry, rw, rh'''使图像旋转15, 30, 45, 60, 75, 90, 105, 120度
'''
# imgpath = './images/train/'          #源图像路径
imgpath = './train_example/'          #源图像路径
xmlpath = './Annotations/'            #源图像所对应的xml文件路径
rotated_imgpath = './train_example_out/'
rotated_xmlpath = './Annotations_out/'
if not (os.path.exists(rotated_imgpath) and os.path.exists(rotated_xmlpath)):os.mkdir(rotated_imgpath)os.mkdir(rotated_xmlpath)
for angle in (15,30, 45, 60, 75, 90, 105, 120):for i in os.listdir(imgpath):a, b = os.path.splitext(i)                            #分离出文件名aimg = cv2.imread(imgpath + a + '.jpg')rotated_img = rotate_image(img,angle)cv2.imwrite(rotated_imgpath + a + '_'+ str(angle) +'d.jpg',rotated_img)print (str(i) + ' has been rotated for '+ str(angle)+'°')tree = ET.parse(xmlpath + a + '.xml')root = tree.getroot()for box in root.iter('bndbox'):xmin = float(box.find('xmin').text)ymin = float(box.find('ymin').text)xmax = float(box.find('xmax').text)ymax = float(box.find('ymax').text)x, y, w, h = rotate_xml(img, xmin, ymin, xmax, ymax, angle)cv2.rectangle(rotated_img, (x, y), (x+w, y+h), [0, 0, 255], 2)   #可在该步骤测试新画的框位置是否正确box.find('xmin').text = str(x)box.find('ymin').text = str(y)box.find('xmax').text = str(x+w)box.find('ymax').text = str(y+h)tree.write(rotated_xmlpath + a + '_'+ str(angle) +'d.xml')cv2.imwrite(rotated_imgpath + a + '_' + str(angle) + 'd.jpg', rotated_img)print (str(a) + '.xml has been rotated for '+ str(angle)+'°')

二.文本框旋转

import  random
import cv2
#转换成顺时针的四个点
def order_point(pts):rect = np.zeros((4, 2), dtype="float32")# the top-left point will have the smallest sum, whereas# the bottom-right point will have the largest sums = np.sum(pts, axis=1)rect[0] = pts[np.argmin(s)]rect[2] = pts[np.argmax(s)]# the top-right point will have the smallest difference,# whereas the bottom-left will have the largest differenced = np.diff(pts, axis=1)rect[1] = pts[np.argmin(d)]rect[3] = pts[np.argmax(d)]return rectdef random_rotate(img,random_angle = 20):angle = random.random() * 2 * random_angle - random_anglew, h = img.shape[:2]rotation_matrix = cv2.getRotationMatrix2D((h / 2, w / 2), angle, 1)img_rotation = cv2.warpAffine(img, rotation_matrix, (h, w))return img_rotation,rotation_matrixdef cal_affine_coord(ori_coord,M):x = ori_coord[0]y = ori_coord[1]_x = x * M[0, 0] + y * M[0, 1] + M[0, 2]_y = x * M[1, 0] + y * M[1, 1] + M[1, 2]return [int(_x), int(_y)]def get_rotate_img_boxes(random_angle = 20):img_file = './ch4_test_images/img_1.jpg'txt_file = './test_gts/img_1.txt'img = cv2.imread(img_file)fid = open(txt_file, 'r', encoding='utf-8')bboxes = []for line in fid.readlines():line = line.strip().replace('\ufeff', '').split(',')# print('====line', line)line = line[:8]line = list(map(int,line))line = np.array(line)line = line.reshape(4, 2)# print('====line', line)line = cv2.minAreaRect(line)#中心(x,y), (宽,高), 旋转角度）# print('====line', line)line = cv2.boxPoints(line).astype(np.int)#逆时针从右下角开始走# print('====line', line)line = order_point(line)#顺时针从左下角开始走# print('====line', line)bboxes.append(line)# print('===bboxes:', bboxes)img1, M = random_rotate(img, random_angle)new_all_rects = []for item in bboxes:rect = []for coord in item:# print('==coord', coord)rotate_coord = cal_affine_coord(coord,M)rect.append(rotate_coord)# debug showcv2.polylines(img1, [np.array(rect).reshape(-1, 1, 2)], True, (0, 255, 0), thickness=2)new_all_rects.append(np.array(rect).reshape(8))print('=====new_all_rects', new_all_rects)cv2.imwrite('./img_rotate.jpg', img1)return img1, new_all_rects

txt:

点：

933,255,954,255,956,277,936,277,###
172,323,195,324,195,339,177,339,###
83,270,118,271,115,294,88,291,###
940,310,962,310,962,320,940,320,###
946,356,976,351,978,368,950,374,###
940,322,962,322,964,333,943,334,###
128,344,210,342,206,361,128,362,###
312,303,360,303,360,312,312,312,###

输出：