args.nms_topk = 400
# Whether to use the voc 2007 mAP metrics.
args.use_voc_07_metric = False
# args params
args.anchors = parse_anchors(args.anchor_path)
args.classes = read_class_names(args.class_name_path)
args.class_num = len(args.classes)
args.val_img_cnt = len(open(args.val_file, 'r').readlines())
# setting placeholders
is_training = tf.placeholder(dtype=tf.bool, name="phase_train")
handle_flag = tf.placeholder(tf.string, [], name='iterator_handle_flag')
pred_boxes_flag = tf.placeholder(tf.float32, [1, None, None])
pred_scores_flag = tf.placeholder(tf.float32, [1, None, None])
gpu_nms_op = gpu_nms(pred_boxes_flag, pred_scores_flag, args.class_num, args.nms_topk, args.score_threshold,
args.nms_threshold)
##################
# tf.data pipeline
##################
val_dataset = tf.data.TextLineDataset(args.val_file)
val_dataset = val_dataset.batch(1)
val_dataset = val_dataset.map(
lambda x: tf.py_func(get_batch_data,
[x, args.class_num, args.img_size, args.anchors, 'val', False, False, args.letterbox_resize],
[tf.int64, tf.float32, tf.float32, tf.float32, tf.float32]),
num_parallel_calls=args.num_threads
)
val_dataset.prefetch(args.prefetech_buffer)
iterator = val_dataset.make_one_shot_iterator()
# img = img[np.newaxis, :] / 255.
with tf.Session() as sess:
input_data = tf.placeholder(tf.float32,
[1, args.new_size[1], args.new_size[0], 3],
name='input_data')
yolo_model = yolov3(args.num_class, args.anchors)
with tf.variable_scope('yolov3'):
pred_feature_maps = yolo_model.forward(input_data, False)
pred_boxes, pred_confs, pred_probs = yolo_model.predict(pred_feature_maps)
pred_scores = pred_confs * pred_probs
boxes, scores, labels = gpu_nms(pred_boxes,
pred_scores,
args.num_class,
max_boxes=200,
score_thresh=args.score_thresh,
nms_thresh=args.nms_thresh)
saver = tf.train.Saver()
saver.restore(sess, args.restore_path)
test_images = glob.glob(base_path + "test/*.jpg")
results = ""
for test_image in test_images:
img_ori = cv2.imread(test_image)
name = os.path.splitext(os.path.basename(test_image))[
0] # get the filename without extension
if args.letterbox_resize:
img, resize_ratio, dw, dh = letterbox_resize(
def detect_vehicle(input_image):
img_ori = cv2.imread(input_image)
if letterbox_isResize:
img, resize_ratio, dw, dh = letterbox_resize(img_ori, new_size[0],
new_size[1])
else:
height_ori, width_ori = img_ori.shape[:2]
img = cv2.resize(img_ori, tuple(new_size))
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
img = np.asarray(img, np.float32)
img = img[np.newaxis, :] / 255.
with tf.Session() as sess:
input_data = tf.placeholder(tf.float32,
[1, new_size[1], new_size[0], 3],
name='input_data')
yolo_model = yolov3(args.num_class, args.anchors)
with tf.variable_scope('yolov3'):
pred_feature_maps = yolo_model.forward(input_data, False)
pred_boxes, pred_confs, pred_probs = yolo_model.predict(
pred_feature_maps)
pred_scores = pred_confs * pred_probs
boxes, scores, labels = gpu_nms(pred_boxes,
pred_scores,
args.num_class,
max_boxes=200,
score_thresh=0.3,
nms_thresh=0.45)
saver = tf.train.Saver()
saver.restore(sess, restore_path)
boxes_, scores_, labels_ = sess.run([boxes, scores, labels],
feed_dict={input_data: img})
# rescale the coordinates to the original image
if letterbox_isResize:
boxes_[:, [0, 2]] = (boxes_[:, [0, 2]] - dw) / resize_ratio
boxes_[:, [1, 3]] = (boxes_[:, [1, 3]] - dh) / resize_ratio
else:
boxes_[:, [0, 2]] *= (width_ori / float(new_size[0]))
boxes_[:, [1, 3]] *= (height_ori / float(new_size[1]))
# print("box coords:")
# print(boxes_)
# print('*' * 30)
# print("scores:")
# print(scores_)
# print('*' * 30)
# print("labels:")
# print(labels_)
for i in range(len(boxes_)):
x0, y0, x1, y1 = boxes_[i]
plot_one_box(img_ori, [x0, y0, x1, y1],
label=args.classes[labels_[i]] +
', {:.2f}%'.format(scores_[i] * 100),
color=color_table[labels_[i]])
# cv2.imshow('Detection result', img_ori)
detection_result = './detection_result.jpg'
cv2.imwrite(detection_result, img_ori)
# cv2.waitKey(0)
return detection_result