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

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

#coding:utf-8
#数据集扩增
import cv2
import math
import numpy as np
import xml.etree.ElementTree as ET
import os

def 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 height
    nw = (abs(np.sin(rangle) * h) + abs(np.cos(rangle) * w)) * scale
    nh = (abs(np.cos(rangle) * h) + abs(np.sin(rangle) * w)) * scale
    # ask OpenCV for the rotation matrix
    rot_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 rotation
    rot_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 transform
    rot_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))*scale
    nh = (abs(np.cos(rangle)*h) + abs(np.sin(rangle)*w))*scale
    # ask OpenCV for the rotation matrix
    rot_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 rotation
    rot_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 transform
    rot_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+rh
    return 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)                            #分离出文件名a

        img = 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 sum
    s = 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 difference
    d = np.diff(pts, axis=1)
    rect[1] = pts[np.argmin(d)]
    rect[3] = pts[np.argmax(d)]
    return rect

def random_rotate(img,random_angle = 20):
    angle = random.random() * 2 * random_angle - random_angle
    w, 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_matrix

def 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 show
        cv2.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,###

输出：