python 圖像和bounding box隨機旋轉

import random
import numpy as np
import cv2
import matplotlib.pyplot as plt


class RandomRotate(object):
    def __init__(self, max_degree, rotate_ratio=0.5, mean=(103.939, 116.779, 123.68)):
        assert isinstance(max_degree, int)
        self.max_degree = max_degree
        self.ratio = rotate_ratio
        self.mean = mean

    def bbox_rotate(self, bbox, M, img_shape):
        """Flip bboxes horizontally.
        Args:
            bbox(list): [left, right, up, down]
            img_shape(tuple): (height, width)
        """
        assert len(bbox) == 4
        a = M[:, :2]  ##a.shape (2,2)
        b = M[:, 2:]  ###b.shape(2,1)
        b = np.reshape(b, newshape=(1, 2))
        a = np.transpose(a)

        [left, right, up, down] = bbox
        corner_point = np.array([[left, up], [right, up], [left, down], [right, down]])
        corner_point = np.dot(corner_point, a) + b
        min_left = max(int(np.min(corner_point[:, 0])), 0)
        max_right = min(int(np.max(corner_point[:, 0])), img_shape[1])
        min_up = max(int(np.min(corner_point[:, 1])), 0)
        max_down = min(int(np.max(corner_point[:, 1])), img_shape[0])

        return [min_left, max_right, min_up, max_down]

    def __call__(self, image, bbox=None):
        height, width, _ = image.shape
        if random.random() < self.ratio:
            angle = random.uniform(-self.max_degree, self.max_degree)
            rotate_matrix = cv2.getRotationMatrix2D((width / 2, height / 2), angle, 1.0)
            image = cv2.warpAffine(image, rotate_matrix, (width, height), flags=cv2.INTER_LINEAR,
                                   borderMode=cv2.BORDER_CONSTANT, borderValue=self.mean)

            if bbox:
                bbox = self.bbox_rotate(bbox, rotate_matrix, (height, width))
        return image, bbox


img = cv2.imread(r'C:\Users\nickccnie\Desktop\pdf_processs\page_1.png')
bbox = [180, 1350, 480, 1200]

ro = RandomRotate(max_degree=50, rotate_ratio=1, mean=(255, 0, 0))

new_img, new_bbox = ro(img, bbox=bbox)
img = cv2.rectangle(img, (bbox[0], bbox[2]), (bbox[1], bbox[3]), (0, 255, 0), 2)
new_img = cv2.rectangle(new_img, (new_bbox[0], new_bbox[2]), (new_bbox[1], new_bbox[3]), (0, 255, 0), 2)


plt.subplot(121)
plt.imshow(img)
plt.subplot(122)
plt.imshow(new_img)
plt.show()