CoCalc -- dota_poly2rbox.py

hukaixuan19970627
Real-time collaboration for Jupyter Notebooks, Linux Terminals, LaTeX, VS Code, R IDE, and more,
all in one place. Commercial Alternative to JupyterHub.

GitHub Repository: hukaixuan19970627/yolov5_obb
Path: blob/master/DOTA_devkit/dota_poly2rbox.py
Views: ⁴⁷⁵
1
import os
2
import math
3
import argparse
4
import os.path as osp
5

6
import numpy as np
7

8
def cal_line_length(point1, point2):
9
    return math.sqrt(math.pow(point1[0] - point2[0], 2) + math.pow(point1[1] - point2[1], 2))
10

11
def get_best_begin_point_single(coordinate):
12
    x1, y1, x2, y2, x3, y3, x4, y4 = coordinate
13
    xmin = min(x1, x2, x3, x4)
14
    ymin = min(y1, y2, y3, y4)
15
    xmax = max(x1, x2, x3, x4)
16
    ymax = max(y1, y2, y3, y4)
17
    combinate = [[[x1, y1], [x2, y2], [x3, y3], [x4, y4]], [[x2, y2], [x3, y3], [x4, y4], [x1, y1]],
18
                 [[x3, y3], [x4, y4], [x1, y1], [x2, y2]], [[x4, y4], [x1, y1], [x2, y2], [x3, y3]]]
19
    dst_coordinate = [[xmin, ymin], [xmax, ymin], [xmax, ymax], [xmin, ymax]]
20
    force = 100000000.0
21
    force_flag = 0
22
    for i in range(4):
23
        temp_force = cal_line_length(combinate[i][0], dst_coordinate[0]) \
24
            + cal_line_length(combinate[i][1], dst_coordinate[1]) \
25
            + cal_line_length(combinate[i][2], dst_coordinate[2]) \
26
            + cal_line_length(combinate[i][3], dst_coordinate[3])
27
        if temp_force < force:
28
            force = temp_force
29
            force_flag = i
30
    if force_flag != 0:
31
        pass
32
        # print("choose one direction!")
33
    return np.array(combinate[force_flag]).reshape(8)
34

35
def poly2rbox_single(poly):
36
    """
37
    poly:[x0,y0,x1,y1,x2,y2,x3,y3]
38
    to
39
    rrect:[x_ctr,y_ctr,w,h,angle]
40
    """
41
    poly = np.array(poly[:8], dtype=np.float32)
42

43
    pt1 = (poly[0], poly[1])
44
    pt2 = (poly[2], poly[3])
45
    pt3 = (poly[4], poly[5])
46
    pt4 = (poly[6], poly[7])
47

48
    edge1 = np.sqrt((pt1[0] - pt2[0]) * (pt1[0] - pt2[0]) +
49
                    (pt1[1] - pt2[1]) * (pt1[1] - pt2[1]))
50
    edge2 = np.sqrt((pt2[0] - pt3[0]) * (pt2[0] - pt3[0]) +
51
                    (pt2[1] - pt3[1]) * (pt2[1] - pt3[1]))
52

53
    angle = 0
54
    width = 0
55
    height = 0
56

57
    if edge1 > edge2:
58
        width = edge1
59
        height = edge2
60
        angle = np.arctan2(
61
            np.float(pt2[1] - pt1[1]), np.float(pt2[0] - pt1[0]))
62
    elif edge2 >= edge1:
63
        width = edge2
64
        height = edge1
65
        angle = np.arctan2(
66
            np.float(pt4[1] - pt1[1]), np.float(pt4[0] - pt1[0]))
67

68
    if angle > np.pi*3/4:
69
        angle -= np.pi
70
    if angle < -np.pi/4:
71
        angle += np.pi
72

73
    x_ctr = np.float(pt1[0] + pt3[0]) / 2
74
    y_ctr = np.float(pt1[1] + pt3[1]) / 2
75
    rbox = np.array([x_ctr, y_ctr, width, height, angle])
76

77
    return rbox
78

79
def norm_angle(angle, range=[-np.pi / 4, np.pi]):
80
    return (angle - range[0]) % range[1] + range[0]
81

82

83
def poly2rbox_single_v2(poly):
84
    """
85
    poly:[x0,y0,x1,y1,x2,y2,x3,y3]
86
    to
87
    rrect:[x_ctr,y_ctr,w,h,angle]
88
    """
89
    poly = np.array(poly[:8], dtype=np.float32)
90

91
    pt1 = (poly[0], poly[1])
92
    pt2 = (poly[2], poly[3])
93
    pt3 = (poly[4], poly[5])
94
    pt4 = (poly[6], poly[7])
95

96
    edge1 = np.sqrt((pt1[0] - pt2[0]) * (pt1[0] - pt2[0]) +
97
                    (pt1[1] - pt2[1]) * (pt1[1] - pt2[1]))
98
    edge2 = np.sqrt((pt2[0] - pt3[0]) * (pt2[0] - pt3[0]) +
99
                    (pt2[1] - pt3[1]) * (pt2[1] - pt3[1]))
100

101
    angle = 0
102
    width = 0
103
    height = 0
104

105
    if edge1 > edge2:
106
        width = edge1
107
        height = edge2
108
        angle = np.arctan2(
109
            np.float(pt2[1] - pt1[1]), np.float(pt2[0] - pt1[0]))
110
    elif edge2 >= edge1:
111
        width = edge2
112
        height = edge1
113
        angle = np.arctan2(
114
            np.float(pt4[1] - pt1[1]), np.float(pt4[0] - pt1[0]))
115

116
    # if angle > np.pi*3/4:
117
    #     angle -= np.pi
118
    # if angle < -np.pi/4:
119
    #     angle += np.pi
120
    angle = norm_angle(angle)
121

122
    x_ctr = np.float(pt1[0] + pt3[0]) / 2
123
    y_ctr = np.float(pt1[1] + pt3[1]) / 2
124
    
125
    return float(x_ctr), float(y_ctr), float(width), float(height), float(angle)
126

127

128
def poly2rbox_single_v3(poly):
129
    """
130
    poly:[x0,y0,x1,y1,x2,y2,x3,y3]
131
    to
132
    rrect:[x_ctr,y_ctr,w,h,angle]
133
    """
134
    poly = np.array(poly[:8], dtype=np.float32)
135

136
    pt1 = (poly[0], poly[1])
137
    pt2 = (poly[2], poly[3])
138
    pt3 = (poly[4], poly[5])
139
    pt4 = (poly[6], poly[7])
140

141
    edge1 = np.sqrt((pt1[0] - pt2[0]) * (pt1[0] - pt2[0]) +
142
                    (pt1[1] - pt2[1]) * (pt1[1] - pt2[1]))
143
    edge2 = np.sqrt((pt2[0] - pt3[0]) * (pt2[0] - pt3[0]) +
144
                    (pt2[1] - pt3[1]) * (pt2[1] - pt3[1]))
145

146
    max_edge = max(edge1, edge2)
147
    min_edge = min(edge1, edge2)
148
    ratio = max_edge / min_edge
149
    # print(ratio)
150
    if ratio < 1.15:
151

152
        width = max_edge
153
        height = min_edge
154
        angle1 = np.arctan2(np.float(pt2[1] - pt1[1]), np.float(pt2[0] - pt1[0]))
155
        # elif edge2 >= edge1:
156
        angle2 = np.arctan2(np.float(pt4[1] - pt1[1]), np.float(pt4[0] - pt1[0]))
157

158
        angle1_norm = norm_angle(angle1)
159
        angle2_norm = norm_angle(angle2)
160
        # if abs(angle1_norm) > abs(angle2_norm):
161
        #     final_angle = angle2_norm
162
        # else:
163
        #     final_angle = angle1_norm
164
        if abs(angle1_norm) > abs(angle2_norm):
165
            final_angle = angle2_norm
166
        else:
167
            final_angle = angle1_norm
168

169
    else:
170
        final_angle = 0
171
        width = 0
172
        height = 0
173

174
        if edge1 > edge2:
175
            width = edge1
176
            height = edge2
177
            final_angle = np.arctan2(
178
                np.float(pt2[1] - pt1[1]), np.float(pt2[0] - pt1[0]))
179
        elif edge2 >= edge1:
180
            width = edge2
181
            height = edge1
182
            final_angle = np.arctan2(
183
                np.float(pt4[1] - pt1[1]), np.float(pt4[0] - pt1[0]))
184

185
        final_angle = norm_angle(final_angle)
186

187
    x_ctr = np.float(pt1[0] + pt3[0]) / 2
188
    y_ctr = np.float(pt1[1] + pt3[1]) / 2
189

190
    return float(x_ctr), float(y_ctr), float(width), float(height), float(final_angle)
191

192

193
def rbox2poly_single(rrect):
194
    """
195
    rrect:[x_ctr,y_ctr,w,h,angle]
196
    to
197
    poly:[x0,y0,x1,y1,x2,y2,x3,y3]
198
    """
199
    x_ctr, y_ctr, width, height, angle = rrect[:5]
200
    tl_x, tl_y, br_x, br_y = -width/2, -height/2, width/2, height/2
201
    rect = np.array([[tl_x, br_x, br_x, tl_x], [tl_y, tl_y, br_y, br_y]])
202
    R = np.array([[np.cos(angle), -np.sin(angle)],
203
                  [np.sin(angle), np.cos(angle)]])
204
    poly = R.dot(rect)
205
    x0, x1, x2, x3 = poly[0, :4] + x_ctr
206
    y0, y1, y2, y3 = poly[1, :4] + y_ctr
207
    poly = np.array([x0, y0, x1, y1, x2, y2, x3, y3], dtype=np.float32)
208
    poly = get_best_begin_point_single(poly)
209
    return poly
210

211
def convert2rbox(src_path):
212
    image_path = osp.join(src_path, 'images/')
213
    src_label_path = osp.join(src_path, 'labelTxt/')
214
    dst_label_path = osp.join(src_path, 'labelTxtRbox/')
215
    if not osp.exists(dst_label_path):
216
        os.mkdir(dst_label_path)
217

218
    image_list = os.listdir(image_path)
219
    image_list.sort()
220

221
    for image in image_list:
222
        img_name = osp.basename(image)
223
        print(img_name)
224
        ann_name = img_name.split('.')[0]+'.txt'
225
        lab_path = osp.join(src_label_path, ann_name)
226
        dst_path = osp.join(dst_label_path, ann_name)
227
        out_str = ''
228

229
        # import time
230
        # half the time used by poly2rbox
231
        with open(lab_path, 'r') as f:
232
            for ann_line in f.readlines():
233
                ann_line = ann_line.strip().split(' ')
234
                bbox = [np.float32(ann_line[i]) for i in range(8)]
235
                # 8 point to 5 point xywha
236
                x_ctr, y_ctr, width, height, angle = poly2rbox_single(bbox)
237
                class_name = ann_line[8]
238
                difficult = int(ann_line[9])
239

240
                out_str += "{} {} {} {} {} {} {}\n".format(str(x_ctr), str(
241
                    y_ctr), str(width), str(height), str(angle), class_name, difficult)
242
        with open(dst_path, 'w') as fdst:
243
            fdst.write(out_str)
244

245

246
if __name__ == '__main__':
247
    parser = argparse.ArgumentParser()
248
    parser.add_argument('-p', '--path', type=str, required=True)
249
    args = parser.parse_args()
250
    convert2rbox(args.path)
251

252
Real-time collaboration for Jupyter Notebooks, Linux Terminals, LaTeX, VS Code, R IDE, and more,
all in one place. Commercial Alternative to JupyterHub.

Product

Resources

Company

Real-time collaboration for Jupyter Notebooks, Linux Terminals, LaTeX, VS Code, R IDE, and more, all in one place. Commercial Alternative to JupyterHub.

Real-time collaboration for Jupyter Notebooks, Linux Terminals, LaTeX, VS Code, R IDE, and more,
all in one place. Commercial Alternative to JupyterHub.
