def rnms_gpu(det_boxes, iou_threshold, device_id):
if det_boxes.shape[0] == 0:
return np.array([], np.int64)
else:
keep = rotate_gpu_nms(det_boxes, iou_threshold, device_id)
keep = np.reshape(keep, [-1])
return np.array(keep, np.int64)
def worker(gpu_id, images, det_net, args, result_queue):
os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_id)
img_plac = tf.placeholder(dtype=tf.uint8, shape=[None, None,
3]) # is RGB. not BGR
img_batch = tf.cast(img_plac, tf.float32)
img_batch = short_side_resize_for_inference_data(
img_tensor=img_batch,
target_shortside_len=cfgs.IMG_SHORT_SIDE_LEN,
length_limitation=cfgs.IMG_MAX_LENGTH)
if cfgs.NET_NAME in ['resnet152_v1d', 'resnet101_v1d', 'resnet50_v1d']:
img_batch = (img_batch / 255 - tf.constant(
cfgs.PIXEL_MEAN_)) / tf.constant(cfgs.PIXEL_STD)
else:
img_batch = img_batch - tf.constant(cfgs.PIXEL_MEAN)
img_batch = tf.expand_dims(img_batch, axis=0)
detection_boxes, detection_scores, detection_category = det_net.build_whole_detection_network(
input_img_batch=img_batch, gtboxes_batch_h=None, gtboxes_batch_r=None)
init_op = tf.group(tf.global_variables_initializer(),
tf.local_variables_initializer())
restorer, restore_ckpt = det_net.get_restorer()
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
with tf.Session(config=config) as sess:
sess.run(init_op)
if not restorer is None:
restorer.restore(sess, restore_ckpt)
print('restore model %d ...' % gpu_id)
for img_path in images:
# if 'P0016' not in img_path:
# continue
img = cv2.imread(img_path)
box_res_rotate = []
label_res_rotate = []
score_res_rotate = []
imgH = img.shape[0]
imgW = img.shape[1]
if imgH < args.h_len:
temp = np.zeros([args.h_len, imgW, 3], np.float32)
temp[0:imgH, :, :] = img
img = temp
imgH = args.h_len
if imgW < args.w_len:
temp = np.zeros([imgH, args.w_len, 3], np.float32)
temp[:, 0:imgW, :] = img
img = temp
imgW = args.w_len
for hh in range(0, imgH, args.h_len - args.h_overlap):
if imgH - hh - 1 < args.h_len:
hh_ = imgH - args.h_len
else:
hh_ = hh
for ww in range(0, imgW, args.w_len - args.w_overlap):
if imgW - ww - 1 < args.w_len:
ww_ = imgW - args.w_len
else:
ww_ = ww
src_img = img[hh_:(hh_ + args.h_len),
ww_:(ww_ + args.w_len), :]
resized_img, det_boxes_r_, det_scores_r_, det_category_r_ = \
sess.run(
[img_batch, detection_boxes, detection_scores, detection_category],
feed_dict={img_plac: src_img[:, :, ::-1]}
)
resized_h, resized_w = resized_img.shape[
1], resized_img.shape[2]
src_h, src_w = src_img.shape[0], src_img.shape[1]
if len(det_boxes_r_) > 0:
det_boxes_r_ = forward_convert(det_boxes_r_, False)
det_boxes_r_[:, 0::2] *= (src_w / resized_w)
det_boxes_r_[:, 1::2] *= (src_h / resized_h)
det_boxes_r_ = backward_convert(det_boxes_r_, False)
for ii in range(len(det_boxes_r_)):
box_rotate = det_boxes_r_[ii]
box_rotate[0] = box_rotate[0] + ww_
box_rotate[1] = box_rotate[1] + hh_
box_res_rotate.append(box_rotate)
label_res_rotate.append(det_category_r_[ii])
score_res_rotate.append(det_scores_r_[ii])
box_res_rotate = np.array(box_res_rotate)
label_res_rotate = np.array(label_res_rotate)
score_res_rotate = np.array(score_res_rotate)
box_res_rotate_ = []
label_res_rotate_ = []
score_res_rotate_ = []
threshold = {
'roundabout': 0.1,
'tennis-court': 0.3,
'swimming-pool': 0.1,
'storage-tank': 0.2,
'soccer-ball-field': 0.3,
'small-vehicle': 0.2,
'ship': 0.05,
'plane': 0.3,
'large-vehicle': 0.1,
'helicopter': 0.2,
'harbor': 0.0001,
'ground-track-field': 0.3,
'bridge': 0.0001,
'basketball-court': 0.3,
'baseball-diamond': 0.3
}
for sub_class in range(1, cfgs.CLASS_NUM + 1):
index = np.where(label_res_rotate == sub_class)[0]
if len(index) == 0:
continue
tmp_boxes_r = box_res_rotate[index]
tmp_label_r = label_res_rotate[index]
tmp_score_r = score_res_rotate[index]
tmp_boxes_r = np.array(tmp_boxes_r)
tmp = np.zeros(
[tmp_boxes_r.shape[0], tmp_boxes_r.shape[1] + 1])
tmp[:, 0:-1] = tmp_boxes_r
tmp[:, -1] = np.array(tmp_score_r)
try:
inx = nms_rotate.nms_rotate_cpu(
boxes=np.array(tmp_boxes_r),
scores=np.array(tmp_score_r),
iou_threshold=threshold[LABEL_NAME_MAP[sub_class]],
max_output_size=500)
except:
# Note: the IoU of two same rectangles is 0, which is calculated by rotate_gpu_nms
jitter = np.zeros(
[tmp_boxes_r.shape[0], tmp_boxes_r.shape[1] + 1])
jitter[:,
0] += np.random.rand(tmp_boxes_r.shape[0], ) / 1000
inx = rotate_gpu_nms(
np.array(tmp, np.float32) +
np.array(jitter, np.float32),
float(threshold[LABEL_NAME_MAP[sub_class]]), 0)
box_res_rotate_.extend(np.array(tmp_boxes_r)[inx])
score_res_rotate_.extend(np.array(tmp_score_r)[inx])
label_res_rotate_.extend(np.array(tmp_label_r)[inx])
result_dict = {
'boxes': np.array(box_res_rotate_),
'scores': np.array(score_res_rotate_),
'labels': np.array(label_res_rotate_),
'image_id': img_path
}
result_queue.put_nowait(result_dict)
def inference(det_net,
file_paths,
des_folder,
h_len,
w_len,
h_overlap,
w_overlap,
save_res=False):
if save_res:
assert cfgs.SHOW_SCORE_THRSHOLD >= 0.5, \
'please set score threshold (example: SHOW_SCORE_THRSHOLD = 0.5) in cfgs.py'
else:
assert cfgs.SHOW_SCORE_THRSHOLD < 0.005, \
'please set score threshold (example: SHOW_SCORE_THRSHOLD = 0.00) in cfgs.py'
# 1. preprocess img
img_plac = tf.placeholder(dtype=tf.uint8, shape=[None, None, 3])
img_batch = tf.cast(img_plac, tf.float32)
img_batch = img_batch - tf.constant(cfgs.PIXEL_MEAN)
img_batch = short_side_resize_for_inference_data(
img_tensor=img_batch,
target_shortside_len=cfgs.IMG_SHORT_SIDE_LEN,
is_resize=False)
det_boxes_h, det_scores_h, det_category_h, \
det_boxes_r, det_scores_r, det_category_r = det_net.build_whole_detection_network(input_img_batch=img_batch,
gtboxes_h_batch=None,
gtboxes_r_batch=None)
init_op = tf.group(tf.global_variables_initializer(),
tf.local_variables_initializer())
restorer, restore_ckpt = det_net.get_restorer()
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
with tf.Session(config=config) as sess:
sess.run(init_op)
if not restorer is None:
restorer.restore(sess, restore_ckpt)
print('restore model')
for count, img_path in enumerate(file_paths):
start = timer()
img = cv2.imread(img_path)
box_res = []
label_res = []
score_res = []
box_res_rotate = []
label_res_rotate = []
score_res_rotate = []
imgH = img.shape[0]
imgW = img.shape[1]
if imgH < h_len:
temp = np.zeros([h_len, imgW, 3], np.float32)
temp[0:imgH, :, :] = img
img = temp
imgH = h_len
if imgW < w_len:
temp = np.zeros([imgH, w_len, 3], np.float32)
temp[:, 0:imgW, :] = img
img = temp
imgW = w_len
for hh in range(0, imgH, h_len - h_overlap):
if imgH - hh - 1 < h_len:
hh_ = imgH - h_len
else:
hh_ = hh
for ww in range(0, imgW, w_len - w_overlap):
if imgW - ww - 1 < w_len:
ww_ = imgW - w_len
else:
ww_ = ww
src_img = img[hh_:(hh_ + h_len), ww_:(ww_ + w_len), :]
det_boxes_h_, det_scores_h_, det_category_h_, \
det_boxes_r_, det_scores_r_, det_category_r_ = \
sess.run(
[det_boxes_h, det_scores_h, det_category_h,
det_boxes_r, det_scores_r, det_category_r],
feed_dict={img_plac: src_img[:, :, ::-1]}
)
if len(det_boxes_h_) > 0:
for ii in range(len(det_boxes_h_)):
box = det_boxes_h_[ii]
box[0] = box[0] + ww_
box[1] = box[1] + hh_
box[2] = box[2] + ww_
box[3] = box[3] + hh_
box_res.append(box)
label_res.append(det_category_h_[ii])
score_res.append(det_scores_h_[ii])
if len(det_boxes_r_) > 0:
for ii in range(len(det_boxes_r_)):
box_rotate = det_boxes_r_[ii]
box_rotate[0] = box_rotate[0] + ww_
box_rotate[1] = box_rotate[1] + hh_
box_res_rotate.append(box_rotate)
label_res_rotate.append(det_category_r_[ii])
score_res_rotate.append(det_scores_r_[ii])
box_res = np.array(box_res)
label_res = np.array(label_res)
score_res = np.array(score_res)
box_res_rotate = np.array(box_res_rotate)
label_res_rotate = np.array(label_res_rotate)
score_res_rotate = np.array(score_res_rotate)
box_res_rotate_, label_res_rotate_, score_res_rotate_ = [], [], []
box_res_, label_res_, score_res_ = [], [], []
r_threshold = {
'roundabout': 0.1,
'tennis-court': 0.3,
'swimming-pool': 0.1,
'storage-tank': 0.2,
'soccer-ball-field': 0.3,
'small-vehicle': 0.2,
'ship': 0.05,
'plane': 0.3,
'large-vehicle': 0.1,
'helicopter': 0.2,
'harbor': 0.0001,
'ground-track-field': 0.3,
'bridge': 0.0001,
'basketball-court': 0.3,
'baseball-diamond': 0.3
}
h_threshold = {
'roundabout': 0.35,
'tennis-court': 0.35,
'swimming-pool': 0.4,
'storage-tank': 0.3,
'soccer-ball-field': 0.3,
'small-vehicle': 0.4,
'ship': 0.35,
'plane': 0.35,
'large-vehicle': 0.4,
'helicopter': 0.4,
'harbor': 0.3,
'ground-track-field': 0.4,
'bridge': 0.3,
'basketball-court': 0.4,
'baseball-diamond': 0.3
}
for sub_class in range(1, cfgs.CLASS_NUM + 1):
index = np.where(label_res_rotate == sub_class)[0]
if len(index) == 0:
continue
tmp_boxes_r = box_res_rotate[index]
tmp_label_r = label_res_rotate[index]
tmp_score_r = score_res_rotate[index]
tmp_boxes_r = np.array(tmp_boxes_r)
tmp = np.zeros(
[tmp_boxes_r.shape[0], tmp_boxes_r.shape[1] + 1])
tmp[:, 0:-1] = tmp_boxes_r
tmp[:, -1] = np.array(tmp_score_r)
try:
inx = nms_rotate.nms_rotate_cpu(
boxes=np.array(tmp_boxes_r),
scores=np.array(tmp_score_r),
iou_threshold=r_threshold[LABEl_NAME_MAP[sub_class]],
max_output_size=500)
except:
# Note: the IoU of two same rectangles is 0, which is calculated by rotate_gpu_nms
jitter = np.zeros(
[tmp_boxes_r.shape[0], tmp_boxes_r.shape[1] + 1])
jitter[:,
0] += np.random.rand(tmp_boxes_r.shape[0], ) / 1000
inx = rotate_gpu_nms(
np.array(tmp, np.float32) +
np.array(jitter, np.float32),
float(r_threshold[LABEl_NAME_MAP[sub_class]]), 0)
box_res_rotate_.extend(np.array(tmp_boxes_r)[inx])
score_res_rotate_.extend(np.array(tmp_score_r)[inx])
label_res_rotate_.extend(np.array(tmp_label_r)[inx])
for sub_class in range(1, cfgs.CLASS_NUM + 1):
index = np.where(label_res == sub_class)[0]
if len(index) == 0:
continue
tmp_boxes_h = box_res[index]
tmp_label_h = label_res[index]
tmp_score_h = score_res[index]
tmp_boxes_h = np.array(tmp_boxes_h)
tmp = np.zeros(
[tmp_boxes_h.shape[0], tmp_boxes_h.shape[1] + 1])
tmp[:, 0:-1] = tmp_boxes_h
tmp[:, -1] = np.array(tmp_score_h)
inx = nms.py_cpu_nms(
dets=np.array(tmp, np.float32),
thresh=h_threshold[LABEl_NAME_MAP[sub_class]],
max_output_size=500)
box_res_.extend(np.array(tmp_boxes_h)[inx])
score_res_.extend(np.array(tmp_score_h)[inx])
label_res_.extend(np.array(tmp_label_h)[inx])
time_elapsed = timer() - start
if save_res:
det_detections_h = draw_box_in_img.draw_box_cv(
np.array(img, np.float32) - np.array(cfgs.PIXEL_MEAN),
boxes=np.array(box_res_),
labels=np.array(label_res_),
scores=np.array(score_res_))
det_detections_r = draw_box_in_img.draw_rotate_box_cv(
np.array(img, np.float32) - np.array(cfgs.PIXEL_MEAN),
boxes=np.array(box_res_rotate_),
labels=np.array(label_res_rotate_),
scores=np.array(score_res_rotate_))
save_dir = os.path.join(des_folder, cfgs.VERSION)
tools.mkdir(save_dir)
cv2.imwrite(
save_dir + '/' + img_path.split('/')[-1].split('.')[0] +
'_h.jpg', det_detections_h)
cv2.imwrite(
save_dir + '/' + img_path.split('/')[-1].split('.')[0] +
'_r.jpg', det_detections_r)
view_bar(
'{} cost {}s'.format(
img_path.split('/')[-1].split('.')[0], time_elapsed),
count + 1, len(file_paths))
else:
# eval txt
CLASS_DOTA = NAME_LABEL_MAP.keys()
# Task1
write_handle_r = {}
write_handle_h_ = {}
txt_dir_r = os.path.join('txt_output', cfgs.VERSION + '_r')
txt_dir_h_minAreaRect = os.path.join(
'txt_output', cfgs.VERSION + '_h_minAreaRect')
tools.mkdir(txt_dir_r)
tools.mkdir(txt_dir_h_minAreaRect)
for sub_class in CLASS_DOTA:
if sub_class == 'back_ground':
continue
write_handle_r[sub_class] = open(
os.path.join(txt_dir_r, 'Task1_%s.txt' % sub_class),
'a+')
write_handle_h_[sub_class] = open(
os.path.join(txt_dir_h_minAreaRect,
'Task2_%s.txt' % sub_class), 'a+')
rboxes = coordinate_convert.forward_convert(box_res_rotate_,
with_label=False)
for i, rbox in enumerate(rboxes):
command = '%s %.3f %.1f %.1f %.1f %.1f %.1f %.1f %.1f %.1f\n' % (
img_path.split('/')[-1].split('.')[0],
score_res_rotate_[i],
rbox[0],
rbox[1],
rbox[2],
rbox[3],
rbox[4],
rbox[5],
rbox[6],
rbox[7],
)
command_ = '%s %.3f %.1f %.1f %.1f %.1f\n' % (
img_path.split('/')[-1].split('.')[0],
score_res_rotate_[i], min(rbox[::2]), min(
rbox[1::2]), max(rbox[::2]), max(rbox[1::2]))
write_handle_r[LABEl_NAME_MAP[label_res_rotate_[i]]].write(
command)
write_handle_h_[LABEl_NAME_MAP[
label_res_rotate_[i]]].write(command_)
for sub_class in CLASS_DOTA:
if sub_class == 'back_ground':
continue
write_handle_r[sub_class].close()
# Task2
write_handle_h = {}
txt_dir_h = os.path.join('txt_output', cfgs.VERSION + '_h')
tools.mkdir(txt_dir_h)
for sub_class in CLASS_DOTA:
if sub_class == 'back_ground':
continue
write_handle_h[sub_class] = open(
os.path.join(txt_dir_h, 'Task2_%s.txt' % sub_class),
'a+')
for i, hbox in enumerate(box_res_):
command = '%s %.3f %.1f %.1f %.1f %.1f\n' % (
img_path.split('/')[-1].split('.')[0], score_res_[i],
hbox[0], hbox[1], hbox[2], hbox[3])
write_handle_h[LABEl_NAME_MAP[label_res_[i]]].write(
command)
for sub_class in CLASS_DOTA:
if sub_class == 'back_ground':
continue
write_handle_h[sub_class].close()
view_bar(
'{} cost {}s'.format(
img_path.split('/')[-1].split('.')[0], time_elapsed),
count + 1, len(file_paths))
def eval_with_plac(det_net, args):
# 1. preprocess img
img_plac = tf.placeholder(dtype=tf.uint8, shape=[None, None,
3]) # is RGB. not BGR
img_batch = tf.cast(img_plac, tf.float32)
if cfgs.NET_NAME in ['resnet152_v1d', 'resnet101_v1d', 'resnet50_v1d']:
img_batch = (img_batch / 255 - tf.constant(
cfgs.PIXEL_MEAN_)) / tf.constant(cfgs.PIXEL_STD)
else:
img_batch = img_batch - tf.constant(cfgs.PIXEL_MEAN)
img_batch = tf.expand_dims(img_batch, axis=0)
detection_boxes, detection_scores, detection_category, detection_boxes_angle = det_net.build_whole_detection_network(
input_img_batch=img_batch,
gtboxes_batch_h=None,
gtboxes_batch_r=None,
gt_smooth_label=None)
init_op = tf.group(tf.global_variables_initializer(),
tf.local_variables_initializer())
restorer, restore_ckpt = det_net.get_restorer()
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
with tf.Session(config=config) as sess:
sess.run(init_op)
if not restorer is None:
restorer.restore(sess, restore_ckpt)
print('restore model')
all_boxes_r = []
img_short_side_len_list = cfgs.IMG_SHORT_SIDE_LEN if isinstance(
cfgs.IMG_SHORT_SIDE_LEN, list) else [cfgs.IMG_SHORT_SIDE_LEN]
img_short_side_len_list = [
img_short_side_len_list[0]
] if not args.multi_scale else img_short_side_len_list
imgs = os.listdir(args.img_dir)
pbar = tqdm(imgs)
for a_img_name in pbar:
a_img_name = a_img_name.split(args.image_ext)[0]
raw_img = cv2.imread(
os.path.join(args.img_dir, a_img_name + args.image_ext))
raw_h, raw_w = raw_img.shape[0], raw_img.shape[1]
box_res_rotate = []
label_res_rotate = []
score_res_rotate = []
for short_size in img_short_side_len_list:
max_len = cfgs.IMG_MAX_LENGTH
if raw_h < raw_w:
new_h, new_w = short_size, min(
int(short_size * float(raw_w) / raw_h), max_len)
else:
new_h, new_w = min(int(short_size * float(raw_h) / raw_w),
max_len), short_size
img_resize = cv2.resize(raw_img, (new_w, new_h))
resized_img, det_boxes_r_, det_scores_r_, det_category_r_ = \
sess.run(
[img_batch, detection_boxes_angle, detection_scores, detection_category],
feed_dict={img_plac: img_resize[:, :, ::-1]}
)
resized_h, resized_w = resized_img.shape[1], resized_img.shape[
2]
if len(det_boxes_r_) > 0:
det_boxes_r_ = forward_convert(det_boxes_r_, False)
det_boxes_r_[:, 0::2] *= (raw_w / resized_w)
det_boxes_r_[:, 1::2] *= (raw_h / resized_h)
for ii in range(len(det_boxes_r_)):
box_rotate = det_boxes_r_[ii]
box_res_rotate.append(box_rotate)
label_res_rotate.append(det_category_r_[ii])
score_res_rotate.append(det_scores_r_[ii])
box_res_rotate = np.array(box_res_rotate)
label_res_rotate = np.array(label_res_rotate)
score_res_rotate = np.array(score_res_rotate)
box_res_rotate_ = []
label_res_rotate_ = []
score_res_rotate_ = []
threshold = {'car': 0.2, 'plane': 0.3}
for sub_class in range(1, cfgs.CLASS_NUM + 1):
index = np.where(label_res_rotate == sub_class)[0]
if len(index) == 0:
continue
tmp_boxes_r = box_res_rotate[index]
tmp_label_r = label_res_rotate[index]
tmp_score_r = score_res_rotate[index]
tmp_boxes_r_ = backward_convert(tmp_boxes_r, False)
try:
inx = nms_rotate.nms_rotate_cpu(
boxes=np.array(tmp_boxes_r_),
scores=np.array(tmp_score_r),
iou_threshold=threshold[LABEL_NAME_MAP[sub_class]],
max_output_size=150)
except:
tmp_boxes_r_ = np.array(tmp_boxes_r_)
tmp = np.zeros(
[tmp_boxes_r_.shape[0], tmp_boxes_r_.shape[1] + 1])
tmp[:, 0:-1] = tmp_boxes_r_
tmp[:, -1] = np.array(tmp_score_r)
# Note: the IoU of two same rectangles is 0, which is calculated by rotate_gpu_nms
jitter = np.zeros(
[tmp_boxes_r_.shape[0], tmp_boxes_r_.shape[1] + 1])
jitter[:,
0] += np.random.rand(tmp_boxes_r_.shape[0], ) / 1000
inx = rotate_gpu_nms(
np.array(tmp, np.float32) +
np.array(jitter, np.float32),
float(threshold[LABEL_NAME_MAP[sub_class]]), 0)
box_res_rotate_.extend(np.array(tmp_boxes_r)[inx])
score_res_rotate_.extend(np.array(tmp_score_r)[inx])
label_res_rotate_.extend(np.array(tmp_label_r)[inx])
box_res_rotate_ = np.array(box_res_rotate_)
score_res_rotate_ = np.array(score_res_rotate_)
label_res_rotate_ = np.array(label_res_rotate_)
if args.draw_imgs:
detected_indices = score_res_rotate_ >= cfgs.VIS_SCORE
detected_scores = score_res_rotate_[detected_indices]
detected_boxes = box_res_rotate_[detected_indices]
detected_boxes = backward_convert(detected_boxes,
with_label=False)
detected_categories = label_res_rotate_[detected_indices]
det_detections_r = draw_box_in_img.draw_boxes_with_label_and_scores(
np.array(raw_img, np.float32),
boxes=detected_boxes,
labels=detected_categories,
scores=detected_scores,
method=1,
in_graph=False,
is_csl=True)
save_dir = os.path.join('test_ucas_aod', cfgs.VERSION,
'ucas_aod_img_vis')
tools.mkdir(save_dir)
cv2.imwrite(save_dir + '/{}.jpg'.format(a_img_name),
det_detections_r[:, :, ::-1])
if box_res_rotate_.shape[0] != 0:
box_res_rotate_ = backward_convert(box_res_rotate_, False)
x_c, y_c, w, h, theta = box_res_rotate_[:, 0], box_res_rotate_[:, 1], box_res_rotate_[:, 2], \
box_res_rotate_[:, 3], box_res_rotate_[:, 4]
boxes_r = np.transpose(np.stack([x_c, y_c, w, h, theta]))
dets_r = np.hstack((label_res_rotate_.reshape(-1, 1),
score_res_rotate_.reshape(-1, 1), boxes_r))
all_boxes_r.append(dets_r)
pbar.set_description("Eval image %s" % a_img_name)
# fw1 = open(cfgs.VERSION + '_detections_r.pkl', 'wb')
# pickle.dump(all_boxes_r, fw1)
return all_boxes_r
int_area = cv2.contourArea(order_pts)
inter = int_area * 1.0 / (
area_r1 + area_r2 - int_area + 1e-5) # 防止除以0
if inter >= iou_threshold:
suppressed[j] = 1
inx = np.array(keep, np.int64)
######################################################
except:
# Note: the IoU of two same rectangles is 0, which is calculated by rotate_gpu_nms
jitter = np.zeros([tmp_boxes_r.shape[0], tmp_boxes_r.shape[1] + 1])
jitter[:, 0] += np.random.rand(tmp_boxes_r.shape[0], ) / 1000
inx = rotate_gpu_nms(
np.array(tmp, np.float32) + np.array(jitter, np.float32),
float(r_threshold[LABEl_NAME_MAP[sub_class]]), 0)
######################################################
box_res_rotate_.extend(np.array(tmp_boxes_r)[inx])
score_res_rotate_.extend(np.array(tmp_score_r)[inx])
label_res_rotate_.extend(np.array(tmp_label_r)[inx])
################################################################################
#########################################
##clw modify: 测试使用
##import matplotlib.pyplot as plt
##plt.imshow(predictions[:, :, [2, 1, 0]])
##plt.axis("off")
##plt.show()
#########################################
def worker(gpu_id, images, det_net, result_queue):
os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_id)
# 1. preprocess img
img_plac = tf.placeholder(dtype=tf.uint8, shape=[None, None,
3]) # is RGB. not BGR
img_batch = tf.cast(img_plac, tf.float32)
if cfgs.NET_NAME in ['resnet152_v1d', 'resnet101_v1d', 'resnet50_v1d']:
img_batch = (img_batch / 255 - tf.constant(
cfgs.PIXEL_MEAN_)) / tf.constant(cfgs.PIXEL_STD)
else:
img_batch = img_batch - tf.constant(cfgs.PIXEL_MEAN)
img_batch = tf.expand_dims(img_batch, axis=0)
detection_boxes, detection_scores, detection_category, detection_boxes_angle = det_net.build_whole_detection_network(
input_img_batch=img_batch,
gtboxes_batch_h=None,
gtboxes_batch_r=None,
gt_smooth_label=None,
gpu_id=0)
init_op = tf.group(tf.global_variables_initializer(),
tf.local_variables_initializer())
restorer, restore_ckpt = det_net.get_restorer()
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
with tf.Session(config=config) as sess:
sess.run(init_op)
if not restorer is None:
restorer.restore(sess, restore_ckpt)
print('restore model %d ...' % gpu_id)
for a_img in images:
raw_img = cv2.imread(a_img)
raw_h, raw_w = raw_img.shape[0], raw_img.shape[1]
det_boxes_r_all, det_scores_r_all, det_category_r_all = [], [], []
img_short_side_len_list = cfgs.IMG_SHORT_SIDE_LEN if isinstance(
cfgs.IMG_SHORT_SIDE_LEN, list) else [cfgs.IMG_SHORT_SIDE_LEN]
img_short_side_len_list = [
img_short_side_len_list[0]
] if not args.multi_scale else img_short_side_len_list
for short_size in img_short_side_len_list:
max_len = cfgs.IMG_MAX_LENGTH
if raw_h < raw_w:
new_h, new_w = short_size, min(
int(short_size * float(raw_w) / raw_h), max_len)
else:
new_h, new_w = min(int(short_size * float(raw_h) / raw_w),
max_len), short_size
img_resize = cv2.resize(raw_img, (new_w, new_h))
resized_img, detected_boxes, detected_scores, detected_categories = \
sess.run(
[img_batch, detection_boxes_angle, detection_scores, detection_category],
feed_dict={img_plac: img_resize[:, :, ::-1]}
)
detected_indices = detected_scores >= cfgs.VIS_SCORE
detected_scores = detected_scores[detected_indices]
detected_boxes = detected_boxes[detected_indices]
detected_categories = detected_categories[detected_indices]
if detected_boxes.shape[0] == 0:
continue
resized_h, resized_w = resized_img.shape[1], resized_img.shape[
2]
detected_boxes = forward_convert(detected_boxes, False)
detected_boxes[:, 0::2] *= (raw_w / resized_w)
detected_boxes[:, 1::2] *= (raw_h / resized_h)
# detected_boxes = backward_convert(detected_boxes, False)
det_boxes_r_all.extend(detected_boxes)
det_scores_r_all.extend(detected_scores)
det_category_r_all.extend(detected_categories)
det_boxes_r_all = np.array(det_boxes_r_all)
det_scores_r_all = np.array(det_scores_r_all)
det_category_r_all = np.array(det_category_r_all)
box_res_rotate_ = []
label_res_rotate_ = []
score_res_rotate_ = []
if det_scores_r_all.shape[0] != 0:
for sub_class in range(1, cfgs.CLASS_NUM + 1):
index = np.where(det_category_r_all == sub_class)[0]
if len(index) == 0:
continue
tmp_boxes_r = det_boxes_r_all[index]
tmp_label_r = det_category_r_all[index]
tmp_score_r = det_scores_r_all[index]
tmp_boxes_r_ = backward_convert(tmp_boxes_r, False)
try:
inx = nms_rotate.nms_rotate_cpu(
boxes=np.array(tmp_boxes_r_),
scores=np.array(tmp_score_r),
iou_threshold=cfgs.NMS_IOU_THRESHOLD,
max_output_size=5000)
except:
tmp_boxes_r_ = np.array(tmp_boxes_r_)
tmp = np.zeros(
[tmp_boxes_r_.shape[0], tmp_boxes_r_.shape[1] + 1])
tmp[:, 0:-1] = tmp_boxes_r_
tmp[:, -1] = np.array(tmp_score_r)
# Note: the IoU of two same rectangles is 0, which is calculated by rotate_gpu_nms
jitter = np.zeros(
[tmp_boxes_r_.shape[0], tmp_boxes_r_.shape[1] + 1])
jitter[:, 0] += np.random.rand(
tmp_boxes_r_.shape[0], ) / 1000
inx = rotate_gpu_nms(
np.array(tmp, np.float32) +
np.array(jitter, np.float32),
float(cfgs.NMS_IOU_THRESHOLD), 0)
box_res_rotate_.extend(np.array(tmp_boxes_r)[inx])
score_res_rotate_.extend(np.array(tmp_score_r)[inx])
label_res_rotate_.extend(np.array(tmp_label_r)[inx])
box_res_rotate_ = np.array(box_res_rotate_)
score_res_rotate_ = np.array(score_res_rotate_)
label_res_rotate_ = np.array(label_res_rotate_)
result_dict = {
'scales': [1, 1],
'boxes': box_res_rotate_,
'scores': score_res_rotate_,
'labels': label_res_rotate_,
'image_id': a_img
}
result_queue.put_nowait(result_dict)