# img_ori = cv2.imread(args.input_image)
# if args.letterbox_resize:
# img, resize_ratio, dw, dh = letterbox_resize(img_ori, args.new_size[0], args.new_size[1])
# else:
# height_ori, width_ori = img_ori.shape[:2]
# img = cv2.resize(img_ori, tuple(args.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, 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)
num_parallel_calls=args.num_threads
)
val_dataset.prefetch(args.prefetech_buffer)
iterator = val_dataset.make_one_shot_iterator()
image_ids, image, y_true_13, y_true_26, y_true_52 = iterator.get_next()
image_ids.set_shape([None])
y_true = [y_true_13, y_true_26, y_true_52]
image.set_shape([None, args.img_size[1], args.img_size[0], 3])
for y in y_true:
y.set_shape([None, None, None, None, None])
##################
# Model definition
##################
yolo_model = yolov3(args.class_num, args.anchors)
with tf.variable_scope('yolov3'):
pred_feature_maps = yolo_model.forward(image, is_training=is_training)
loss = yolo_model.compute_loss(pred_feature_maps, y_true)
y_pred = yolo_model.predict(pred_feature_maps)
saver_to_restore = tf.train.Saver()
with tf.Session() as sess:
sess.run([tf.global_variables_initializer()])
saver_to_restore.restore(sess, args.restore_path)
print('\n----------- start to eval -----------\n')
val_loss_total, val_loss_xy, val_loss_wh, val_loss_conf, val_loss_class = \
AverageMeter(), AverageMeter(), AverageMeter(), AverageMeter(), AverageMeter()
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
image_ids, image, y_true_13, y_true_26, y_true_52 = iterator.get_next()
y_true = [y_true_13, y_true_26, y_true_52]
# tf.data pipeline will lose the data `static` shape, so we need to set it manually
image_ids.set_shape([None])
image.set_shape([None, None, None, 3])
for y in y_true:
y.set_shape([None, None, None, None, None])
##################
# Model definition
##################
yolo_model = yolov3(args.class_num,
args.anchors,
args.use_label_smooth,
args.use_focal_loss,
args.batch_norm_decay,
args.weight_decay,
use_static_shape=False)
with tf.variable_scope('yolov3'):
pred_feature_maps = yolo_model.forward(image, is_training=is_training)
loss = yolo_model.compute_loss(pred_feature_maps, y_true)
y_pred = yolo_model.predict(pred_feature_maps)
l2_loss = tf.losses.get_regularization_loss()
# setting restore parts and vars to update
saver_to_restore = tf.train.Saver(
var_list=tf.contrib.framework.get_variables_to_restore(
include=args.restore_include, exclude=args.restore_exclude))
update_vars = tf.contrib.framework.get_variables_to_restore(
import os
import sys
import tensorflow as tf
import numpy as np
from YOLOv3.model import yolov3
from YOLOv3.utils.misc_utils import parse_anchors, load_weights
num_class = 80
img_size = 416
weight_path = './data/darknet_weights/yolov3.weights'
save_path = './data/darknet_weights/yolov3.ckpt'
anchors = parse_anchors('./data/yolo_anchors.txt')
model = yolov3(80, anchors)
with tf.Session() as sess:
inputs = tf.placeholder(tf.float32, [1, img_size, img_size, 3])
with tf.variable_scope('yolov3'):
feature_map = model.forward(inputs)
saver = tf.train.Saver(var_list=tf.global_variables(scope='yolov3'))
load_ops = load_weights(tf.global_variables(scope='yolov3'), weight_path)
sess.run(load_ops)
saver.save(sess, save_path=save_path)
print('TensorFlow model checkpoint has been saved to {}'.format(save_path))