def main():
yolo = YOLO(class_num, anchors)
# Placeholder:0
inputs = tf.compat.v1.placeholder(dtype=tf.float32, shape=[1, None, None, 3])
feature_y1, feature_y2, feature_y3 = yolo.forward(inputs, isTrain=False)
# concat_9:0, concat_10:0, concat_11:0
pre_boxes, pre_score, pre_label = yolo.get_predict_result(feature_y1, feature_y2, feature_y3, class_num,
score_thresh=config.val_score_thresh, iou_thresh=config.iou_thresh, max_box=config.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 __init_args(self):
Log.add_log("data agument strategy : "+str(self.agument))
# data augment strategy
self.multi_scale_img = self.agument[0] # multiscale zoom the image
self.keep_img_shape = self.agument[1] # keep image's shape when we reshape the image
self.flip_img = self.agument[2] # flip image
self.gray_img = self.agument[3] # gray image
self.label_smooth = self.agument[4] # label smooth strategy
self.erase_img = self.agument[5] # random erase image
self.invert_img = self.agument[6] # invert image pixel
self.rotate_img = self.agument[7] # random rotate image
Log.add_log("message: begin to initial images path")
# init imgs path
for voc_dir in self.data_dirs:
for img_name in os.listdir(voc_dir):
img_path = path.join(voc_dir, img_name)
label_path = img_path.replace("JPEGImages", "Annotations")
label_path = label_path.replace(img_name.split('.')[-1], "xml")
if not path.isfile(img_path):
Log.add_log("warning:VOC image'"+str(img_path)+"'is not a file")
continue
if not path.isfile(label_path):
Log.add_log("warning:VOC label'"+str(label_path)+"'if not a file")
continue
self.imgs_path.append(img_path)
self.labels_path.append(label_path)
self.num_imgs += 1
Log.add_log("message:initialize VOC dataset complete, there are "+str(self.num_imgs)+" pictures in all")
if self.num_imgs <= 0:
raise ValueError("there are 0 pictures to train in all")
return
def __init_args(self):
Log.add_log("data agument strategy : "+str(self.agument))
# 初始化数据增强策略的参数
self.multi_scale_img = self.agument[0] # 多尺度缩放图片
self.keep_img_shape = self.agument[1] # keep image's shape when we reshape the image
self.flip_img = self.agument[2] # flip image
self.gray_img = self.agument[3] # gray image
self.label_smooth = self.agument[4] # label smooth strategy
self.erase_img = self.agument[5] # random erase image
self.invert_img = self.agument[6] # invert image pixel
self.rotate_img = self.agument[7] # rotate image
Log.add_log("message:开始初始化路径")
# init imgs path
self.imgs_path = tools.read_file(self.data_file)
if not self.imgs_path:
Log.add_log("error:imgs_path文件不存在")
raise ValueError(str(self.data_file) + ":文件中不存在路径")
Log.add_log("message:一共有 "+str(len(self.imgs_path)) + " 张图片")
# init labels path
for img in self.imgs_path:
label_path = img.replace("JPEGImages", "labels")
label_path = label_path.replace(img.split('.')[-1], "txt")
self.labels_path.append(label_path)
Log.add_log("message:数据路径初始化完成")
return
def __init_args(self):
Log.add_log("data agument strategy : " + str(self.agument))
# parameters of data augment strategy
self.multi_scale_img = self.agument[0] # multiscale zoom the image
self.keep_img_shape = self.agument[
1] # keep image's shape when we reshape the image
self.flip_img = self.agument[2] # flip image
self.gray_img = self.agument[3] # gray image
self.label_smooth = self.agument[4] # label smooth strategy
self.erase_img = self.agument[5] # random erase image
self.invert_img = self.agument[6] # invert image pixel
self.rotate_img = self.agument[7] # rotate image
Log.add_log("message: begin to initial images path")
# init imgs path
self.imgs_path = tools.read_file(self.data_file)
if not self.imgs_path:
Log.add_log("error:imgs_path not exist")
raise ValueError("no file find in :" + str(self.data_file))
Log.add_log("message:there are " + str(len(self.imgs_path)) +
" pictures in all")
# init labels path
for img in self.imgs_path:
label_path = img.replace("JPEGImages", "labels")
label_path = label_path.replace(img.split('.')[-1], "txt")
self.labels_path.append(label_path)
Log.add_log(
"message: initialize images path and corresponding labels complete"
)
return
def __get_data(self):
'''
加载 batch_size 的标签和数据
return:imgs, label_y1, label_y2, label_y3
'''
# 十个 batch 随机一次 size
if self.multi_scale_img and (self.num_batch % 10 == 0):
random_size = random.randint(13, 23) * 32
self.width = self.height = random_size
imgs = []
labels_y1, labels_y2, labels_y3 = [], [], []
count = 0
while count < self.batch_size:
curr_index = random.randint(0, self.num_imgs - 1)
img_name = self.imgs_path[curr_index]
label_name = self.labels_path[curr_index]
# probility of flip image is 0.5
if self.agument[2] and (np.random.random() < 0.5):
self.flip_img = True
# print("flip")
else:
self.flip_img = False
img, new_w, new_h, test_img = self.read_img(img_name)
label_y1, label_y2, label_y3, test_result = self.read_label(
label_name, self.names_dict, self.anchors, new_w, new_h)
# 显示最后的结果
if self.data_debug:
test_img = tools.draw_img(test_img, test_result, None, None,
None, None)
cv2.imshow("letterbox_img", test_img)
cv2.waitKey(0)
if img is None:
Log.add_log("VOC 文件'" + img_name + "'读取异常")
continue
if label_y1 is None:
Log.add_log("VOC 文件'" + label_name + "'读取异常")
continue
imgs.append(img)
labels_y1.append(label_y1)
labels_y2.append(label_y2)
labels_y3.append(label_y3)
count += 1
self.num_batch += 1
imgs = np.asarray(imgs)
labels_y1 = np.asarray(labels_y1)
labels_y2 = np.asarray(labels_y2)
labels_y3 = np.asarray(labels_y3)
return imgs, labels_y1, labels_y2, labels_y3
def __get_data(self):
'''
load batch_size labels and images data
return:imgs, label_y1, label_y2, label_y3
'''
# random resize image per ten batch
if self.multi_scale_img and (self.num_batch % 10 == 0):
random_size = random.randint(13, 23) * 32
self.width = self.height = random_size
imgs = []
labels_y1, labels_y2, labels_y3 = [], [], []
count = 0
while count < self.batch_size:
curr_index = random.randint(0, len(self.imgs_path) - 1)
img_name = self.imgs_path[curr_index]
label_name = self.labels_path[curr_index]
# probility of flip image is 0.5
if self.agument[2] and (np.random.random() < 0.5):
self.flip_img = True
else:
self.flip_img = False
img, new_w, new_h, test_img = self.read_img(img_name)
label_y1, label_y2, label_y3, test_result = self.read_label(
label_name, self.anchors, new_w, new_h)
# show data augment result
if self.data_debug:
test_img = tools.draw_img(test_img, test_result, None, None,
None, None)
cv2.imshow("letterbox_img", test_img)
cv2.waitKey(0)
if img is None:
Log.add_log("file '" + img_name + "' is None")
continue
if label_y1 is None:
Log.add_log("file'" + label_name + "' is None")
continue
imgs.append(img)
labels_y1.append(label_y1)
labels_y2.append(label_y2)
labels_y3.append(label_y3)
count += 1
self.num_batch += 1
imgs = np.asarray(imgs)
labels_y1 = np.asarray(labels_y1)
labels_y2 = np.asarray(labels_y2)
labels_y3 = np.asarray(labels_y3)
return imgs, labels_y1, labels_y2, labels_y3
def __remove(self, img_file, xml_file):
self.imgs_path.remove(img_file)
self.labels_path.remove(xml_file)
self.num_imgs -= 1
if not len(self.imgs_path) == len(self.labels_path):
Log.add_log(
"error:after remove the file:{} , the number of label and image is not equal"
.format(img_file))
assert (0)
return
def config_optimizer(optimizer_name, lr_init, momentum=0.99):
Log.add_log("message:配置优化器:'" + str(optimizer_name) + "'")
if optimizer_name == 'momentum':
return tf.compat.v1.train.MomentumOptimizer(lr_init, momentum=momentum)
elif optimizer_name == 'adam':
return tf.compat.v1.train.AdamOptimizer(learning_rate=lr_init)
elif optimizer_name == 'sgd':
return tf.compat.v1.train.GradientDescentOptimizer(learning_rate=lr_init)
else:
Log.add_log("error:不支持的优化器类型:'" + str(optimizer_name) + "'")
raise ValueError(str(optimizer_name) + ":不支持的优化器类型")
def save_img(img, name):
'''
img:需要保存的 mat 图片
name:保存的图片名
'''
if not path.isdir(config.save_dir):
Log.add_log("message:创建文件夹'" + str(config.save_dir) + "'")
os.mkdir(config.save_dir)
img_name = path.join(config.save_dir, name)
cv2.imwrite(img_name, img)
return
def save_img(img, name):
'''
img: mat
name:
'''
if not path.isdir(config.save_dir):
Log.add_log("message: create floder'" + str(config.save_dir) + "'")
os.mkdir(config.save_dir)
img_name = path.join(config.save_dir, name)
cv2.imwrite(img_name, img)
return
def save_img(img, name):
'''
img: mat
name: saved name
'''
voc_save_dir = config.voc_save_dir
if not path.isdir(voc_save_dir):
Log.add_log("message: create folder'" + str(voc_save_dir) + "'")
os.mkdir(voc_save_dir)
img_name = path.join(voc_save_dir, name)
cv2.imwrite(img_name, img)
return
def read_label(self, label_file):
'''
读取 label_file, 并生成 label_y1, label_y2, label_y3
return:label_y1, label_y2, label_y3
'''
contents = tools.read_file(label_file)
if not contents:
return None, None, None
label_y1 = np.zeros(
(self.height // 32, self.width // 32, 3, 5 + self.class_num),
np.float32)
label_y2 = np.zeros(
(self.height // 16, self.width // 16, 3, 5 + self.class_num),
np.float32)
label_y3 = np.zeros(
(self.height // 8, self.width // 8, 3, 5 + self.class_num),
np.float32)
y_true = [label_y3, label_y2, label_y1]
ratio = {0: 8, 1: 16, 2: 32}
for label in contents:
label = label.split()
if len(label) != 5:
Log.add_log("error:文件'" + str(label_file) + "'标签数量错误")
raise ValueError(str(label_file) + ":标签数量错误")
label_id = int(label[0])
box = np.asarray(label[1:5]).astype(
np.float32) # label中保存的就是 x,y,w,h
best_giou = 0
best_index = 0
for i in range(len(self.anchors)):
min_wh = np.minimum(box[2:4], self.anchors[i])
max_wh = np.maximum(box[2:4], self.anchors[i])
giou = (min_wh[0] * min_wh[1]) / (max_wh[0] * max_wh[1])
if giou > best_giou:
best_giou = giou
best_index = i
# 012->0, 345->1, 678->2
x = int(np.floor(box[0] * self.width / ratio[best_index // 3]))
y = int(np.floor(box[1] * self.height / ratio[best_index // 3]))
k = best_index % 3
y_true[best_index // 3][y, x, k, 0:4] = box
y_true[best_index // 3][y, x, k, 4:5] = 1.0
y_true[best_index // 3][y, x, k, 5 + label_id] = 1.0
return label_y1, label_y2, label_y3
def config_optimizer(optimizer_name, lr, momentum=0.99):
Log.add_log("message: configuration optimizer:'" + str(optimizer_name) +
"'")
if optimizer_name == 'momentum':
return tf.compat.v1.train.MomentumOptimizer(lr, momentum=momentum)
elif optimizer_name == 'adam':
return tf.compat.v1.train.AdamOptimizer(learning_rate=lr)
elif optimizer_name == 'sgd':
return tf.compat.v1.train.GradientDescentOptimizer(learning_rate=lr)
else:
Log.add_log("error:unsupported type of optimizer:'" +
str(optimizer_name) + "'")
raise ValueError(
str(optimizer_name) + ":unsupported type of optimizer")
def __init_args(self):
Log.add_log("message: begin to initial images path")
# init imgs path
for voc_dir in self.data_dirs:
for img_name in os.listdir(voc_dir):
img_path = path.join(voc_dir, img_name)
label_path = img_path.replace("JPEGImages", "Annotations")
label_path = label_path.replace(img_name.split('.')[-1], "xml")
if not path.isfile(img_path):
Log.add_log("warning:VOC image'" + str(img_path) +
"'is not a file")
continue
if not path.isfile(label_path):
Log.add_log("warning:VOC label'" + str(label_path) +
"'if not a file")
continue
self.imgs_path.append(img_path)
self.labels_path.append(label_path)
self.num_imgs += 1
self.steps_per_epoch = int(self.num_imgs / self.batch_size)
Log.add_log("message:initialize VOC dataset complete, there are " +
str(self.num_imgs) + " pictures in all")
if self.num_imgs <= 0:
raise ValueError("there are 0 pictures to train in all")
return
def __get_data(self):
'''
load batch_size labels and images
return:imgs, label_y1, label_y2, label_y3
'''
imgs = []
labels_y1, labels_y2, labels_y3 = [], [], []
count = 0
while count < self.batch_size:
curr_index = random.randint(0, self.num_imgs - 1)
img_name = self.imgs_path[curr_index]
label_name = self.labels_path[curr_index]
img, ori_w, ori_h, test_img = self.read_img(img_name)
if img is None:
Log.add_log("img file'" + img_name +
"'reading exception, will be deleted")
self.__remove(img_name, label_name)
continue
label_y1, label_y2, label_y3, test_result = self.read_label(
label_name, self.names_dict, self.anchors)
if label_y1 is None:
Log.add_log("label file'" + label_name +
"'reading exception, will be deleted")
self.__remove(img_name, label_name)
continue
# show data agument result
if self.data_debug:
test_img = tools.draw_img(test_img, test_result, None, None,
None, None)
cv2.imshow("letterbox_img", test_img)
cv2.waitKey(0)
imgs.append(img)
labels_y1.append(label_y1)
labels_y2.append(label_y2)
labels_y3.append(label_y3)
count += 1
self.num_batch += 1
imgs = np.asarray(imgs)
labels_y1 = np.asarray(labels_y1)
labels_y2 = np.asarray(labels_y2)
labels_y3 = np.asarray(labels_y3)
return imgs, labels_y1, labels_y2, labels_y3
def config_lr(lr_type, lr_init, epoch_globalstep=0):
Log.add_log("message:配置学习率:'" + str(lr_type) + "', 初始学习率:"+str(lr_init))
if lr_type == 'piecewise':
lr = tf.compat.v1.train.piecewise_constant(epoch_globalstep,
config.piecewise_boundaries, config.piecewise_values)
elif lr_type == 'exponential':
lr = tf.compat.v1.train.exponential_decay(learning_rate=lr_init,
global_step=epoch_globalstep, decay_steps=10, decay_rate=0.99, staircase=True)
elif lr_type =='constant':
lr = lr_init
else:
Log.add_log("error:不支持的学习率类型:'" + str(lr_type) + "'")
raise ValueError(str(lr_type) + ":不支持的学习率类型")
return tf.maximum(lr, config.lr_lower)
def __get_data(self):
'''
加载 batch_size 的标签和数据
return:imgs, label_y1, label_y2, label_y3
'''
# 十个 batch 随机一次 size
if self.multi_scale_img and (self.num_batch % 10 == 0):
random_size = random.randint(10, 19) * 32
self.width = self.height = random_size
imgs = []
labels_y1, labels_y2, labels_y3 = [], [], []
count = 0
while count < self.batch_size:
curr_index = random.randint(0, len(self.imgs_path) - 1)
img_name = self.imgs_path[curr_index]
label_name = self.labels_path[curr_index]
img = self.read_img(img_name)
label_y1, label_y2, label_y3 = self.read_label(label_name)
if img is None:
Log.add_log("文件'" + img_name + "'读取异常")
continue
if label_y1 is None:
Log.add_log("文件'" + label_name + "'读取异常")
continue
imgs.append(img)
labels_y1.append(label_y1)
labels_y2.append(label_y2)
labels_y3.append(label_y3)
count += 1
self.num_batch += 1
imgs = np.asarray(imgs)
labels_y1 = np.asarray(labels_y1)
labels_y2 = np.asarray(labels_y2)
labels_y3 = np.asarray(labels_y3)
return imgs, labels_y1, labels_y2, labels_y3
def __init_args(self):
Log.add_log("data agument strategy : " + str(self.agument))
# 初始化数据增强策略的参数
self.multi_scale_img = self.agument[0] # 多尺度缩放图片
self.keep_img_shape = self.agument[
1] # keep image's shape when we reshape the image
self.flip_img = self.agument[2] # flip image
self.gray_img = self.agument[3] # gray image
self.label_smooth = self.agument[4] # label smooth strategy
self.erase_img = self.agument[5] # random erase image
self.invert_img = self.agument[6] # invert image pixel
self.rotate_img = self.agument[7] # random rotate image
Log.add_log("message:开始初始化路径")
# init imgs path
for voc_dir in self.data_dirs:
for img_name in os.listdir(voc_dir):
img_path = path.join(voc_dir, img_name)
label_path = img_path.replace("JPEGImages", "Annotations")
label_path = label_path.replace(img_name.split('.')[-1], "xml")
if not path.isfile(img_path):
Log.add_log("warning:VOC 图片文件'" + str(img_path) + "'不存在")
continue
if not path.isfile(label_path):
Log.add_log("warning:VOC 标签文件'" + str(label_path) + "'不存在")
continue
self.imgs_path.append(img_path)
self.labels_path.append(label_path)
self.num_imgs += 1
Log.add_log("message:VOC 数据初始化完成,一共有 " + str(self.num_imgs) + " 张图片")
if self.num_imgs <= 0:
raise ValueError("没有可训练的图片, 程序退出")
return
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 config_lr(lr_type,
lr_init,
lr_lower=1e-6,
piecewise_boundaries=None,
piecewise_values=None,
epoch=0):
Log.add_log("message:configuration lr:'" + str(lr_type) +
"', initial lr:" + str(lr_init))
if lr_type == 'piecewise':
lr = tf.compat.v1.train.piecewise_constant(epoch, piecewise_boundaries,
piecewise_values)
elif lr_type == 'exponential':
lr = tf.compat.v1.train.exponential_decay(learning_rate=lr_init,
global_step=epoch,
decay_steps=10,
decay_rate=0.99,
staircase=True)
elif lr_type == 'constant':
lr = lr_init
else:
Log.add_log("error:unsupported lr type:'" + str(lr_type) + "'")
raise ValueError(str(lr_type) + ": unsupported type of learning rate")
return tf.maximum(lr, lr_lower)
def main():
yolo = YOLO(config.class_num, config.anchors)
inputs = tf.compat.v1.placeholder(dtype=tf.float32, shape=[1, None, None, 3])
feature_y1, feature_y2, feature_y3 = yolo.forward(inputs, isTrain=False)
pre_boxes, pre_score, pre_label = yolo.get_predict_result(feature_y1, feature_y2, feature_y3, config.class_num,
score_thresh=config.val_score_thresh, iou_thresh=config.iou_thresh, max_box=config.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(config.model_path)
if ckpt and ckpt.model_checkpoint_path:
saver.restore(sess, ckpt.model_checkpoint_path)
Log.add_log("message:存在 ckpt 模型:'"+str(ckpt.model_checkpoint_path)+"'")
else:
Log.add_log("message:不存在 ckpt 模型")
# exit(1)
# 名字字典
word_dict = tools.get_word_dict(config.name_file)
# 色表
color_table = tools.get_color_table(config.class_num)
width = config.width
height = config.height
for name in os.listdir(config.val_dir):
img_name = path.join(config.val_dir, name)
if not path.isfile(img_name):
print("'%s'不是图片" %img_name)
continue
start = time.perf_counter()
img, img_ori = read_img(img_name, width, height)
if img is None:
Log.add_log("message:'"+str(img)+"'图片读取错误")
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 config.save_img:
save_img(img_ori, name)
def __init_args(self):
Log.add_log("message: begin to initial images path")
# init imgs path
self.imgs_path = tools.read_file(self.data_file)
if not self.imgs_path:
Log.add_log("error:imgs_path not exist")
raise ValueError("no file find in :" + str(self.data_file))
self.num_imgs = len(self.imgs_path)
self.steps_per_epoch = int(self.num_imgs / self.batch_size)
Log.add_log("message:there are " + str(self.num_imgs) +
" pictures in all")
# init labels path
for img in self.imgs_path:
label_path = img.replace("JPEGImages", "labels")
label_path = label_path.replace(img.split('.')[-1], "txt")
self.labels_path.append(label_path)
Log.add_log(
"message: initialize images path and corresponding labels complete"
)
return
def __init_args(self):
Log.add_log("message:开始初始化路径")
# init imgs path
for voc_dir in self.data_dirs:
for img_name in os.listdir(voc_dir):
img_path = path.join(voc_dir, img_name)
label_path = img_path.replace("JPEGImages", "Annotations")
label_path = label_path.replace(img_name.split('.')[-1], "xml")
if not path.isfile(img_path):
Log.add_log("warning:VOC 图片文件'"+str(img_path)+"'不存在")
continue
if not path.isfile(label_path):
Log.add_log("warning:VOC 标签文件'"+str(label_path)+"'不存在")
continue
self.imgs_path.append(img_path)
self.labels_path.append(label_path)
self.num_imgs += 1
Log.add_log("message:VOC 数据初始化完成,一共有 "+str(self.num_imgs)+" 张图片")
if self.num_imgs <= 0:
raise ValueError("没有可训练的图片, 程序退出")
return
def __init_args(self):
Log.add_log("message:开始初始化路径")
# init imgs path
self.imgs_path = tools.read_file(self.data_file)
if not self.imgs_path:
Log.add_log("error:imgs_path文件不存在")
raise ValueError(str(self.data_file) + ":文件中不存在路径")
print("data:一共有 %d 张图片" %(len(self.imgs_path)))
# init labels path
for img in self.imgs_path:
label_path = img.replace("JPEGImages", "labels")
label_path = label_path.replace(img.split('.')[-1], "txt")
self.labels_path.append(label_path)
Log.add_log("message:数据路径初始化完成")
return
score_thresh=config.val_score_thresh,
iou_thresh=config.iou_thresh,
max_box=config.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)
if __name__ == "__main__":
Log.add_log("message: convert ckpt model to pb model")
main()
def backward():
yolo = YOLO()
data = Data(voc_root_dir,
names_file,
class_num,
batch_size,
anchors,
is_tiny=False,
size=size)
inputs = tf.compat.v1.placeholder(dtype=tf.float32,
shape=[batch_size, None, None, 3])
y1_true = tf.compat.v1.placeholder(
dtype=tf.float32, shape=[batch_size, None, None, 3, 4 + 1 + class_num])
y2_true = tf.compat.v1.placeholder(
dtype=tf.float32, shape=[batch_size, None, None, 3, 4 + 1 + class_num])
y3_true = tf.compat.v1.placeholder(
dtype=tf.float32, shape=[batch_size, None, None, 3, 4 + 1 + class_num])
feature_y1, feature_y2, feature_y3 = yolo.forward(
inputs, class_num, weight_decay=weight_decay, isTrain=True)
global_step = tf.Variable(0, trainable=False)
# loss value of yolov4
loss = Loss().yolo_loss([feature_y1, feature_y2, feature_y3],
[y1_true, y2_true, y3_true],
[anchors[2], anchors[1], anchors[0]],
width,
height,
class_num,
cls_normalizer=cls_normalizer,
iou_normalizer=iou_normalizer,
iou_thresh=iou_thresh,
prob_thresh=prob_thresh,
score_thresh=score_thresh)
l2_loss = tf.compat.v1.losses.get_regularization_loss()
epoch = compute_curr_epoch(global_step, batch_size, len(data.imgs_path))
lr = Lr.config_lr(lr_type, lr_init, lr_lower=lr_lower, \
piecewise_boundaries=piecewise_boundaries, \
piecewise_values=piecewise_values, epoch=epoch)
optimizer = Optimizer.config_optimizer(optimizer_type, lr, momentum)
update_ops = tf.compat.v1.get_collection(tf.compat.v1.GraphKeys.UPDATE_OPS)
with tf.control_dependencies(update_ops):
gvs = optimizer.compute_gradients(loss + l2_loss)
clip_grad_var = [
gv if gv[0] is None else [tf.clip_by_norm(gv[0], 100.), gv[1]]
for gv in gvs
]
train_step = optimizer.apply_gradients(clip_grad_var,
global_step=global_step)
# initialize
init = tf.compat.v1.global_variables_initializer()
saver = tf.compat.v1.train.Saver()
with tf.compat.v1.Session() as sess:
sess.run(init)
step = 0
ckpt = tf.compat.v1.train.get_checkpoint_state(model_path)
if ckpt and ckpt.model_checkpoint_path:
saver.restore(sess, ckpt.model_checkpoint_path)
step = ckpt.model_checkpoint_path.split('/')[-1].split('-')[-1]
step = eval(step)
Log.add_log("message: load ckpt model, global_step=" + str(step))
else:
Log.add_log("message:can not fint ckpt model")
curr_epoch = step // data.steps_per_epoch
while curr_epoch < total_epoch:
for _ in range(data.steps_per_epoch):
start = time.perf_counter()
batch_img, y1, y2, y3 = next(data)
_, loss_, step, lr_ = sess.run(
[train_step, loss, global_step, lr],
feed_dict={
inputs: batch_img,
y1_true: y1,
y2_true: y2,
y3_true: y3
})
end = time.perf_counter()
if (loss_ > 1e3) and (step > 1e3):
Log.add_log("error:loss exception, loss_value = " +
str(loss_))
''' break the process or lower learning rate '''
raise ValueError("error:loss exception, loss_value = " +
str(loss_) +
", please lower your learning rate")
# lr = tf.math.maximum(tf.math.divide(lr, 10), config.lr_lower)
if step % 5 == 2:
print(
"step: %6d, epoch: %3d, loss: %.5g\t, wh: %3d, lr:%.5g\t, time: %5f s"
% (step, curr_epoch, loss_, width, lr_, end - start))
Log.add_loss(str(step) + "\t" + str(loss_))
curr_epoch += 1
if curr_epoch % save_per_epoch == 0:
# save ckpt model
Log.add_log("message: save ckpt model, step=" + str(step) +
", lr=" + str(lr_))
saver.save(sess,
path.join(model_path, model_name),
global_step=step)
Log.add_log("message: save final ckpt model, step=" + str(step))
saver.save(sess, path.join(model_path, model_name), global_step=step)
return 0
raise ValueError("error:loss exception, loss_value = " +
str(loss_) +
", please lower your learning rate")
# lr = tf.math.maximum(tf.math.divide(lr, 10), config.lr_lower)
if step % 5 == 2:
print(
"step: %6d, epoch: %3d, loss: %.5g\t, wh: %3d, lr:%.5g\t, time: %5f s"
% (step, curr_epoch, loss_, width, lr_, end - start))
Log.add_loss(str(step) + "\t" + str(loss_))
curr_epoch += 1
if curr_epoch % save_per_epoch == 0:
# save ckpt model
Log.add_log("message: save ckpt model, step=" + str(step) +
", lr=" + str(lr_))
saver.save(sess,
path.join(model_path, model_name),
global_step=step)
Log.add_log("message: save final ckpt model, step=" + str(step))
saver.save(sess, path.join(model_path, model_name), global_step=step)
return 0
if __name__ == "__main__":
Log.add_log("message: into VOC backward function")
# Log.add_loss("###########")
backward()
inputs: batch_img,
y1_true: y1,
y2_true: y2,
y3_true: y3
})
end = time.perf_counter()
print(
"step: %6d, loss: %.5g\t, w: %3d, h: %3d, lr:%.5g\t, time: %5f s"
% (step, loss_, data.width, data.height, lr_, end - start))
if step % 5 == 2:
Log.add_loss(str(step) + "\t" + str(loss_))
if (step + 1) % config.save_step == 0:
Log.add_log("message:当前运行模型保存, step=" + str(step) + ", lr=" +
str(lr_))
saver.save(sess,
path.join(config.model_path, config.model_name),
global_step=0)
Log.add_log("message:训练完成保存模型, step=" + str(step))
saver.save(sess,
path.join(config.model_path, config.model_name),
global_step=step)
return 0
if __name__ == "__main__":
Log.add_log("message:进入 backward 函数")
Log.add_loss("###########")
backward()
def backward():
yolo = YOLO(config.class_num,
config.anchors,
width=config.width,
height=config.height)
data = Data(config.train_file,
config.class_num,
config.batch_size,
config.anchors,
config.multi_scale_img,
width=config.width,
height=config.height)
inputs = tf.compat.v1.placeholder(dtype=tf.float32,
shape=[config.batch_size, None, None, 3])
y1_true = tf.compat.v1.placeholder(
dtype=tf.float32,
shape=[config.batch_size, None, None, 3, 4 + 1 + config.class_num])
y2_true = tf.compat.v1.placeholder(
dtype=tf.float32,
shape=[config.batch_size, None, None, 3, 4 + 1 + config.class_num])
y3_true = tf.compat.v1.placeholder(
dtype=tf.float32,
shape=[config.batch_size, None, None, 3, 4 + 1 + config.class_num])
feature_y1, feature_y2, feature_y3 = yolo.forward(
inputs, weight_decay=config.weight_decay, isTrain=True)
global_step = tf.Variable(0, trainable=False)
# 损失 yolov4
loss = yolo.get_loss_v4(feature_y1, feature_y2, feature_y3, y1_true,
y2_true, y3_true, config.cls_normalizer,
config.ignore_thresh, config.prob_thresh,
config.score_thresh)
l2_loss = tf.compat.v1.losses.get_regularization_loss()
epoch = compute_curr_epoch(global_step, config.batch_size,
len(data.imgs_path))
lr = config_lr(config.lr_type, config.lr_init, epoch)
optimizer = config_optimizer(config.optimizer_type, lr, config.momentum)
update_ops = tf.compat.v1.get_collection(tf.compat.v1.GraphKeys.UPDATE_OPS)
with tf.control_dependencies(update_ops):
gvs = optimizer.compute_gradients(loss + l2_loss)
clip_grad_var = [
gv if gv[0] is None else [tf.clip_by_norm(gv[0], 100.), gv[1]]
for gv in gvs
]
train_step = optimizer.apply_gradients(clip_grad_var,
global_step=global_step)
# 初始化
init = tf.compat.v1.global_variables_initializer()
saver = tf.compat.v1.train.Saver()
with tf.compat.v1.Session() as sess:
sess.run(init)
step = 0
ckpt = tf.compat.v1.train.get_checkpoint_state(config.model_path)
if ckpt and ckpt.model_checkpoint_path:
saver.restore(sess, ckpt.model_checkpoint_path)
step = ckpt.model_checkpoint_path.split('/')[-1].split('-')[-1]
step = eval(step)
Log.add_log("message:存在 ckpt 模型, global_step=" + str(step))
else:
Log.add_log("message:不存在 ckpt 模型")
# 一共迭代这么多次
total_steps = np.ceil(config.total_epoch * len(data.imgs_path) /
config.batch_size)
while step < total_steps:
start = time.perf_counter()
batch_img, y1, y2, y3 = next(data)
_, loss_, step, lr_ = sess.run([train_step, loss, global_step, lr],
feed_dict={
inputs: batch_img,
y1_true: y1,
y2_true: y2,
y3_true: y3
})
end = time.perf_counter()
print(
"step: %6d, loss: %.5g\t, w: %3d, h: %3d, lr:%.5g\t, time: %5f s"
% (step, loss_, data.width, data.height, lr_, end - start))
if step % 5 == 2:
Log.add_loss(str(step) + "\t" + str(loss_))
if (step + 1) % config.save_step == 0:
Log.add_log("message:当前运行模型保存, step=" + str(step) + ", lr=" +
str(lr_))
saver.save(sess,
path.join(config.model_path, config.model_name),
global_step=0)
Log.add_log("message:训练完成保存模型, step=" + str(step))
saver.save(sess,
path.join(config.model_path, config.model_name),
global_step=step)
return 0
def backward():
class_num = config.voc_class_num
yolo = YOLO(class_num,
config.voc_anchors,
width=config.width,
height=config.height)
data = Data(config.voc_root_dir,
config.voc_dir_ls,
"./data/voc.names",
class_num,
config.batch_size,
config.voc_anchors,
config.data_augment,
config.width,
config.height,
data_debug=config.data_debug)
inputs = tf.compat.v1.placeholder(dtype=tf.float32,
shape=[config.batch_size, None, None, 3])
y1_true = tf.compat.v1.placeholder(
dtype=tf.float32,
shape=[config.batch_size, None, None, 3, 4 + 1 + class_num])
y2_true = tf.compat.v1.placeholder(
dtype=tf.float32,
shape=[config.batch_size, None, None, 3, 4 + 1 + class_num])
y3_true = tf.compat.v1.placeholder(
dtype=tf.float32,
shape=[config.batch_size, None, None, 3, 4 + 1 + class_num])
feature_y1, feature_y2, feature_y3 = yolo.forward(
inputs, weight_decay=config.weight_decay, isTrain=True)
global_step = tf.Variable(0, trainable=False)
# 损失 yolov4
loss = yolo.get_loss_v4(feature_y1, feature_y2, feature_y3, y1_true,
y2_true, y3_true, config.cls_normalizer,
config.ignore_thresh, config.prob_thresh,
config.score_thresh)
l2_loss = tf.compat.v1.losses.get_regularization_loss()
epoch = compute_curr_epoch(global_step, config.batch_size, data.num_imgs)
lr = config_lr(config.lr_type, config.lr_init, epoch)
optimizer = config_optimizer(config.optimizer_type, lr, config.momentum)
update_ops = tf.compat.v1.get_collection(tf.compat.v1.GraphKeys.UPDATE_OPS)
with tf.control_dependencies(update_ops):
gvs = optimizer.compute_gradients(loss + l2_loss)
clip_grad_var = [
gv if gv[0] is None else [tf.clip_by_norm(gv[0], 100.), gv[1]]
for gv in gvs
]
train_step = optimizer.apply_gradients(clip_grad_var,
global_step=global_step)
# 初始化
init = tf.compat.v1.global_variables_initializer()
saver = tf.compat.v1.train.Saver()
with tf.compat.v1.Session() as sess:
sess.run(init)
step = 0
ckpt = tf.compat.v1.train.get_checkpoint_state(config.voc_model_path)
if ckpt and ckpt.model_checkpoint_path:
saver.restore(sess, ckpt.model_checkpoint_path)
step = ckpt.model_checkpoint_path.split('/')[-1].split('-')[-1]
step = eval(step)
Log.add_log("message:存在 ckpt 模型, global_step=" + str(step))
else:
Log.add_log("message:不存在 ckpt 模型")
# 一共迭代这么多次
total_steps = np.ceil(config.total_epoch * data.num_imgs /
config.batch_size)
while step < total_steps:
start = time.perf_counter()
batch_img, y1, y2, y3 = next(data)
_, loss_, step, lr_ = sess.run([train_step, loss, global_step, lr],
feed_dict={
inputs: batch_img,
y1_true: y1,
y2_true: y2,
y3_true: y3
})
end = time.perf_counter()
print(
"step: %6d, loss: %.5g\t, w: %3d, h: %3d, lr:%.5g\t, time: %5f s"
% (step, loss_, data.width, data.height, lr_, end - start))
if (loss_ > 1e3) and (step > 1e3):
Log.add_log("error:loss exception, loss_value = " + str(loss_))
''' break the process or lower learning rate '''
raise ValueError("error:loss exception, loss_value = " +
str(loss_) +
", please lower your learning rate")
# lr = tf.math.maximum(tf.math.divide(lr, 10), config.lr_lower)
if step % 5 == 2:
Log.add_loss(str(step) + "\t" + str(loss_))
if (step + 1) % config.save_step == 0:
if config.save_ckpt_model:
# save ckpt model
Log.add_log("message: save ckpt model, step=" + str(step) +
", lr=" + str(lr_))
saver.save(sess,
path.join(config.voc_model_path,
config.voc_model_name),
global_step=step)
if config.save_pb_model:
Log.add_log("message: save pb model, step=" + str(step) +
", lr=" + str(lr_))
pb_model_name = path.join(
config.voc_model_path,
config.voc_model_name) + '-' + str(step) + ".pb"
constant_graph = graph_util.convert_variables_to_constants(
sess, sess.graph_def, [
'yolo/Conv_1/BiasAdd', 'yolo/Conv_9/BiasAdd',
'yolo/Conv_17/BiasAdd'
])
# save PB model
with tf.gfile.FastGFile(pb_model_name, mode='wb') as f:
f.write(constant_graph.SerializeToString())
# save ckpt model
if config.save_ckpt_model:
Log.add_log("message: save final ckpt model, step=" + str(step))
saver.save(sess,
path.join(config.voc_model_path, config.voc_model_name),
global_step=step)
# save pb model
if config.save_pb_model:
Log.add_log("message: save final pb model, step=" + str(step))
pb_model_name = path.join(
config.voc_model_path,
config.voc_model_name) + '-' + str(step) + ".pb"
constant_graph = graph_util.convert_variables_to_constants(
sess, sess.graph_def, [
'yolo/Conv_1/BiasAdd', 'yolo/Conv_9/BiasAdd',
'yolo/Conv_17/BiasAdd'
])
# save PB model
with tf.gfile.FastGFile(pb_model_name, mode='wb') as f:
f.write(constant_graph.SerializeToString())
return 0
