def main(args):
set_gpu_growth()
# 加载标注
annotation_files = file_utils.get_sub_files(config.IMAGE_GT_DIR)
image_annotations = [
reader.load_annotation(file, config.IMAGE_DIR)
for file in annotation_files
]
# 过滤不存在的图像，ICDAR2017中部分图像找不到
image_annotations = [
ann for ann in image_annotations if os.path.exists(ann['image_path'])
]
# 加载模型
m = models.ctpn_net(config, 'train')
models.compile(m,
config,
loss_names=['ctpn_regress_loss', 'ctpn_class_loss'])
if args.init_epochs > 0:
m.load_weights(args.weight_path, by_name=True)
else:
m.load_weights(config.PRE_TRAINED_WEIGHT, by_name=True)
m.summary()
# 生成器
gen = generator(image_annotations, config.IMAGES_PER_GPU,
config.IMAGE_SHAPE, config.ANCHORS_WIDTH,
config.MAX_GT_INSTANCES)
# 训练
m.fit_generator(gen,
steps_per_epoch=len(image_annotations) //
config.IMAGES_PER_GPU,
epochs=args.epochs,
initial_epoch=args.init_epochs,
verbose=True,
callbacks=get_call_back(),
use_multiprocessing=True)
# 保存模型
m.save(config.WEIGHT_PATH)
if __name__ == '__main__':
parse = argparse.ArgumentParser()
parse.add_argument("--epochs", type=int, default=50, help="epochs")
parse.add_argument("--init_epochs", type=int, default=0, help="epochs")
parse.add_argument("--weight_path",
type=str,
default=None,
help="weight path")
argments = parse.parse_args(sys.argv[1:])
main(argments)
def main(args):
set_gpu_growth()
# 加载标注 load the annotations
annotation_files = file_utils.get_sub_files(config.IMAGE_GT_DIR)
image_annotations = [
reader.load_annotation(file, config.IMAGE_DIR)
for file in annotation_files
]
# 过滤不存在的图像，ICDAR2017中部分图像找不到 remove the missing images
image_annotations = [
ann for ann in image_annotations if os.path.exists(ann['image_path'])
]
# 加载模型 load the model
m = models.ctpn_net(config, 'train')
models.compile(m,
config,
loss_names=[
'ctpn_regress_loss', 'ctpn_class_loss',
'side_regress_loss'
])
# 增加度量 increasing the metrics
output = models.get_layer(m, 'ctpn_target').output
models.add_metrics(
m, ['gt_num', 'pos_num', 'neg_num', 'gt_min_iou', 'gt_avg_iou'],
output[-5:])
if args.init_epochs > 0:
m.load_weights(args.weight_path, by_name=True)
else:
m.load_weights(config.PRE_TRAINED_WEIGHT, by_name=True)
m.summary()
# 生成器 generator
gen = generator(image_annotations, config.IMAGES_PER_GPU,
config.IMAGE_SHAPE, config.ANCHORS_WIDTH,
config.MAX_GT_INSTANCES)
# 训练 training
m.fit_generator(gen,
steps_per_epoch=len(image_annotations) //
config.IMAGES_PER_GPU * 2,
epochs=args.epochs,
initial_epoch=args.init_epochs,
verbose=True,
callbacks=get_call_back(),
workers=2,
use_multiprocessing=True)
# 保存模型 model saving
m.save(config.WEIGHT_PATH)
def main(args):
set_gpu_growth()
# 加载标注
annotation_files = file_utils.get_sub_files(config.IMAGE_GT_DIR)
# annotation_files_1 = [file for file in annotation_files if 'rctw' in file]
annotation_files = [file for file in annotation_files if 'rects' in file]
# annotation_files = annotation_files_1 + annotation_files_2
image_annotations = [
reader.load_annotation(file, config.IMAGE_DIR)
for file in annotation_files
]
# 过滤不存在的图像，ICDAR2017中部分图像找不到
image_annotations = [
ann for ann in image_annotations if os.path.exists(ann['image_path'])
]
# 加载模型
m = models.ctpn_net(config, 'train')
models.compile(m,
config,
loss_names=[
'ctpn_regress_loss', 'ctpn_class_loss',
'side_regress_loss'
])
# 增加度量
output = models.get_layer(m, 'ctpn_target').output
models.add_metrics(
m, ['gt_num', 'pos_num', 'neg_num', 'gt_min_iou', 'gt_avg_iou'],
output[-5:])
if args.init_epochs > 0:
m.load_weights(args.weight_path, by_name=True)
else:
m.load_weights(config.PRE_TRAINED_WEIGHT, by_name=True)
m.summary()
# 生成器
gen = generator(image_annotations[:-100],
config.IMAGES_PER_GPU,
config.IMAGE_SHAPE,
config.ANCHORS_WIDTH,
config.MAX_GT_INSTANCES,
horizontal_flip=False,
random_crop=False)
val_gen = generator(image_annotations[-100:], config.IMAGES_PER_GPU,
config.IMAGE_SHAPE, config.ANCHORS_WIDTH,
config.MAX_GT_INSTANCES)
# 训练
m.fit_generator(gen,
steps_per_epoch=len(image_annotations) //
config.IMAGES_PER_GPU * 2,
epochs=args.epochs,
initial_epoch=args.init_epochs,
validation_data=val_gen,
validation_steps=100 // config.IMAGES_PER_GPU,
verbose=True,
callbacks=get_call_back(),
workers=2,
use_multiprocessing=True)
# 保存模型
m.save(config.WEIGHT_PATH)
def main(args):
set_gpu_growth()
config.set_root(args.root)
image_annotations = load_folder_annotation(args.root)
if len(image_annotations) < 5:
print("Too small dataset...")
return
# gen = generator(image_annotations[:-100],
# config.IMAGES_PER_GPU,
# config.IMAGE_SHAPE,
# config.ANCHORS_WIDTH,
# config.MAX_GT_INSTANCES,
# horizontal_flip=False,
# random_crop=False)
# val_gen = generator(image_annotations[-100:],
# config.IMAGES_PER_GPU,
# config.IMAGE_SHAPE,
# config.ANCHORS_WIDTH,
# config.MAX_GT_INSTANCES)
# for bat in range(100):
# print(bat)
# val, _ = next(gen)
# for key in val.keys():
# print( val[key].shape, end="," )
# print()
# val, _ = next(val_gen)
# for key in val.keys():
# print( val[key].shape, end="," )
# print()
# print()
# exit(1)
# 加载模型
m = models.ctpn_net(config, 'train')
models.compile(m,
config,
loss_names=[
'ctpn_regress_loss', 'ctpn_class_loss',
'side_regress_loss'
])
# 增加度量
output = models.get_layer(m, 'ctpn_target').output
models.add_metrics(
m, ['gt_num', 'pos_num', 'neg_num', 'gt_min_iou', 'gt_avg_iou'],
output[-5:])
# 从0开始的话，用resnet50
if args.weight_path is None:
args.weight_path = config.WEIGHT_PATH
# get current epoch from file name.
if args.init_epochs == 0:
res = re.match(r".*ctpn\.(\d+)\.h5", args.weight_path)
if res is not None:
args.init_epochs = int(res.group(1))
m.load_weights(args.weight_path, by_name=True)
m.summary()
# print( len( image_annotations[:-100]), len(image_annotations[-100:]) )
# 生成器
# 前面100条作为训练集，后面100条做成测试集。
gen = generator(image_annotations[:-10],
config.IMAGES_PER_GPU,
config.IMAGE_SHAPE,
config.ANCHORS_WIDTH,
config.MAX_GT_INSTANCES,
horizontal_flip=False,
random_crop=False)
val_gen = generator(image_annotations[-10:], config.IMAGES_PER_GPU,
config.IMAGE_SHAPE, config.ANCHORS_WIDTH,
config.MAX_GT_INSTANCES)
# 训练
m.fit_generator(gen,
steps_per_epoch=len(image_annotations) //
config.IMAGES_PER_GPU * 2,
epochs=args.epochs,
initial_epoch=args.init_epochs,
validation_data=val_gen,
validation_steps=100 // config.IMAGES_PER_GPU,
verbose=True,
callbacks=get_call_back(),
workers=args.jobs,
use_multiprocessing=True)
#
# # 保存模型
path = os.path.split(config.WEIGHT_PATH)
m.save(os.sep.join([path[0], "ctpn.%03d.h5" % (args.epochs)]))