图片和box旋转

import os,cv2,shutil,math
import numpy as np

root = '/home/data/general_ocr_det_data'

#root = '/home/xx/data/general_det_data'
def rotate_about_center(src, angle, scale=1.):
    """https://www.oschina.net/translate/opencv-rotation"""
    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]
    return cv2.warpAffine(src, rot_mat, (int(math.ceil(nw)), int(math.ceil(nh))), flags=cv2.INTER_LANCZOS4), rot_mat

def rotate_image(image, xs, ys):
    rotation_angle = np.random.randint(low = -30, high = 30);
    #scale = np.random.uniform(low = MIN_ROTATION_SCLAE, high = MAX_ROTATION_SCLAE)
    scale = 1.0
    h, w = image.shape[0:2]
    # rotate image
    image, M = rotate_about_center(image, rotation_angle, scale = scale)
    nh, nw = image.shape[0:2]

    #rotate bboxes
    #axs = xs * w
    #ys = ys * h
    def rotate_xys(xs, ys):
        xs = np.reshape(xs, -1)
        ys = np.reshape(ys, -1)
        xs, ys = np.dot(M, np.transpose([xs, ys, 1]))
        xs = np.reshape(xs, (-1, 4))
        ys = np.reshape(ys, (-1, 4))
        return xs, ys
    #scale xs,ys to [0,1]
    xs, ys = rotate_xys(xs, ys)
    #xs = xs * 1.0 / nw
    #ys = ys * 1.0 / nh
    ys = ys * 1.0
    ys = ys * 1.0
    xmin = np.min(xs, axis = 1)
    xmin[np.where(xmin < 0)] = 0

    xmax = np.max(xs, axis = 1)
    xmax[np.where(xmax > 1)] = 1

    ymin = np.min(ys, axis = 1)
    ymin[np.where(ymin < 0)] = 0

    ymax = np.max(ys, axis = 1)
    ymax[np.where(ymax > 1)] = 1
    
    bboxes = np.transpose(np.asarray([ymin, xmin, ymax, xmax]))
    image = np.asarray(image, np.uint8)
    return image, bboxes, xs, ys


folders = [os.path.join(root,d) for d in os.listdir(root)]
for dd in folders:
    print(dd)
    files = [f for f in os.listdir(dd) if f.endswith('.jpg')]
    if not os.path.exists(os.path.join('/home/xx/data/general_det_data_rotate/',dd.split('/')[-1])):
        os.makedirs(os.path.join('/home/xx/data/general_det_data_rotate/',dd.split('/')[-1]))
    for f in files:
        path = os.path.join(dd,f)
        #try:
        #    im = Image.open(path)
        #except:
        #    print(path)
        img =cv2.imread(path)
        if img is None:
            print(path)
        oriented_boxes = []
        xs_all = []
        ys_all = []
        texts = []
        with open(path.replace('.jpg','.txt'),'r') as txt:
            lines = txt.readlines()
            for line in lines:
                line = line.replace('\xef\xbb\xbf','').strip()
                gt = line.split(',')
                oriented_box = [int(gt[i]) for i in range(8)]
                oriented_box = np.asarray(oriented_box)
                oriented_boxes.append(oriented_box)
                xs = oriented_box.reshape(4, 2)[:, 0]
                ys = oriented_box.reshape(4, 2)[:, 1]
                texts.append(line.split(',')[-1])
                xs_all.append(xs)
                ys_all.append(ys)
        xs_all = np.array(xs_all)
        ys_all = np.array(ys_all)
        im,bboxes, xs_out,ys_out = rotate_image(img,xs_all,ys_all)
        #oriented_bboxes = np.transpose(np.stack([xs_out,ys_out],1)).reshape((-1,8))
        output_path = '/home/xx/data/general_det_data_rotate/'+dd.split('/')[-1]+'/'
        cv2.imwrite(output_path + f, im)
        
        with open(output_path + f.replace('.jpg','.txt'),'w') as out:
            for idx, (xx,yy) in enumerate(zip(xs_out,ys_out)):
                #print(box,texts[idx])
                out.write(str(int(xx[0]))+','+str(int(yy[0]))+','+\
                        str(int(xx[1]))+','+str(int(yy[1]))+','+\
                        str(int(xx[2]))+','+str(int(yy[2]))+','+\
                        str(int(xx[3]))+','+str(int(yy[3]))+ ','+ texts[idx]+'\n')