目标检测透视变换数据增强包含label(yolov5格式)

# -*- coding:utf-8 -*-
import cv2
import numpy as np
import sys


def rad(x):
    return x * np.pi / 180


def get_warpR():
    global anglex, angley, anglez, fov, W, H, r
    # 镜头与图像间的距离，21为半可视角，算z的距离是为了保证在此可视角度下恰好显示整幅图像
    z = np.sqrt(H ** 2 + W ** 2) / 2 / np.tan(rad(fov / 2))
    # 齐次变换矩阵
    rx = np.array([[1, 0, 0, 0],
                   [0, np.cos(rad(anglex)), -np.sin(rad(anglex)), 0],
                   [0, -np.sin(rad(anglex)), np.cos(rad(anglex)), 0, ],
                   [0, 0, 0, 1]], np.float32)

    ry = np.array([[np.cos(rad(angley)), 0, np.sin(rad(angley)), 0],
                   [0, 1, 0, 0],
                   [-np.sin(rad(angley)), 0, np.cos(rad(angley)), 0, ],
                   [0, 0, 0, 1]], np.float32)

    rz = np.array([[np.cos(rad(anglez)), np.sin(rad(anglez)), 0, 0],
                   [-np.sin(rad(anglez)), np.cos(rad(anglez)), 0, 0],
                   [0, 0, 1, 0],
                   [0, 0, 0, 1]], np.float32)

    r = rx.dot(ry).dot(rz)

    # 四对点的生成
    pcenter = np.array([W / 2, H / 2, 0, 0], np.float32)

    p1 = np.array([0, 0, 0, 0], np.float32) - pcenter
    p2 = np.array([H, 0, 0, 0], np.float32) - pcenter
    p3 = np.array([0, W, 0, 0], np.float32) - pcenter
    p4 = np.array([H, W, 0, 0], np.float32) - pcenter

    dst1 = r.dot(p1)
    dst2 = r.dot(p2)
    dst3 = r.dot(p3)
    dst4 = r.dot(p4)

    list_dst = [dst1, dst2, dst3, dst4]

    org = np.array([[0, 0],
                    [H, 0],
                    [0, W],
                    [H, W]], np.float32)

    dst = np.zeros((4, 2), np.float32)

    # 投影至成像平面
    for i in range(4):
        dst[i, 0] = list_dst[i][0] * z / (z - list_dst[i][2]) + pcenter[0]
        dst[i, 1] = list_dst[i][1] * z / (z - list_dst[i][2]) + pcenter[1]

    warpR = cv2.getPerspectiveTransform(org, dst)
    return warpR

# ccc = 0
# for i in range(1,8):
#     for angley in [i*10,-i*10]:
#         labels = []
#         with open('/home/lixuan/Arduino/sku_recogniztion/imgs/000368.txt') as f:
#             boxs = []
#             datas = f.readlines()
#             for data in datas:
#                 label = []
#                 data = data.strip()
#                 data = data.split(' ')[1:]
#                 x = float(data[0]) * W
#                 y = float(data[1]) * H
#                 w = float(data[2]) * W
#                 h = float(data[3]) * H
#                 x1 = int(x - w / 2)
#                 y1 = int(y - h / 2)
#                 x2 = int(x + w / 2)
#                 y2 = int(y + h / 2)
#                 label.append([x1, y1])
#                 label.append([x2, y1])
#                 label.append([x2, y2])
#                 label.append([x1, y2])
#                 xs = []
#                 ys = []
#                 for xys in label:
#                     xs.append(xys[0])
#                     ys.append(xys[1])
#                 contour = np.vstack((xs, ys))  # 将x，y竖着拼接
#                 contour = np.array(contour.T, dtype=np.float32)
#                 box = cv2.boundingRect(contour)
#                 boxs.append([int(box[0]), int(box[1]), int(box[0] + box[2]), int(box[1] + box[3])])
#             bboxes = []
#             for box in boxs:
#                 bboxes.append([[box[0],box[1]],[box[2],box[3]]])
#             #     # img = cv2.rectangle(img, (box[0], box[1]), (box[2], box[3]), (0, 255, 0), 2)
#
#             warpR = get_warpR()
#
#             result = cv2.warpPerspective(img, warpR, (W, H))
#
#             matrix = warpR
#             for label in bboxes:
#                 box = []
#                 for p in label:
#                     px = (matrix[0][0] * p[0] + matrix[0][1] * p[1] + matrix[0][2]) / (
#                         (matrix[2][0] * p[0] + matrix[2][1] * p[1] + matrix[2][2]))
#                     py = (matrix[1][0] * p[0] + matrix[1][1] * p[1] + matrix[1][2]) / (
#                         (matrix[2][0] * p[0] + matrix[2][1] * p[1] + matrix[2][2]))
#                     p_after = (int(px), int(py))  # after transformation
#                     box.append(list(p_after))
#                 if box[0][0] < 0:
#                     box[0][0] = 0
#                 if box[0][1] < 0:
#                     box[0][1] = 0
#                 if box[1][0] < 0:
#                     box[1][0] = 0
#                 if box[1][1] < 0:
#                     box[1][1] = 0
#                 if box[0][0] > W:
#                     box[0][0] = W
#                 if box[1][0] > W:
#                     box[1][0] = W
#                 if box[0][1] > H:
#                     box[0][1] = H
#                 if box[1][1] > H:
#                     box[1][1] = H
#                 result = cv2.rectangle(result, (box[0][0], box[0][1]), (box[1][0], box[1][1]), (0, 255, 0), 2)
#                     # result = cv2.circle(result, p_after, 10, (0, 0, 255), -1)
#
#             # cv2.namedWindow('result', 2)
#             # cv2.imshow("result", result)
#             # c = cv2.waitKey(0)
#             cv2.imwrite('/home/lixuan/sku_data/cailiao/{}.jpg'.format(ccc),result)
#             ccc += 1



import os
ccc = 0
anglex = 0
angley = 0
anglez = 0  # 是旋转
fov = 42
r = 0
for imgfile in os.listdir('/home/lixuan/sku_data/local/images'):
    img = cv2.imread('/home/lixuan/sku_data/local/images/{}'.format(imgfile))
    H, W = img.shape[0:2]
    for i in range(1,8):
        for angley in [i*10,-i*10]:
            labels = []
            with open(os.path.join('/home/lixuan/sku_data/local/labels',imgfile.replace('jpg','txt'))) as f:
                boxs = []
                datas = f.readlines()
                for data in datas:
                    label = []
                    data = data.strip()
                    data = data.split(' ')[1:]
                    x = float(data[0]) * W
                    y = float(data[1]) * H
                    w = float(data[2]) * W
                    h = float(data[3]) * H
                    x1 = int(x - w / 2)
                    y1 = int(y - h / 2)
                    x2 = int(x + w / 2)
                    y2 = int(y + h / 2)
                    label.append([x1, y1])
                    label.append([x2, y1])
                    label.append([x2, y2])
                    label.append([x1, y2])
                    xs = []
                    ys = []
                    for xys in label:
                        xs.append(xys[0])
                        ys.append(xys[1])
                    contour = np.vstack((xs, ys))  # 将x，y竖着拼接
                    contour = np.array(contour.T, dtype=np.float32)
                    box = cv2.boundingRect(contour)
                    boxs.append([int(box[0]), int(box[1]), int(box[0] + box[2]), int(box[1] + box[3])])
                bboxes = []
                for box in boxs:
                    bboxes.append([[box[0],box[1]],[box[2],box[3]]])
                #     # img = cv2.rectangle(img, (box[0], box[1]), (box[2], box[3]), (0, 255, 0), 2)

                warpR = get_warpR()

                result = cv2.warpPerspective(img, warpR, (W, H))

                matrix = warpR
                fw = open('/home/lixuan/sku_data/local/labelscopy/sku_{}.txt'.format(ccc),'w')
                for label in bboxes:
                    box = []
                    for p in label:
                        px = (matrix[0][0] * p[0] + matrix[0][1] * p[1] + matrix[0][2]) / (
                            (matrix[2][0] * p[0] + matrix[2][1] * p[1] + matrix[2][2]))
                        py = (matrix[1][0] * p[0] + matrix[1][1] * p[1] + matrix[1][2]) / (
                            (matrix[2][0] * p[0] + matrix[2][1] * p[1] + matrix[2][2]))
                        p_after = (int(px), int(py))  # after transformation
                        box.append(list(p_after))
                    if box[0][0] < 0:
                        box[0][0] = 0
                    if box[0][1] < 0:
                        box[0][1] = 0
                    if box[1][0] < 0:
                        box[1][0] = 0
                    if box[1][1] < 0:
                        box[1][1] = 0
                    if box[0][0] > W:
                        box[0][0] = W
                    if box[1][0] > W:
                        box[1][0] = W
                    if box[0][1] > H:
                        box[0][1] = H
                    if box[1][1] > H:
                        box[1][1] = H
                    x1 = box[0][0] / W
                    y1 = box[0][1] / H
                    x2 = box[1][0] / W
                    y2 = box[1][1] / H
                    x = (x1+x2)/2
                    y = (y1+y2)/2
                    w = x2-x1
                    h = y2-y1
                    fw.write('0 ' + str(x) + ' ' + str(y) + ' ' + str(w) + ' ' + str(h) + '\n')
                    # result = cv2.rectangle(result, (box[0][0], box[0][1]), (box[1][0], box[1][1]), (0, 255, 0), 2)
                        # result = cv2.circle(result, p_after, 10, (0, 0, 255), -1)

                # cv2.namedWindow('result', 2)
                # cv2.imshow("result", result)
                # c = cv2.waitKey(0)
                fw.close()
                cv2.imwrite('/home/lixuan/sku_data/local/imagescopy/sku_{}.jpg'.format(ccc),result)
                ccc += 1