DOTA2Train.py

import os
import numpy as np
import cv2

label_path = '/opt/datasets/DOTA/labelTxt/'
store_path = '/opt/datasets/DOTA/Annotations/'

if not os.path.exists(store_path):
    os.makedirs(store_path)

classes = ['plane', 'baseball-diamond', 'bridge', 'ground-track-field', 'small-vehicle', 'large-vehicle',
           'ship', 'tennis-court', 'basketball-court', 'storage-tank', 'soccer-ball-field', 'roundabout',
           'harbor', 'swimming-pool', 'helicopter']

label_list = os.listdir(label_path)

def minAreaRect2longSideFormat(rectangle_inf):
    width = rectangle_inf[1][0]
    height = rectangle_inf[1][1]
    theta = rectangle_inf[-1]
    longSide = max(width, height)
    shortSide = min(width, height)
    if theta == 90:
        if longSide == width:
            pass
        else:
            theta = 0
        if np.around(longSide, 2) == np.around(shortSide, 2):
            theta = 0
    else:
        if np.around(longSide, 2) == np.around(shortSide, 2):
            pass
        else:
            if longSide == width:
                pass
            else:
                theta += 90

    if 179 < theta <= 180:
        theta = 179

    return (rectangle_inf[0], (longSide, shortSide), theta)


for label in label_list:
    boxes = []
    label_info_path = label_path + label
    label_store_path = store_path + label
    with open(label_info_path) as label_info:
        new_boxes = []
        new_boxes1 = []
        for box in label_info.readlines():
            box = list(map(str, box.split()))
            corners = np.array([[int(float(box[0])), int(float(box[1]))], [int(float(box[2])), int(float(box[3]))],
                                [int(float(box[4])), int(float(box[5]))], [int(float(box[6])), int(float(box[7]))]])

            corners_sort = corners[np.argsort(corners[:, 1])]
            corners_sorted = corners_sort[np.append(np.argsort(
                corners_sort[0:2, 0]), np.argsort(corners_sort[2:4, 0])[::-1] + 2)]

            if corners_sorted[0, 0] == corners_sorted[1, 0]:
                temp_max = corners_sorted[0, 1]
                if corners_sorted[0, 0] < min(corners_sorted[2, 0], corners_sorted[3, 0]):
                    if temp_max > corners_sorted[1, 1]:
                        pass
                    else:
                        corners_sorted = corners_sorted[[1, 0, 2, 3]]
                else:
                    if temp_max > corners_sorted[1, 1]:
                        corners_sorted = corners_sorted[[1, 0, 2, 3]]
                    else:
                        pass

            if corners_sorted[2, 0] == corners_sorted[3, 0]:
                temp_max = corners_sorted[2, 1]
                if corners_sorted[2, 0] < min(corners_sorted[0, 0], corners_sorted[1, 0]):
                    if temp_max > corners_sorted[3, 1]:
                        pass
                    else:
                        corners_sorted = corners_sorted[[0, 1, 3, 2]]
                else:
                    if temp_max > corners_sorted[3, 1]:
                        corners_sorted = corners_sorted[[0, 1, 3, 2]]
                    else:
                        pass

            corners_x_max = corners[np.argmax(corners[:, 0])][0]
            corners_y_max = corners[np.argmax(corners[:, 1])][1]
            corners_x_min = corners[np.argmin(corners[:, 0])][0]
            corners_y_min = corners[np.argmin(corners[:, 1])][1]
            corners_fixed = corners_sorted

            if corners_x_max == corners_sorted[1, 0] and corners_x_min == corners_sorted[
                3, 0] and corners_y_max == corners_sorted[2, 1] and corners_y_min == corners_sorted[0, 1]:
                # 0度 正矩形
                pass
                # corners_fixed = corners_fixed[[1, 2, 3, 0]]

            elif corners_x_max == corners_sorted[2, 0] and corners_x_min == corners_sorted[
                0, 0] and corners_y_max == corners_sorted[3, 1] and corners_y_min == corners_sorted[1, 1]:

                corners_fixed = corners_fixed[[1, 2, 3, 0]]


            elif corners_x_max == corners_sorted[1, 0] and corners_x_min == corners_sorted[
                3, 0] and corners_y_max == corners_sorted[3, 1] and corners_y_min == corners_sorted[1, 1]:

                corners_fixed[0, 1] = corners_y_min
                corners_fixed[2, 1] = corners_y_max

            elif corners_x_max == corners_sorted[2, 0] and corners_x_min == corners_sorted[
                0, 0] and corners_y_max == corners_sorted[2, 1] and corners_y_min == corners_sorted[0, 1]:

                corners_fixed[1, 0] = corners_x_max
                corners_fixed[3, 0] = corners_x_min

            elif corners_x_max == corners_sorted[1, 0] and corners_x_min == corners_sorted[
                3, 0] and corners_y_max == corners_sorted[3, 1] and corners_y_min == corners_sorted[0, 1]:

                corners_fixed[2, 1] = corners_y_max

            elif corners_x_max == corners_sorted[1, 0] and corners_x_min == corners_sorted[
                3, 0] and corners_y_max == corners_sorted[2, 1] and corners_y_min == corners_sorted[1, 1]:

                corners_fixed[0, 1] = corners_y_min

            elif corners_x_max == corners_sorted[2, 0] and corners_x_min == corners_sorted[
                0, 0] and corners_y_max == corners_sorted[3, 1] and corners_y_min == corners_sorted[0, 1]:

                corners_fixed[1, 1] = corners_y_min

                corners_fixed = corners_fixed[[1, 2, 3, 0]]


            elif corners_x_max == corners_sorted[2, 0] and corners_x_min == corners_sorted[
                0, 0] and corners_y_max == corners_sorted[2, 1] and corners_y_min == corners_sorted[1, 1]:

                corners_fixed[3, 1] = corners_y_max

                corners_fixed = corners_fixed[[1, 2, 3, 0]]

            elif corners_x_max == corners_sorted[1, 0] and corners_x_min == corners_sorted[
                0, 0] and corners_y_max == corners_sorted[2, 1] and corners_y_min == corners_sorted[0, 1]:

                corners_fixed[3, 0] = corners_x_min

            elif corners_x_max == corners_sorted[1, 0] and corners_x_min == corners_sorted[
                0, 0] and corners_y_max == corners_sorted[2, 1] and corners_y_min == corners_sorted[1, 1]:

                corners_fixed[0, 1] = corners_y_min
                corners_fixed[3, 0] = corners_x_min

            elif corners_x_max == corners_sorted[1, 0] and corners_x_min == corners_sorted[
                0, 0] and corners_y_max == corners_sorted[3, 1] and corners_y_min == corners_sorted[0, 1]:

                corners_fixed[2, 1] = corners_y_max
                corners_fixed[3, 0] = corners_x_min

            elif corners_x_max == corners_sorted[1, 0] and corners_x_min == corners_sorted[
                0, 0] and corners_y_max == corners_sorted[3, 1] and corners_y_min == corners_sorted[1, 1]:

                corners_fixed[2, 0] = corners_x_max

                corners_fixed = corners_fixed[[1, 2, 3, 0]]

            # X_max Xmin Y_max Y_min => X3 X4 X3 X1 @ 13
            elif corners_x_max == corners_sorted[2, 0] and corners_x_min == corners_sorted[
                3, 0] and corners_y_max == corners_sorted[2, 1] and corners_y_min == corners_sorted[0, 1]:

                corners_fixed[1, 0] = corners_x_max

            elif corners_x_max == corners_sorted[2, 0] and corners_x_min == corners_sorted[
                3, 0] and corners_y_max == corners_sorted[2, 1] and corners_y_min == corners_sorted[1, 1]:

                corners_fixed[0, 1] = corners_y_min
                corners_fixed[1, 0] = corners_x_max

            elif corners_x_max == corners_sorted[2, 0] and corners_x_min == corners_sorted[
                3, 0] and corners_y_max == corners_sorted[3, 1] and corners_y_min == corners_sorted[0, 1]:

                corners_fixed[1, 0] = corners_x_max
                corners_fixed[2, 1] = corners_y_max

            elif corners_x_max == corners_sorted[2, 0] and corners_x_min == corners_sorted[
                3, 0] and corners_y_max == corners_sorted[3, 1] and corners_y_min == corners_sorted[1, 1]:

                corners_fixed[0, 0] = corners_x_min

                corners_fixed = corners_fixed[[1, 2, 3, 0]]
                
            basic_info = [corners_x_min, corners_y_min,
                          corners_x_max, corners_y_max, classes.index(box[8])]
            corners_info = [x for corner in corners_fixed.tolist()
                            for x in corner]
            all_info = basic_info + corners_info

            if box[9] != '2':
                r = cv2.contourArea(corners_fixed) / (float(corners_x_max - corners_x_min) * float(corners_y_max - corners_y_min))
                all_info.append(r)

                cnt = np.array([[int(corners_fixed[0][0]), int(str(corners_fixed[0][1]))],
                                [int(corners_fixed[1][0]), int(corners_fixed[1][1])],
                                [int(corners_fixed[2][0]), int(corners_fixed[2][1])],
                                [int(corners_fixed[3][0]), int(corners_fixed[3][1])]])
                rect = cv2.minAreaRect(cnt)
                longSide_inf = minAreaRect2longSideFormat(rect)
                angle = longSide_inf[-1]

                all_info.append(angle)

                new_boxes.append(",".join([str(x) for x in all_info]))

                labelout = str(classes.index(box[8])) + ' ' + str(corners_fixed[0][0]) + ' ' + \
                           str(corners_fixed[0][1]) + ' ' + str(corners_fixed[1][0]) + ' ' + \
                           str(corners_fixed[1][1]) + ' ' + str(corners_fixed[2][0]) + ' ' + \
                           str(corners_fixed[2][1]) + ' ' + str(corners_fixed[3][0]) + ' ' + \
                           str(corners_fixed[3][1]) + '\n'
                
                with open(store_path, 'a') as file1:
                    file1.writelines(str(labelout))

        each_boxes = ("/opt/datasets/DOTA/JPEGImages/" + label.split('.')[0] + '.'
                      + label.split('.')[1] + '.png' + " " + " ".join(new_boxes) + "\n")
        temp = "/opt/datasets/DOTA/JPEGImages/" + label.split('.')[0] + '.' + label.split('.')[1] + '.png' + " " + "\n"
        if each_boxes != temp:
            boxes.append(each_boxes)

    with open("/opt/datasets/DOTA/train.txt", 'a') as file:
        file.writelines(boxes)