y_true = [y_true_13, y_true_26, y_true_52]
else:
y_true = [y_true_13, y_true_26]
# tf.data pipeline will lose the data shape, so we need to set it manually
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
##################
# define yolo-v3 model here
#yolo_model = yolov3(args.class_num, args.anchors)
yolo_model = yolov3_tiny(args.class_num, args.anchors)
# the input variables name of .pb
image = tf.identity(image, name='inputs')
with tf.variable_scope(net_name):
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)
# the output variables name of .pb
boxes, confs, probs = y_pred
boxes = tf.identity(boxes, name='boxes')
confs = tf.identity(confs, name='confs')
probs = tf.identity(probs, name='probs')
################
vid = cv2.VideoCapture(args.input_video)
video_width = int(vid.get(3))
video_height = int(vid.get(4))
#print(video_width,video_height)
if args.save_video:
fourcc = cv2.VideoWriter_fourcc('m', 'p', '4', 'v')
videoWriter = cv2.VideoWriter('video_result.mp4', fourcc, 20.0,
(video_width, video_height))
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_tiny(args.num_class, args.anchors)
with tf.variable_scope('yolov3_tiny'):
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=30,
score_thresh=0.5,
iou_thresh=0.5)
saver = tf.train.Saver()
saver.restore(sess, args.restore_path)
