forked from hoangkhoiLE/VisDrone_FPN_ensemble_tensorflow
-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy patheval_visdrone.py
148 lines (130 loc) · 6.76 KB
/
eval_visdrone.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
from models.meta_architectures import BuildFrameworks
import tensorflow as tf
import cv2
import os
import time
import numpy as np
from core.tools import tools
from core.config import cfgs
def txt_visdrone(img_dir, save_path, ensemble_path,normalized_result=False):
# initialize
net = BuildFrameworks(False)
img_plac = tf.placeholder(dtype=tf.uint8, shape=[None, None, 3])
if cfgs.PRETRAIN_BACKBONE_NAME in ['resnet101_v1d', 'resnet50_v1d']:
img_tensor = img_plac / 255 - tf.constant([[cfgs.PIXEL_MEAN_]])
img_tensor = img_tensor / tf.constant([cfgs.PIXEL_STD])
else:
img_tensor = img_plac - tf.constant([[cfgs.PIXEL_MEAN]])
img_tensor = tf.expand_dims(img_tensor, axis=0)
# detection result dict
det_dict = net.build_framworks(input_img_batch=img_tensor)
final_boxes, final_scores, final_category = det_dict[
"final_boxes"], det_dict["final_scores"], det_dict["final_category"]
restorer, restore_ckpt = net.get_restorer()
if restore_ckpt == cfgs.PRETRAIN_BACKBONE_WEIGHTS:
raise IndexError("----->Please select the trained weights rather than pretrained weights.")
img_name = os.listdir(path=img_dir)
with tf.Session() as sess:
restorer.restore(sess, restore_ckpt)
for i in range(len(img_name)):
# read images
img_array = cv2.imread(os.path.join(img_dir, img_name[i]))
raw_h = img_array.shape[0]
raw_w = img_array.shape[1]
detected_scores_, detected_boxes_, detected_categories_ = [], [], []
start=time.time()
if cfgs.MULTI_SCALSE_TESTING:
short_size_len=cfgs.SHORT_SIZE_LIMITATION_LIST
else:
short_size_len=[cfgs.SHORT_SIZE_LIMITATION]
for short_size in short_size_len:
if raw_h < raw_w:
new_h, new_w = short_size, min(int(short_size * float(raw_w) / raw_h), cfgs.MAX_LENGTH)
else:
new_h, new_w = min(int(short_size * float(raw_h) / raw_w), cfgs.MAX_LENGTH), short_size
img_resize = cv2.resize(img_array, (new_w, new_h))
resized_img_, final_boxes_, final_scores_, final_category_ = sess.run(
[img_plac, final_boxes, final_scores, final_category],
feed_dict={img_plac: img_resize[:, :, ::-1]})
#剔除变成一条线的框
inds_inside = np.where(
(final_boxes_[:, 2] > final_boxes_[:, 0]) &
(final_boxes_[:, 3] > final_boxes_[:, 1])
)[0]
final_boxes_ = final_boxes_[inds_inside, :]
final_scores_ = final_scores_[inds_inside]
final_category_ = final_category_[inds_inside]
xmin, ymin, xmax, ymax = final_boxes_[:, 0], final_boxes_[:, 1], \
final_boxes_[:, 2], final_boxes_[:, 3]
resized_h, resized_w = resized_img_.shape[0], resized_img_.shape[1]
# normalized boxes
xmin = xmin / resized_w
xmax = xmax / resized_w
ymin = ymin / resized_h
ymax = ymax / resized_h
resized_boxes = np.transpose(np.stack([xmin, ymin, xmax, ymax]))
detected_scores_.append(final_scores_)
detected_boxes_.append(resized_boxes)
detected_categories_.append(final_category_)
end=time.time()
from ensemble_boxes import weighted_boxes_fusion
boxes, scores, labels = weighted_boxes_fusion(detected_boxes_, detected_scores_, detected_categories_,
weights=None,
iou_thr=0.6, skip_box_thr=0.001, conf_type='avg')
detected_scores = np.array(scores)
detected_boxes = np.array(boxes)
detected_categories = np.array(labels)
if not normalized_result:
# map coordinate to raw picture
xmin, ymin, xmax, ymax = detected_boxes[:, 0] * raw_w, detected_boxes[:, 1] * raw_h, \
detected_boxes[:, 2] * raw_w, detected_boxes[:, 3] * raw_h
boxes = np.transpose(np.stack([xmin, ymin, xmax - xmin, ymax - ymin]))
result_full_path = save_path + \
str('/' + img_name[i].split('.')[0]) + '.txt'
else:
result_full_path = ensemble_path +\
str('/' + img_name[i].split('.')[0]) + '.txt'
if not os.path.exists(save_path):
os.mkdir(save_path)
if not os.path.exists(ensemble_path):
os.makedirs(ensemble_path)
with open(result_full_path, 'w') as fout:
num = 0
if normalized_result:
fout.write(str(raw_w))
fout.write(',')
fout.write(str(raw_h))
fout.write('\n')
for j, box in enumerate(boxes):
if num < 500:
if detected_categories[j] != 1 and detected_categories[j] != 12:
fout.write(str(box[0]))
fout.write(',')
fout.write(str(box[1]))
fout.write(',')
fout.write(str(box[2]))
fout.write(',')
fout.write(str(box[3]))
fout.write(',')
fout.write(str(detected_scores[j]))
fout.write(',')
fout.write(str(detected_categories[j] - 1))
fout.write(',-1,')
fout.write('-1' + '\n')
num += 1
else:
break
fout.close()
tools.view_bar(
'{}, time cost: {}s'.format(
img_name[i],
(end - start)),
i + 1,
len(img_name))
if __name__ == '__main__':
# os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
img_dir = cfgs.IMG_INFERENCE_DIR
save_path = r'./test_result_txt'
ensemble_path = './test_model_ensemble' + '/' + cfgs.PRETRAIN_BACKBONE_NAME +'/'
#if you want to get nomalized results for model ensemble, set normalized_result = True
txt_visdrone(img_dir, save_path,ensemble_path,normalized_result=cfgs.NORMALIZED_RESULTS_FOR_MODEL_ENSEMBLE)