def main():
yolo = YOLO()
inputs = tf.compat.v1.placeholder(dtype=tf.float32, shape=[1, None, None, 3])
feature_y1, feature_y2, feature_y3 = yolo.forward(inputs, class_num, isTrain=False)
pre_boxes, pre_score, pre_label = get_predict_result(feature_y1, feature_y2, feature_y3,
anchors[2], anchors[1], anchors[0],
width, height, class_num,
score_thresh=score_thresh,
iou_thresh=iou_thresh,
max_box=max_box)
init = tf.compat.v1.global_variables_initializer()
saver = tf.train.Saver()
with tf.compat.v1.Session() as sess:
sess.run(init)
ckpt = tf.compat.v1.train.get_checkpoint_state(model_path)
if ckpt and ckpt.model_checkpoint_path:
saver.restore(sess, ckpt.model_checkpoint_path)
Log.add_log("message: load ckpt model:'"+str(ckpt.model_checkpoint_path)+"'")
else:
Log.add_log("message:can not find ckpt model")
# exit(1)
# assert(0)
# dictionary of name of corresponding id
word_dict = tools.get_word_dict(name_file)
# dictionary of per names
color_table = tools.get_color_table(class_num)
for name in os.listdir(val_dir):
img_name = path.join(val_dir, name)
if not path.isfile(img_name):
print("'%s' is not file" %img_name)
continue
img, nw, nh, img_ori, show_img = read_img(img_name, width, height)
if img is None:
Log.add_log("message:'"+str(img)+"' is None")
continue
start = time.perf_counter()
boxes, score, label = sess.run([pre_boxes, pre_score, pre_label], feed_dict={inputs:img})
end = time.perf_counter()
print("%s\t, time:%f s" %(img_name, end-start))
img_ori = tools.draw_img(img_ori, boxes, score, label, word_dict, color_table)
cv2.imshow('img', img_ori)
cv2.waitKey(0)
if save_img:
write_img(img_ori, name)
pass
def main():
yolo = YOLO()
# Placeholder:0
inputs = tf.compat.v1.placeholder(dtype=tf.float32,
shape=[1, None, None, 3])
# yolo/Conv_17/BiasAdd:0, yolo/Conv_9/BiasAdd:0, yolo/Conv_1/BiasAdd:0
feature_y1, feature_y2, feature_y3 = yolo.forward(inputs,
class_num,
isTrain=False)
# concat_9:0, concat_10:0, concat_11:0
pre_boxes, pre_score, pre_label = get_predict_result(
feature_y1,
feature_y2,
feature_y3,
anchors[2],
anchors[1],
anchors[0],
width,
height,
class_num,
score_thresh=score_thresh,
iou_thresh=iou_thresh,
max_box=max_box)
init = tf.compat.v1.global_variables_initializer()
saver = tf.train.Saver()
with tf.compat.v1.Session() as sess:
sess.run(init)
ckpt = tf.compat.v1.train.get_checkpoint_state(ckpt_file_dir)
if ckpt and ckpt.model_checkpoint_path:
saver.restore(sess, ckpt.model_checkpoint_path)
Log.add_log("message: ckpt model:'" +
str(ckpt.model_checkpoint_path) + "'")
else:
Log.add_log("message:no ckpt model")
exit(1)
# save PB model
out_graph = tf.graph_util.convert_variables_to_constants(
sess, sess.graph_def, [
'Placeholder', 'yolo/Conv_1/BiasAdd', 'yolo/Conv_9/BiasAdd',
'yolo/Conv_17/BiasAdd', 'concat_9', 'concat_10', 'concat_11'
]) # "yolo/Conv_13/BiasAdd"
saver_path = tf.train.write_graph(out_graph, "", pd_dir, as_text=False)
print("saver path: ", saver_path)
def main():
class_num = 80
width = 608
height = 608
score_thresh = 0.5
iou_thresh = 0.213
max_box = 50
anchors = 12, 16, 19, 36, 40, 28, 36, 75, 76, 55, 72, 146, 142, 110, 192, 243, 459, 401
anchors = np.asarray(anchors).astype(np.float32).reshape([-1, 3, 2])
model_dir = "./yolo_weights"
name_file = "./data/coco.names"
val_dir = "./coco_test_img"
yolo = YOLO()
inputs = tf.compat.v1.placeholder(dtype=tf.float32,
shape=[1, None, None, 3])
feature_y1, feature_y2, feature_y3 = yolo.forward(inputs,
class_num,
isTrain=False)
pre_boxes, pre_score, pre_label = get_predict_result(
feature_y1,
feature_y2,
feature_y3,
anchors[2],
anchors[1],
anchors[0],
width,
height,
class_num,
score_thresh=score_thresh,
iou_thresh=iou_thresh,
max_box=max_box)
init = tf.compat.v1.global_variables_initializer()
saver = tf.train.Saver()
with tf.compat.v1.Session() as sess:
sess.run(init)
ckpt = tf.compat.v1.train.get_checkpoint_state(model_dir)
if ckpt and ckpt.model_checkpoint_path:
print("restore model from ", ckpt.model_checkpoint_path)
saver.restore(sess, ckpt.model_checkpoint_path)
else:
print("can not find ckpt model")
assert (0)
# id to names
word_dict = tools.get_word_dict(name_file)
# color of corresponding names
color_table = tools.get_color_table(class_num)
for name in os.listdir(val_dir):
img_name = path.join(val_dir, name)
if not path.isfile(img_name):
print("'%s' is not a file" % img_name)
continue
start = time.perf_counter()
img, img_ori = read_img(img_name, width, height)
if img is None:
continue
boxes, score, label = sess.run([pre_boxes, pre_score, pre_label],
feed_dict={inputs: img})
end = time.perf_counter()
print("%s\t, time:%f s" % (img_name, end - start))
img_ori = tools.draw_img(img_ori, boxes, score, label, word_dict,
color_table)
cv2.imshow('img', img_ori)
cv2.waitKey(0)
save_img(img_ori, name)
def __init__(self,model_path,GPU_ratio=0.2):
#----var
JB_flow = 0
class_num = 80
height = 608#416, 608
width = 608#416, 608
score_thresh = 0.5
iou_thresh = 0.213
max_box = 50
anchors = 12, 16, 19, 36, 40, 28, 36, 75, 76, 55, 72, 146, 142, 110, 192, 243, 459, 401
anchors = np.asarray(anchors).astype(np.float32).reshape([-1, 3, 2])
name_file = "./coco.names"
node_dict = {"input": "Placeholder:0",
"pre_boxes": "concat_9:0",
"pre_score": "concat_10:0",
"pre_label": "concat_11:0",
}
#----model extension check
if model_path[-2:] == 'pb':
sess, tf_dict = model_restore_from_pb(model_path, node_dict,GPU_ratio=GPU_ratio)
tf_input = tf_dict['input']
tf_pre_boxes = tf_dict["pre_boxes"]
tf_pre_score = tf_dict['pre_score']
tf_pre_label = tf_dict['pre_label']
else:
width = int(model_path.split("\\")[-1].split(".")[0].split("_")[-1]) # 416, 608
height = width # 416, 608
yolo = YOLO()
tf_input = tf.compat.v1.placeholder(dtype=tf.float32, shape=[1, None, None, 3])
feature_y1, feature_y2, feature_y3 = yolo.forward(tf_input, class_num, isTrain=False)
tf_pre_boxes, tf_pre_score, tf_pre_label = get_predict_result(feature_y1, feature_y2, feature_y3,
anchors[2], anchors[1], anchors[0],
width, height, class_num,
score_thresh=score_thresh,
iou_thresh=iou_thresh,
max_box=max_box)
init = tf.compat.v1.global_variables_initializer()
saver = tf.train.Saver()
#----GPU ratio setting
config = tf.ConfigProto(log_device_placement=True, # 印出目前的運算是使用CPU或GPU
allow_soft_placement=True, # 當設備不存在時允許tf選擇一个存在且可用的設備來繼續執行程式
)
if GPU_ratio is None:
config.gpu_options.allow_growth = True # 依照程式執行所需要的資料來自動調整
else:
config.gpu_options.per_process_gpu_memory_fraction = GPU_ratio # 手動限制GPU資源的使用
sess = tf.compat.v1.Session(config=config)
sess.run(init)
saver.restore(sess, model_path[:-5])
print("JB> Height: {}, width: {}".format(height, width))
#----label to class name
label_dict = tools.get_word_dict(name_file)
#----color of corresponding names
color_table = tools.get_color_table(class_num)
#----local var to global
self.JB_flow = JB_flow # show flow
self.width = width
self.height = height
self.tf_input = tf_input
self.pre_boxes = tf_pre_boxes
self.pre_score = tf_pre_score
self.pre_label = tf_pre_label
self.sess = sess
self.label_dict = label_dict
self.color_table = color_table
