-
Notifications
You must be signed in to change notification settings - Fork 43
/
show_learning_points_and_boxes.py
211 lines (176 loc) · 7.12 KB
/
show_learning_points_and_boxes.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
# -*- coding: utf-8 -*-
"""
@author: liwentong
"""
import pickle
import cv2
import json
import os
import shutil
import numpy as np
detection_pkl_path = './orientedreppoints_r50_demo.pkl'
val_json = './test_dota.json'
outpath = './vis_result/'
inputpath = './test_images/'
if os.path.exists(outpath):
shutil.rmtree(outpath) # delete output folderX
os.makedirs(outpath)
def transfer_to_order_point(corner_points):
a_x = corner_points[0][0]
a_y = corner_points[0][1]
b_x = corner_points[1][0]
b_y = corner_points[1][1]
c_x = corner_points[2][0]
c_y = corner_points[2][1]
d_x = corner_points[3][0]
d_y = corner_points[3][1]
# top x0 y0
if a_y == min(a_y, b_y, c_y, d_y) :
x0 = a_x
y0 = a_y
if b_y == min(a_y, b_y, c_y, d_y) :
x0 = b_x
y0 = b_y
if c_y == min(a_y, b_y, c_y, d_y) :
x0 = c_x
y0 = c_y
if d_y == min(a_y, b_y, c_y, d_y) :
x0 = d_x
y0 = d_y
# right x1 y1
if a_x == max(a_x, b_x, c_x, d_x) :
x1 = a_x
y1 = a_y
if b_x == max(a_x, b_x, c_x, d_x) :
x1 = b_x
y1 = b_y
if c_x == max(a_x, b_x, c_x, d_x) :
x1 = c_x
y1 = c_y
if d_x == max(a_x, b_x, c_x, d_x) :
x1 = d_x
y1 = d_y
# bottom x2 y2
if a_y == max(a_y, b_y, c_y, d_y) :
x2 = a_x
y2 = a_y
if b_y == max(a_y, b_y, c_y, d_y) :
x2 = b_x
y2 = b_y
if c_y == max(a_y, b_y, c_y, d_y) :
x2 = c_x
y2 = c_y
if d_y == max(a_y, b_y, c_y, d_y) :
x2 = d_x
y2 = d_y
# left x3 y3
if a_x == min(a_x, b_x, c_x, d_x) :
x3 = a_x
y3 = a_y
if b_x == min(a_x, b_x, c_x, d_x) :
x3 = b_x
y3 = b_y
if c_x == min(a_x, b_x, c_x, d_x) :
x3 = c_x
y3 = c_y
if d_x == min(a_x, b_x, c_x, d_x) :
x3 = d_x
y3 = d_y
order_points = np.array([(x0,y0),(x1,y1),(x2,y2),(x3,y3)], np.int32)
return order_points
#open result pkl
with open(detection_pkl_path, 'rb') as file:
while True:
try:
data = pickle.load(file)
except EOFError:
break
num_img = len(data)
for iter_img in range(num_img):
# print(iter_img)
img_results = data[iter_img]
#open json
with open(val_json) as f:
ann=json.load(f)
for img_item in ann['images']:
print(img_item['id'])
if iter_img + 1 == img_item['id']:
img_name = img_item['file_name']
img_base_name = img_name.split('.png')[0]
bboxes = img_results
num_bboxes = len(bboxes)
for iter_box in range(num_bboxes):
if len(bboxes[iter_box])>0:
print(bboxes[iter_box])
if iter_box == 0:
class_name = 'plane'
r_color = (0, 0, 255) #纯红
elif iter_box == 1:
class_name = 'baseball-diamond'
r_color = (0, 255, 0) #纯绿
elif iter_box == 2:
class_name = 'bridge'
r_color = (255, 0, 0) # 纯蓝
elif iter_box == 3:
class_name = 'ground-track-field'
r_color = (169, 169, 169) # 深灰色
elif iter_box == 4:
class_name = 'small-vehicle'
r_color = (0, 0, 139) #深红色
elif iter_box == 5:
class_name = 'large-vehicle'
r_color = (0, 69, 255) #橙红色
elif iter_box == 6:
class_name = 'ship'
r_color = (30, 105, 210)# 巧克力色
elif iter_box == 7:
class_name = 'tennis-court'
r_color = (10,215,255) #金色
elif iter_box == 8:
class_name = 'basketball-court'
r_color = (0,128,128) #橄榄色
elif iter_box == 9:
class_name = 'storage-tank'
r_color = (203,192,255) #粉色
elif iter_box == 10:
class_name = 'soccer-ball-field'
r_color = (230,216,173) #天蓝色
elif iter_box == 11:
class_name = 'roundabout'
r_color = (238,130,238) #紫罗兰
elif iter_box == 12:
class_name = 'harbor'
r_color = (144,238,144) #浅绿
elif iter_box == 13:
class_name = 'swimming-pool'
r_color = (130,0,75) #藏青色
elif iter_box == 14:
class_name = 'helicopter'
r_color = (147,20,255) #深粉色
else:
class_name = None
for bbox in bboxes[iter_box]:
pt_x = bbox[0:-9:2]
pt_y = bbox[1:-9:2]
confidence = float(bbox[-1])
#confidence threshold
if confidence > 0.3:
if os.path.exists(outpath+img_name):
print(img_name)
image = cv2.imread(outpath+img_name)
else:
image = cv2.imread(inputpath+img_name)
# visulize the oriented boxes
box_list = []
box_list.append((float(bbox[-9]), float(bbox[-8])))
box_list.append((float(bbox[-7]), float(bbox[-6])))
box_list.append((float(bbox[-5]), float(bbox[-4])))
box_list.append((float(bbox[-3]), float(bbox[-2])))
box_order = transfer_to_order_point(box_list)
# cv2.polylines(image, [box_order], True, r_color, thickness=3)
for i in range(len(pt_x)):
# learning points
cv2.circle(image, (pt_x[i], pt_y[i]), 1, r_color, thickness=2)
# cv2.putText(image, class_name, (int(box_order[0][0]), int(box_order[0][1]) - 2),
# cv2.FONT_HERSHEY_COMPLEX, 0.5, r_color)
cv2.imwrite(outpath+img_name, image)