数据增强-随机旋转

---

前言

目标检测中对图像进行随机旋转的时候，要考虑到bbox坐标的改变
网上查资料发现都是一头雾水，自己亲自写了一套易于理解的

1.导包

import cv2
from PIL import Image
import torch
import random
import numpy as np
import matplotlib.pyplot as plt

2.核心代码

class RandomRotation(object):
    def __call__(self, image, boxes): #boxes,tensor数据，格式n*4,n为目标个数，4为左上右下坐标 xyxy
        degree = random.uniform(-30, 30)
        print(degree)
        h, w, c = image.shape
        retval = cv2.getRotationMatrix2D((h / 2, w / 2), degree, 0.8)
        image = cv2.warpAffine(image, retval, (w, h), borderValue=(114, 114, 114))
        M11, M12, M13 = retval[0]
        M21, M22, M23 = retval[1]
        x1 = boxes[:, [0]] * M11 + boxes[:, [1]] * M12 + M13  # lt
        y1 = boxes[:, [0]] * M21 + boxes[:, [1]] * M22 + M23

        x2 = boxes[:, [2]] * M11 + boxes[:, [3]] * M12 + M13  # rb
        y2 = boxes[:, [2]] * M21 + boxes[:, [3]] * M22 + M23

        x3 = boxes[:, [0]] * M11 + boxes[:, [3]] * M12 + M13  # lb
        y3 = boxes[:, [0]] * M21 + boxes[:, [3]] * M22 + M23

        x4 = boxes[:, [2]] * M11 + boxes[:, [1]] * M12 + M13  # rt
        y4 = boxes[:, [2]] * M21 + boxes[:, [1]] * M22 + M23

        lx = torch.cat([x1, x2, x3, x4], -1).min(-1)[0]
        ly = torch.cat([y1, y2, y3, y4], -1).min(-1)[0]
        rx = torch.cat([x1, x2, x3, x4], -1).max(-1)[0]
        ry = torch.cat([y1, y2, y3, y4], -1).max(-1)[0]
        boxes = torch.stack([lx, ly, rx, ry], -1)
        boxes[:, ::2].clip_(0, w)
        boxes[:, 1::2].clip_(0, h)
        boxes = boxes.numpy()
        return image, boxes

3.测试代码

image = Image.open("test.jpg")  # 读入图片数据
gt_bbox = torch.tensor([[356, 183, 500, 280], [60, 109, 142, 213], [246, 134, 348, 217]])
img = np.asarray(image)
random_obj = RandomRotation()
new, neb = random_obj(img, gt_bbox)
neb = neb.numpy()
print('\n', new.shape, "\n", gt_bbox)
pt0x, pt0y = int(neb[0][0]), int(neb[0][1])
pt1x, pt1y = int(neb[0][2]), int(neb[0][3])
pt2x, pt2y = int(neb[1][0]), int(neb[1][1])
pt3x, pt3y = int(neb[1][2]), int(neb[1][3])
pt4x, pt4y = int(neb[2][0]), int(neb[2][1])
pt5x, pt5y = int(neb[2][2]), int(neb[2][3])
cv2.rectangle(new, (pt0x, pt0y), (pt1x, pt1y), (255, 255, 0), 2)
cv2.rectangle(new, (pt2x, pt2y), (pt3x, pt3y), (255, 0, 0), 2)
cv2.rectangle(new, (pt4x, pt4y), (pt5x, pt5y), (0, 0, 255), 2)
plt.imshow(new)
plt.show()

4.具体效果

---

总结

以上就是今天要讲的内容，本文仅仅简单介绍了目标检测数据增强之随机旋转的使用