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_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_encode_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
def detect(det_net, inference_save_path, real_test_imgname_list):
# 1. preprocess img
img_plac = tf.placeholder(dtype=tf.uint8, shape=[None, None,
3]) # is RGB. not GBR
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)
img_batch = img_batch - tf.constant(cfgs.PIXEL_MEAN)
img_batch = tf.expand_dims(img_batch, axis=0) # [1, None, None, 3]
detection_boxes, detection_scores, detection_category = det_net.build_whole_detection_network(
input_img_batch=img_batch, gtboxes_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 i, a_img_name in enumerate(real_test_imgname_list):
raw_img = cv2.imread(a_img_name)
start = time.time()
resized_img, detected_boxes, detected_scores, detected_categories = \
sess.run(
[img_batch, detection_boxes, detection_scores, detection_category],
feed_dict={img_plac: raw_img}
)
end = time.time()
# print("{} cost time : {} ".format(img_name, (end - start)))
raw_h, raw_w = raw_img.shape[0], raw_img.shape[1]
xmin, ymin, xmax, ymax = detected_boxes[:, 0], detected_boxes[:, 1], \
detected_boxes[:, 2], detected_boxes[:, 3]
resized_h, resized_w = resized_img.shape[1], resized_img.shape[2]
xmin = xmin * raw_w / resized_w
xmax = xmax * raw_w / resized_w
ymin = ymin * raw_h / resized_h
ymax = ymax * raw_h / resized_h
detected_boxes = np.transpose(np.stack([xmin, ymin, xmax, ymax]))
show_indices = detected_scores >= cfgs.SHOW_SCORE_THRSHOLD
show_scores = detected_scores[show_indices]
show_boxes = detected_boxes[show_indices]
show_categories = detected_categories[show_indices]
final_detections = draw_box_in_img.draw_boxes_with_label_and_scores(
raw_img - np.array(cfgs.PIXEL_MEAN),
boxes=show_boxes,
labels=show_categories,
scores=show_scores)
nake_name = a_img_name.split('/')[-1]
# print (inference_save_path + '/' + nake_name)
cv2.imwrite(inference_save_path + '/' + nake_name,
final_detections[:, :, ::-1])
tools.view_bar(
'{} image cost {}s'.format(a_img_name, (end - start)), i + 1,
len(real_test_imgname_list))
def eval_with_plac(det_net, real_test_imgname_list, img_root, draw_imgs=False):
# 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)
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 ['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 = (img_batch - tf.constant(cfgs.PIXEL_MEAN)) / (tf.constant(cfgs.PIXEL_STD)*255)
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=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 i, a_img_name in enumerate(real_test_imgname_list):
raw_img = cv2.imread(os.path.join(img_root, a_img_name))
raw_h, raw_w = raw_img.shape[0], raw_img.shape[1]
start = time.time()
resized_img, detected_boxes, detected_scores, detected_categories = \
sess.run(
[img_batch, detection_boxes, detection_scores, detection_category],
feed_dict={img_plac: raw_img[:, :, ::-1]} # cv is BGR. But need RGB
)
end = time.time()
# print("{} cost time : {} ".format(img_name, (end - start)))
if draw_imgs:
show_indices = detected_scores >= cfgs.SHOW_SCORE_THRSHOLD
show_scores = detected_scores[show_indices]
show_boxes = detected_boxes[show_indices]
show_categories = detected_categories[show_indices]
draw_img = np.squeeze(resized_img, 0)
if cfgs.NET_NAME in ['resnet101_v1d', 'resnet50_v1d']:
draw_img = (draw_img * np.array(cfgs.PIXEL_STD) + np.array(cfgs.PIXEL_MEAN_)) * 255
else:
draw_img = draw_img + np.array(cfgs.PIXEL_MEAN)
# draw_img = draw_img * (np.array(cfgs.PIXEL_STD)*255) + np.array(cfgs.PIXEL_MEAN)
final_detections = draw_box_in_img.draw_boxes_with_label_and_scores(draw_img,
boxes=show_boxes,
labels=show_categories,
scores=show_scores,
in_graph=False)
if not os.path.exists(cfgs.TEST_SAVE_PATH):
os.makedirs(cfgs.TEST_SAVE_PATH)
cv2.imwrite(cfgs.TEST_SAVE_PATH + '/' + a_img_name,
final_detections[:, :, ::-1])
xmin, ymin, xmax, ymax = detected_boxes[:, 0], detected_boxes[:, 1], \
detected_boxes[:, 2], detected_boxes[:, 3]
resized_h, resized_w = resized_img.shape[1], resized_img.shape[2]
xmin = xmin * raw_w / resized_w
xmax = xmax * raw_w / resized_w
ymin = ymin * raw_h / resized_h
ymax = ymax * raw_h / resized_h
boxes = np.transpose(np.stack([xmin, ymin, xmax, ymax]))
dets = np.hstack((detected_categories.reshape(-1, 1),
detected_scores.reshape(-1, 1),
boxes))
# all_boxes.append(dets)
# eval txt
CLASS_VOC = NAME_LABEL_MAP.keys()
write_handle = {}
txt_dir = os.path.join('voc2012_eval', cfgs.VERSION, 'results', 'VOC2012', 'Main')
tools.mkdir(txt_dir)
for sub_class in CLASS_VOC:
if sub_class == 'back_ground':
continue
write_handle[sub_class] = open(os.path.join(txt_dir, 'comp3_det_test_%s.txt' % sub_class), 'a+')
for det in dets:
command = '%s %.6f %.6f %.6f %.6f %.6f\n' % (a_img_name.split('/')[-1].split('.')[0],
det[1],
det[2], det[3], det[4], det[5])
write_handle[LABEl_NAME_MAP[det[0]]].write(command)
for sub_class in CLASS_VOC:
if sub_class == 'back_ground':
continue
write_handle[sub_class].close()
tools.view_bar('%s image cost %.3fs' % (a_img_name, (end - start)), i + 1, len(real_test_imgname_list))
def test_icdar2015(det_net, real_test_img_list, gpu_ids, show_box, txt_name):
save_path = os.path.join('./test_icdar2015', cfgs.VERSION)
tools.mkdir(save_path)
nr_records = len(real_test_img_list)
pbar = tqdm(total=nr_records)
gpu_num = len(gpu_ids.strip().split(','))
nr_image = math.ceil(nr_records / gpu_num)
result_queue = Queue(500)
procs = []
for i, gpu_id in enumerate(gpu_ids.strip().split(',')):
start = i * nr_image
end = min(start + nr_image, nr_records)
split_records = real_test_img_list[start:end]
proc = Process(target=worker,
args=(int(gpu_id), split_records, det_net,
result_queue))
print('process:%d, start:%d, end:%d' % (i, start, end))
proc.start()
procs.append(proc)
for i in range(nr_records):
res = result_queue.get()
if res['boxes'].shape[0] == 0:
fw_txt_dt = open(
os.path.join(
save_path, 'res_{}.txt'.format(
res['image_id'].split('/')[-1].split('.')[0])), 'w')
fw_txt_dt.close()
pbar.update(1)
fw = open(txt_name, 'a+')
fw.write('{}\n'.format(res['image_id'].split('/')[-1]))
fw.close()
continue
x1, y1, x2, y2, x3, y3, x4, y4 = res['boxes'][:, 0], res['boxes'][:, 1], res['boxes'][:, 2], res['boxes'][:, 3],\
res['boxes'][:, 4], res['boxes'][:, 5], res['boxes'][:, 6], res['boxes'][:, 7]
x1, y1 = x1 * res['scales'][0], y1 * res['scales'][1]
x2, y2 = x2 * res['scales'][0], y2 * res['scales'][1]
x3, y3 = x3 * res['scales'][0], y3 * res['scales'][1]
x4, y4 = x4 * res['scales'][0], y4 * res['scales'][1]
boxes = np.transpose(np.stack([x1, y1, x2, y2, x3, y3, x4, y4]))
if show_box:
boxes = backward_convert(boxes, False)
nake_name = res['image_id'].split('/')[-1]
draw_path = os.path.join(save_path, nake_name)
draw_img = np.array(cv2.imread(res['image_id']), np.float32)
final_detections = draw_box_in_img.draw_boxes_with_label_and_scores(
draw_img,
boxes=boxes,
labels=res['labels'],
scores=res['scores'],
method=1,
in_graph=False)
cv2.imwrite(draw_path, final_detections)
else:
fw_txt_dt = open(
os.path.join(
save_path, 'res_{}.txt'.format(
res['image_id'].split('/')[-1].split('.')[0])), 'w')
for box in boxes:
line = '%d,%d,%d,%d,%d,%d,%d,%d\n' % (box[0], box[1], box[2],
box[3], box[4], box[5],
box[6], box[7])
fw_txt_dt.write(line)
fw_txt_dt.close()
fw = open(txt_name, 'a+')
fw.write('{}\n'.format(res['image_id'].split('/')[-1]))
fw.close()
pbar.set_description("Test image %s" % res['image_id'].split('/')[-1])
pbar.update(1)
for p in procs:
p.join()
def test_dota(det_net, real_test_img_list, args, txt_name):
save_path = os.path.join('./test_dota', cfgs.VERSION)
nr_records = len(real_test_img_list)
pbar = tqdm(total=nr_records)
gpu_num = len(args.gpus.strip().split(','))
nr_image = math.ceil(nr_records / gpu_num)
result_queue = Queue(500)
procs = []
for i, gpu_id in enumerate(args.gpus.strip().split(',')):
start = i * nr_image
end = min(start + nr_image, nr_records)
split_records = real_test_img_list[start:end]
proc = Process(target=worker, args=(int(gpu_id), split_records, det_net, args, result_queue))
print('process:%d, start:%d, end:%d' % (i, start, end))
proc.start()
procs.append(proc)
log_dir = './csl_log/{}'.format(cfgs.VERSION)
tools.mkdir(log_dir)
fw_tsv = open(os.path.join(log_dir, 'csl_meta.tsv'), 'w')
# fw_tsv.write("Label\n")
final_logits = []
for i in range(nr_records):
res = result_queue.get()
if args.show_box:
nake_name = res['image_id'].split('/')[-1]
tools.mkdir(os.path.join(save_path, 'dota_img_vis'))
draw_path = os.path.join(save_path, 'dota_img_vis', nake_name)
draw_img = np.array(cv2.imread(res['image_id']), np.float32)
detected_boxes = backward_convert(res['boxes'], with_label=False)
detected_indices = res['scores'] >= cfgs.VIS_SCORE
detected_scores = res['scores'][detected_indices]
detected_boxes = detected_boxes[detected_indices]
detected_categories = res['labels'][detected_indices]
final_detections = draw_box_in_img.draw_boxes_with_label_and_scores(draw_img,
boxes=detected_boxes,
labels=detected_categories,
scores=detected_scores,
method=1,
head=np.ones_like(detected_scores) * -1,
is_csl=True,
in_graph=False)
cv2.imwrite(draw_path, final_detections)
else:
detected_indices = res['scores'] >= cfgs.VIS_SCORE
res['scores'] = res['scores'][detected_indices]
res['boxes'] = res['boxes'][detected_indices]
res['labels'] = res['labels'][detected_indices]
rboxes = backward_convert(res['boxes'], with_label=False)
rboxes = coordinate_present_convert(rboxes, -1, False)
rlogits = res['logits'][detected_indices]
for ii, rb in enumerate(rboxes):
fw_tsv.write("%d\n" % (int(rb[-1])))
final_logits.append(rlogits[ii])
fw = open(txt_name, 'a+')
fw.write('{}\n'.format(res['image_id'].split('/')[-1]))
fw.close()
pbar.set_description("Test image %s" % res['image_id'].split('/')[-1])
pbar.update(1)
for p in procs:
p.join()
fw_tsv.close()
final_logits = np.array(final_logits)
np.save(os.path.join(log_dir, "final_logits.npy"), final_logits)
def detect(det_net, inference_save_path, real_test_imgname_list):
# 1. preprocess img
img_plac = tf.placeholder(dtype=tf.uint8, shape=[None, None,
3]) # is RGB. not GBR
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) # [1, None, None, 3]
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')
for i, a_img_name in enumerate(real_test_imgname_list):
raw_img = cv2.imread(a_img_name)
start = time.time()
resized_img, detected_boxes, detected_scores, detected_categories = \
sess.run(
[img_batch, detection_boxes, detection_scores, detection_category],
feed_dict={img_plac: raw_img[:, :, ::-1]} # cv is BGR. But need RGB
)
end = time.time()
# print("{} cost time : {} ".format(img_name, (end - start)))
show_indices = detected_scores >= cfgs.VIS_SCORE
show_scores = detected_scores[show_indices]
show_boxes = detected_boxes[show_indices]
show_categories = detected_categories[show_indices]
draw_img = np.squeeze(resized_img, 0)
if cfgs.NET_NAME in [
'resnet152_v1d', 'resnet101_v1d', 'resnet50_v1d'
]:
draw_img = (draw_img * np.array(cfgs.PIXEL_STD) +
np.array(cfgs.PIXEL_MEAN_)) * 255
else:
draw_img = draw_img + np.array(cfgs.PIXEL_MEAN)
final_detections = draw_box_in_img.draw_boxes_with_label_and_scores(
draw_img,
boxes=show_boxes,
labels=show_categories,
scores=show_scores,
method=1,
in_graph=False)
nake_name = a_img_name.split('/')[-1]
# print (inference_save_path + '/' + nake_name)
cv2.imwrite(inference_save_path + '/' + nake_name,
final_detections[:, :, ::-1])
tools.view_bar(
'{} image cost {}s'.format(nake_name, (end - start)), i + 1,
len(real_test_imgname_list))
def test_dota(det_net, real_test_img_list, args, txt_name):
save_path = os.path.join('./test_dota', cfgs.VERSION)
nr_records = len(real_test_img_list)
pbar = tqdm(total=nr_records)
gpu_num = len(args.gpus.strip().split(','))
nr_image = math.ceil(nr_records / gpu_num)
result_queue = Queue(500)
procs = []
for i, gpu_id in enumerate(args.gpus.strip().split(',')):
start = i * nr_image
end = min(start + nr_image, nr_records)
split_records = real_test_img_list[start:end]
proc = Process(target=worker,
args=(int(gpu_id), split_records, det_net, args,
result_queue))
print('process:%d, start:%d, end:%d' % (i, start, end))
proc.start()
procs.append(proc)
for i in range(nr_records):
res = result_queue.get()
if args.show_box:
nake_name = res['image_id'].split('/')[-1]
tools.mkdir(os.path.join(save_path, 'dota_img_vis'))
draw_path = os.path.join(save_path, 'dota_img_vis', nake_name)
draw_img = np.array(cv2.imread(res['image_id']), np.float32)
detected_indices = res['scores'] >= cfgs.VIS_SCORE
detected_scores = res['scores'][detected_indices]
detected_boxes = res['boxes'][detected_indices]
detected_categories = res['labels'][detected_indices]
final_detections = draw_box_in_img.draw_boxes_with_label_and_scores(
draw_img,
boxes=detected_boxes,
labels=detected_categories,
scores=detected_scores,
method=1,
in_graph=False)
cv2.imwrite(draw_path, final_detections)
else:
CLASS_DOTA = NAME_LABEL_MAP.keys()
write_handle = {}
tools.mkdir(os.path.join(save_path, 'dota_res'))
for sub_class in CLASS_DOTA:
if sub_class == 'back_ground':
continue
write_handle[sub_class] = open(
os.path.join(save_path, 'dota_res',
'Task1_%s.txt' % sub_class), 'a+')
rboxes = forward_convert(res['boxes'], with_label=False)
for i, rbox in enumerate(rboxes):
command = '%s %.3f %.1f %.1f %.1f %.1f %.1f %.1f %.1f %.1f\n' % (
res['image_id'].split('/')[-1].split('.')[0],
res['scores'][i],
rbox[0],
rbox[1],
rbox[2],
rbox[3],
rbox[4],
rbox[5],
rbox[6],
rbox[7],
)
write_handle[LABEL_NAME_MAP[res['labels'][i]]].write(command)
for sub_class in CLASS_DOTA:
if sub_class == 'back_ground':
continue
write_handle[sub_class].close()
fw = open(txt_name, 'a+')
fw.write('{}\n'.format(res['image_id'].split('/')[-1]))
fw.close()
pbar.set_description("Test image %s" % res['image_id'].split('/')[-1])
pbar.update(1)
for p in procs:
p.join()
def detect(det_net, src_dir, res_dir, draw_imgs):
# 1. preprocess img
img_plac = tf.placeholder(dtype=tf.uint8, shape=[None, None,
3]) # is RGB. not GBR
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)
img_batch = img_batch - tf.constant(cfgs.PIXEL_MEAN)
img_batch = tf.expand_dims(img_batch, axis=0) # [1, None, None, 3]
result_dict = det_net.build_whole_detection_network(
input_img_batch=img_batch, gtboxes_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')
sub_folders = os.listdir(src_dir)
for sub_folder in sub_folders:
folder_dir = os.path.join(src_dir, sub_folder)
real_test_imgname_list = [
os.path.join(folder_dir, img_name)
for img_name in os.listdir(folder_dir)
]
tools.mkdir(os.path.join(res_dir, sub_folder))
for i, a_img_name in enumerate(real_test_imgname_list):
raw_img = cv2.imread(a_img_name)
raw_h, raw_w = raw_img.shape[0], raw_img.shape[1]
start = time.time()
resized_img, result_dict_ = \
sess.run(
[img_batch, result_dict],
feed_dict={img_plac: raw_img[:, :, ::-1]} # cv is BGR. But need RGB
)
end = time.time()
detected_boxes, detected_scores, detected_categories = merge_result(
result_dict_)
nake_name = a_img_name.split('/')[-1]
xmin, ymin, xmax, ymax = detected_boxes[:, 0], detected_boxes[:, 1], \
detected_boxes[:, 2], detected_boxes[:, 3]
resized_h, resized_w = resized_img.shape[1], resized_img.shape[
2]
xmin = xmin * raw_w / resized_w
xmax = xmax * raw_w / resized_w
ymin = ymin * raw_h / resized_h
ymax = ymax * raw_h / resized_h
boxes = np.transpose(np.stack([xmin, ymin, xmax, ymax]))
dets = np.hstack((detected_categories.reshape(-1, 1),
detected_scores.reshape(-1, 1), boxes))
show_indices = detected_scores >= cfgs.SHOW_SCORE_THRSHOLD
show_scores = detected_scores[show_indices]
show_boxes = boxes[show_indices]
show_categories = detected_categories[show_indices]
f = open(
os.path.join(res_dir, sub_folder) + '/' +
nake_name.split('.')[0] + '.txt', 'w')
f.write('{}\n'.format(nake_name.split('.')[0]))
# f.write('{}\n'.format(dets.shape[0]))
# for inx in range(dets.shape[0]):
#
# f.write('%d %d %d %d %.3f\n' % (int(dets[inx][2]),
# int(dets[inx][3]),
# int(dets[inx][4]) - int(dets[inx][2]),
# int(dets[inx][5]) - int(dets[inx][3]),
# dets[inx][1]))
f.write('{}\n'.format(show_boxes.shape[0]))
for inx in range(show_boxes.shape[0]):
f.write('%d %d %d %d %.3f\n' %
(int(show_boxes[inx][0]), int(show_boxes[inx][1]),
int(show_boxes[inx][2]) - int(show_boxes[inx][0]),
int(show_boxes[inx][3]) - int(show_boxes[inx][1]),
show_scores[inx]))
if draw_imgs:
final_detections = draw_box_in_img.draw_boxes_with_label_and_scores(
raw_img - np.array(cfgs.PIXEL_MEAN),
boxes=show_boxes,
labels=show_categories,
scores=show_scores)
tools.mkdir(cfgs.TEST_SAVE_PATH)
cv2.imwrite(cfgs.TEST_SAVE_PATH + '/' + nake_name,
final_detections)
tools.view_bar(
'{} image cost {}s'.format(a_img_name, (end - start)),
i + 1, len(real_test_imgname_list))
def eval_with_plac(det_net, real_test_imgname_list, img_root, draw_imgs=False):
# 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)
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)
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=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 = []
for i, a_img_name in enumerate(real_test_imgname_list):
raw_img = cv2.imread(os.path.join(img_root, a_img_name))
raw_h, raw_w = raw_img.shape[0], raw_img.shape[1]
start = time.time()
resized_img, detected_boxes, detected_scores, detected_categories = \
sess.run(
[img_batch, detection_boxes, detection_scores, detection_category],
feed_dict={img_plac: raw_img[:, :, ::-1]} # cv is BGR. But need RGB
)
end = time.time()
# print("{} cost time : {} ".format(img_name, (end - start)))
if draw_imgs:
show_indices = detected_scores >= cfgs.SHOW_SCORE_THRSHOLD
show_scores = detected_scores[show_indices]
show_boxes = detected_boxes[show_indices]
show_categories = detected_categories[show_indices]
final_detections = draw_box_in_img.draw_boxes_with_label_and_scores(np.squeeze(resized_img, 0),
boxes=show_boxes,
labels=show_categories,
scores=show_scores)
if not os.path.exists(cfgs.TEST_SAVE_PATH):
os.makedirs(cfgs.TEST_SAVE_PATH)
cv2.imwrite(cfgs.TEST_SAVE_PATH + '/' + a_img_name,
final_detections[:, :, ::-1])
xmin, ymin, xmax, ymax = detected_boxes[:, 0], detected_boxes[:, 1], \
detected_boxes[:, 2], detected_boxes[:, 3]
resized_h, resized_w = resized_img.shape[1], resized_img.shape[2]
xmin = xmin * raw_w / resized_w
xmax = xmax * raw_w / resized_w
ymin = ymin * raw_h / resized_h
ymax = ymax * raw_h / resized_h
boxes = np.transpose(np.stack([xmin, ymin, xmax, ymax]))
dets = np.hstack((detected_categories.reshape(-1, 1),
detected_scores.reshape(-1, 1),
boxes))
all_boxes.append(dets)
tools.view_bar('{} image cost {}s'.format(a_img_name, (end - start)), i + 1, len(real_test_imgname_list))
save_dir = os.path.join(cfgs.EVALUATE_DIR, cfgs.VERSION)
if not os.path.exists(save_dir):
os.makedirs(save_dir)
fw1 = open(os.path.join(save_dir, 'detections.pkl'), 'wb')
pickle.dump(all_boxes, fw1)
def eval_coco(det_net, real_test_img_list, draw_imgs=False):
# 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)
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 ['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 = (img_batch - tf.constant(cfgs.PIXEL_MEAN)) / (tf.constant(cfgs.PIXEL_STD)*255)
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=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')
save_path = os.path.join('./eval_coco', cfgs.VERSION)
tools.mkdir(save_path)
fw_json_dt = open(os.path.join(save_path, 'coco_minival.json'), 'w')
coco_det = []
for i, a_img in enumerate(real_test_img_list):
record = json.loads(a_img)
img_path = os.path.join('/data/COCO/val2017',
record['fpath'].split('_')[-1])
raw_img = cv2.imread(img_path)
# raw_img = cv2.imread(record['fpath'])
raw_h, raw_w = raw_img.shape[0], raw_img.shape[1]
start = time.time()
resized_img, detected_boxes, detected_scores, detected_categories = \
sess.run(
[img_batch, detection_boxes, detection_scores, detection_category],
feed_dict={img_plac: raw_img[:, :, ::-1]} # cv is BGR. But need RGB
)
end = time.time()
eval_indices = detected_scores >= 0.01
detected_scores = detected_scores[eval_indices]
detected_boxes = detected_boxes[eval_indices]
detected_categories = detected_categories[eval_indices]
# print("{} cost time : {} ".format(img_name, (end - start)))
if draw_imgs:
show_indices = detected_scores >= cfgs.SHOW_SCORE_THRSHOLD
show_scores = detected_scores[show_indices]
show_boxes = detected_boxes[show_indices]
show_categories = detected_categories[show_indices]
draw_img = np.squeeze(resized_img, 0)
# draw_img = draw_img + np.array(cfgs.PIXEL_MEAN)
# draw_img = draw_img * (np.array(cfgs.PIXEL_STD)*255) + np.array(cfgs.PIXEL_MEAN)
final_detections = draw_box_in_img.draw_boxes_with_label_and_scores(
draw_img,
boxes=show_boxes,
labels=show_categories,
scores=show_scores,
in_graph=False)
if not os.path.exists(cfgs.TEST_SAVE_PATH):
os.makedirs(cfgs.TEST_SAVE_PATH)
cv2.imwrite(cfgs.TEST_SAVE_PATH + '/' + record['ID'],
final_detections[:, :, ::-1])
xmin, ymin, xmax, ymax = detected_boxes[:, 0], detected_boxes[:, 1], \
detected_boxes[:, 2], detected_boxes[:, 3]
resized_h, resized_w = resized_img.shape[1], resized_img.shape[2]
xmin = xmin * raw_w / resized_w
xmax = xmax * raw_w / resized_w
ymin = ymin * raw_h / resized_h
ymax = ymax * raw_h / resized_h
boxes = np.transpose(
np.stack([xmin, ymin, xmax - xmin, ymax - ymin]))
# cost much time
for j, box in enumerate(boxes):
coco_det.append({
'bbox': [
float(box[0]),
float(box[1]),
float(box[2]),
float(box[3])
],
'score':
float(detected_scores[j]),
'image_id':
int(record['ID'].split('.jpg')[0].split('_000000')[-1]),
'category_id':
int(classes_originID[LABEl_NAME_MAP[
detected_categories[j]]])
})
tools.view_bar(
'%s image cost %.3fs' % (record['ID'], (end - start)), i + 1,
len(real_test_img_list))
json.dump(coco_det, fw_json_dt)
fw_json_dt.close()
return os.path.join(save_path, 'coco_minival.json')
img_name = xml_name.replace('.xml', '.jpg')
img = cv2.imread(os.path.join(root_path, 'train', img_name))
img_height, img_width, gtbox_label = read_xml_gtbox_and_label(
os.path.join(root_path, 'Annotations', xml_name))
gtbox_and_label_list = np.array(gtbox_label, dtype=np.int32)
if gtbox_and_label_list.shape[0] == 0:
return next_img(step + 1)
else:
return img_name, img[:, :, ::-1], gtbox_and_label_list
if __name__ == '__main__':
imgid, img, gtbox = next_img(3234)
print("::")
from libs.box_utils.draw_box_in_img import draw_boxes_with_label_and_scores
img = draw_boxes_with_label_and_scores(
img_array=img,
boxes=gtbox[:, :-1],
labels=gtbox[:, -1],
scores=np.ones(shape=(len(gtbox), )))
print("_----")
cv2.imshow("test", img)
cv2.waitKey(0)
def eval_with_plac(det_net, real_test_imgname_list, img_root, draw_imgs=False):
# 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)
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=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 = []
pbar = tqdm(real_test_imgname_list)
for a_img_name in pbar:
raw_img = cv2.imread(os.path.join(img_root, a_img_name))
raw_h, raw_w = raw_img.shape[0], raw_img.shape[1]
resized_img, detected_boxes, detected_scores, detected_categories = \
sess.run(
[img_batch, detection_boxes, detection_scores, detection_category],
feed_dict={img_plac: raw_img[:, :, ::-1]} # cv is BGR. But need RGB
)
if draw_imgs:
show_indices = detected_scores >= cfgs.VIS_SCORE
show_scores = detected_scores[show_indices]
show_boxes = detected_boxes[show_indices]
show_categories = detected_categories[show_indices]
draw_img = np.squeeze(resized_img, 0)
if cfgs.NET_NAME in ['resnet152_v1d', 'resnet101_v1d', 'resnet50_v1d']:
draw_img = (draw_img * np.array(cfgs.PIXEL_STD) + np.array(cfgs.PIXEL_MEAN_)) * 255
else:
draw_img = draw_img + np.array(cfgs.PIXEL_MEAN)
final_detections = draw_box_in_img.draw_boxes_with_label_and_scores(draw_img,
boxes=show_boxes,
labels=show_categories,
scores=show_scores,
in_graph=False)
if not os.path.exists(cfgs.TEST_SAVE_PATH):
os.makedirs(cfgs.TEST_SAVE_PATH)
cv2.imwrite(cfgs.TEST_SAVE_PATH + '/' + a_img_name + '.jpg',
final_detections[:, :, ::-1])
xmin, ymin, xmax, ymax = detected_boxes[:, 0], detected_boxes[:, 1], \
detected_boxes[:, 2], detected_boxes[:, 3]
resized_h, resized_w = resized_img.shape[1], resized_img.shape[2]
xmin = xmin * raw_w / resized_w
xmax = xmax * raw_w / resized_w
ymin = ymin * raw_h / resized_h
ymax = ymax * raw_h / resized_h
boxes = np.transpose(np.stack([xmin, ymin, xmax, ymax]))
dets = np.hstack((detected_categories.reshape(-1, 1),
detected_scores.reshape(-1, 1),
boxes))
all_boxes.append(dets)
pbar.set_description("Eval image %s" % a_img_name)
# save_dir = os.path.join(cfgs.EVALUATE_DIR, cfgs.VERSION)
# if not os.path.exists(save_dir):
# os.makedirs(save_dir)
# fw1 = open(os.path.join(save_dir, 'detections.pkl'), 'w')
# pickle.dump(all_boxes, fw1)
return all_boxes
def exucute_detect(self, image_path, save_path):
"""
execute object detect
:param detect_net:
:param image_path:
:return:
"""
input_image = tf.placeholder(dtype=tf.uint8,
shape=(None, None, 3),
name='inputs_images')
resize_img = self.image_process(input_image)
# expend dimension
image_batch = tf.expand_dims(input=resize_img,
axis=0) # (1, None, None, 3)
self.detect_net.images_batch = image_batch
# img_shape = tf.shape(inputs_img)
# load detect network
detection_boxes, detection_scores, detection_category = self.detect_net.inference(
)
# restore pretrain weight
restorer, restore_ckpt = self.detect_net.get_restorer()
# config gpu to growth train
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
init_op = tf.group(tf.global_variables_initializer(),
tf.local_variables_initializer())
with tf.Session(config=config) as sess:
sess.run(init_op)
if restorer is not None:
restorer.restore(sess, save_path=restore_ckpt)
print('*' * 80 +
'\nSuccessful restore model from {0}\n'.format(
restore_ckpt) + '*' * 80)
# construct image path list
format_list = ('.jpg', '.png', '.jpeg', '.tif', '.tiff')
if os.path.isfile(image_path):
image_name_list = [image_path]
else:
image_name_list = [
img_name for img_name in os.listdir(image_path)
if img_name.endswith(format_list)
and os.path.isfile(os.path.join(image_path, img_name))
]
assert len(image_name_list) != 0
print(
"test_dir has no imgs there. Note that, we only support img format of {0}"
.format(format_list))
#+++++++++++++++++++++++++++++++++++++start detect+++++++++++++++++++++++++++++++++++++++++++++++++++++=++
makedir(save_path)
fw = open(os.path.join(save_path, 'detect_bbox.txt'), 'w')
for index, img_name in enumerate(image_name_list):
detect_dict = {}
bgr_img = cv.imread(os.path.join(image_path, img_name))
rgb_img = cv.cvtColor(
bgr_img, cv.COLOR_BGR2RGB
) # convert channel from BGR to RGB (cv is BGR)
start_time = time.perf_counter()
# image resize and white process
# construct feed_dict
feed_dict = {input_image: rgb_img}
resized_img, detected_boxes, detected_scores, detected_categories = \
sess.run([resize_img, detection_boxes, detection_scores, detection_category],
feed_dict=feed_dict)
end_time = time.perf_counter()
# select object according to threshold
object_indices = detected_scores >= cfgs.SHOW_SCORE_THRSHOLD
object_scores = detected_scores[object_indices]
object_boxes = detected_boxes[object_indices]
object_categories = detected_categories[object_indices]
final_detections_img = draw_box_in_img.draw_boxes_with_label_and_scores(
resized_img,
boxes=object_boxes,
labels=object_categories,
scores=object_scores)
final_detections_img = cv.cvtColor(final_detections_img,
cv.COLOR_RGB2BGR)
cv.imwrite(os.path.join(save_path, img_name),
final_detections_img)
# resize boxes and image according to raw input image
raw_h, raw_w = rgb_img.shape[0], rgb_img.shape[1]
resized_h, resized_w = resized_img.shape[1], resized_img.shape[
2]
x_min, y_min, x_max, y_max = object_boxes[:, 0], object_boxes[:, 1], object_boxes[:, 2], \
object_boxes[:, 3]
x_min = x_min * raw_w / resized_w
y_min = y_min * raw_h / resized_h
x_max = x_max * raw_w / resized_w
y_max = y_max * raw_h / resized_h
object_boxes = np.stack([x_min, y_min, x_max, y_max], axis=1)
# final_detections= cv.resize(final_detections[:, :, ::-1], (raw_w, raw_h))
# recover to raw size
detect_dict['score'] = object_scores
detect_dict['boxes'] = object_boxes
detect_dict['categories'] = object_categories
# convert from RGB to BG
fw.write(f'\n{img_name}')
for score, boxes, categories in zip(object_scores,
object_boxes,
object_categories):
fw.write('\n\tscore:' + str(score))
fw.write('\tbboxes:' + str(boxes))
fw.write('\tcategories:' + str(categories))
view_bar(
'{} image cost {} second'.format(img_name,
(end_time - start_time)),
index + 1, len(image_name_list))
fw.close()
def eval_with_plac(img_dir, det_net, num_imgs, image_ext, draw_imgs=False):
# 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)
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_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_encode_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 = []
imgs = os.listdir(img_dir)
pbar = tqdm(imgs)
for a_img_name in pbar:
a_img_name = a_img_name.split(image_ext)[0]
raw_img = cv2.imread(os.path.join(img_dir,
a_img_name + image_ext))
raw_h, raw_w = raw_img.shape[0], raw_img.shape[1]
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: raw_img[:, :, ::-1]}
)
if draw_imgs:
detected_indices = det_scores_r_ >= cfgs.VIS_SCORE
detected_scores = det_scores_r_[detected_indices]
detected_boxes = det_boxes_r_[detected_indices]
detected_categories = det_category_r_[detected_indices]
det_detections_r = draw_box_in_img.draw_boxes_with_label_and_scores(np.squeeze(resized_img, 0),
boxes=detected_boxes,
labels=detected_categories,
scores=detected_scores,
method=1,
in_graph=True)
save_dir = os.path.join('test_hrsc', cfgs.VERSION, 'hrsc2016_img_vis')
tools.mkdir(save_dir)
cv2.imwrite(save_dir + '/{}.jpg'.format(a_img_name),
det_detections_r[:, :, ::-1])
if det_boxes_r_.shape[0] != 0:
resized_h, resized_w = resized_img.shape[1], resized_img.shape[2]
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)
det_boxes_r_ = backward_convert(det_boxes_r_, False)
x_c, y_c, w, h, theta = det_boxes_r_[:, 0], det_boxes_r_[:, 1], det_boxes_r_[:, 2], \
det_boxes_r_[:, 3], det_boxes_r_[:, 4]
boxes_r = np.transpose(np.stack([x_c, y_c, w, h, theta]))
dets_r = np.hstack((det_category_r_.reshape(-1, 1),
det_scores_r_.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
def exucute_detect(self, img_dir, img_name_list):
"""
execute object detect
:param detect_net: detect network
:param img_list: the image dir of detect
:return:
"""
# config gpu to growth train
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
init_op = tf.group(tf.global_variables_initializer(),
tf.local_variables_initializer())
with tf.Session() as sess:
sess.run(init_op)
# restore pretrain weight
restorer, restore_ckpt = self.detect_net.get_restore(
pretrained_model_dir=self.pretrain_model_dir,
restore_from_rpn=False,
is_pretrained=True)
if not restorer is None:
restorer.restore(sess, save_path=restore_ckpt)
print('Successful restore model from {0}'.format(restore_ckpt))
# +++++++++++++++++++++++++++++++++++++start detect+++++++++++++++++++++++++++++++++++++++++++++++++++++=++
all_boxes = []
img_path_list = [
os.path.join(img_dir, img_name) for img_name in img_name_list
]
for index, img_name in enumerate(img_path_list):
bgr_img = cv.imread(img_name)
raw_img = cv.cvtColor(bgr_img, cv.COLOR_BGR2RGB)
resized_img = self.image_process(raw_img)
# expend dimension
image_batch = tf.expand_dims(input=resized_img,
axis=0) # (1, None, None, 3)
start_time = time.time()
feed_dict = self.detect_net.fill_feed_dict(
image_feed=image_batch.eval())
resized_img, (detected_boxes, detected_scores, detected_categories) = \
sess.run(fetches=[resized_img, self.detect_net.inference],
feed_dict=feed_dict) # convert channel from BGR to RGB (cv is BGR)
end_time = time.time()
print("{} cost time : {} ".format(img_name,
(end_time - start_time)))
# draw object image
if self.draw_img:
# select object according to threshold
object_indices = detected_scores >= cfgs.SHOW_SCORE_THRSHOLD
object_scores = detected_scores[object_indices]
object_boxes = detected_boxes[object_indices]
object_categories = detected_categories[object_indices]
final_detections = draw_box_in_img.draw_boxes_with_label_and_scores(
img_array=resized_img,
boxes=object_boxes,
labels=object_categories,
scores=object_scores)
final_detections = cv.cvtColor(final_detections,
cv.COLOR_RGB2BGR)
object_boxes = self.bbox_resize(
bbox=object_boxes,
inputs_shape=resized_img.shape[1:3],
target_shape=raw_img.shape[1:3])
# resize boxes and image shape size to raw input image
detected_boxes = self.bbox_resize(
bbox=detected_boxes,
inputs_shape=(resized_img.shape[0], resized_img.shape[1]),
target_shape=(raw_img.shape[0], resized_img.shape[1]))
# construct detect array for evaluation
detect_bbox_label = np.hstack(
(detected_categories.reshape(-1, 1).astype(np.int32),
detected_scores.reshape(-1, 1), detected_boxes))
all_boxes.append(detect_bbox_label)
# dump bbox to local
makedir(self.object_bbox_save_path)
with open(
os.path.join(self.object_bbox_save_path, 'detections.pkl'),
'wb') as fw:
pickle.dump(all_boxes, fw)
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)
if cfgs.NET_NAME in ['resnet101_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)
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_net.build_whole_detection_network(
input_img_batch=img_batch, gtboxes_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')
if not os.path.exists('./tmp.txt'):
fw = open('./tmp.txt', 'w')
fw.close()
fr = open('./tmp.txt', 'r')
pass_img = fr.readlines()
fr.close()
for count, img_path in enumerate(file_paths):
fw = open('./tmp.txt', 'a+')
if img_path + '\n' in pass_img:
continue
start = timer()
img = cv2.imread(img_path)
box_res = []
label_res = []
score_res = []
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_ = \
sess.run(
[det_boxes_h, det_scores_h, det_category_h],
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])
box_res = np.array(box_res)
label_res = np.array(label_res)
score_res = np.array(score_res)
box_res_, label_res_, score_res_ = [], [], []
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 == 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:
scores = np.array(score_res_)
labels = np.array(label_res_)
boxes = np.array(box_res_)
valid_show = scores > cfgs.SHOW_SCORE_THRSHOLD
scores = scores[valid_show]
boxes = boxes[valid_show]
labels = labels[valid_show]
det_detections_h = draw_box_in_img.draw_boxes_with_label_and_scores(
np.array(img, np.float32),
boxes=np.array(boxes),
labels=np.array(labels),
scores=np.array(scores),
in_graph=False)
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)
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()
# 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(
'%s cost %.3fs' %
(img_path.split('/')[-1].split('.')[0], time_elapsed),
count + 1, len(file_paths))
fw.write('{}\n'.format(img_path))
fw.close()
os.remove('./tmp.txt')
