def main(args):
# 覆盖参数
config.USE_SIDE_REFINE = bool(args.use_side_refine)
if args.weight_path is None:
args.weight_path = config.WEIGHT_PATH
config.IMAGES_PER_GPU = 1
config.IMAGE_SHAPE = (720, 720, 3)
config.set_root(args.root)
try:
os.makedirs(args.output_dir)
except:
pass
# 图像路径
image_path_list = file_utils.get_sub_files(args.root)
# 加载模型
m = models.ctpn_net(config, 'test')
m.load_weights(args.weight_path, by_name=True)
# 预测
start_time = datetime.datetime.now()
gen = generator(image_path_list, config.IMAGE_SHAPE)
text_boxes, text_scores, image_metas = m.predict_generator(
generator=gen, steps=len(image_path_list), use_multiprocessing=True)
end_time = datetime.datetime.now()
print("======完成{}张图像评估，耗时:{} 秒".format(len(image_path_list),
end_time - start_time))
# 去除padding
text_boxes = [np_utils.remove_pad(text_box) for text_box in text_boxes]
text_scores = [
np_utils.remove_pad(text_score)[:, 0] for text_score in text_scores
]
image_metas = image_utils.batch_parse_image_meta(image_metas)
# 文本行检测
detector = TextDetector(config)
text_lines = [
detector.detect(boxes, scores, config.IMAGE_SHAPE,
window) for boxes, scores, window in zip(
text_boxes, text_scores, image_metas["window"])
]
# 还原检测文本行边框到原始图像坐标
text_lines = [
image_utils.recover_detect_quad(boxes, window, scale)
for boxes, window, scale in zip(text_lines, image_metas["window"],
image_metas["scale"])
]
# 写入文档中
for image_path, boxes in zip(image_path_list, text_lines):
output_filename = os.path.splitext(
'res_' + os.path.basename(image_path))[0] + '.txt'
with open(os.path.join(args.output_dir, output_filename),
mode='w') as f:
for box in boxes.astype(np.int32):
f.write("{},{},{},{},{},{},{},{}\r\n".format(
box[0], box[1], box[2], box[3], box[4], box[5], box[6],
box[7]))
def main(args):
# 覆盖参数 rewrite the parameters
config.USE_SIDE_REFINE = bool(args.use_side_refine)
if args.weight_path is not None:
config.WEIGHT_PATH = args.weight_path
config.IMAGES_PER_GPU = 1
config.IMAGE_SHAPE = (1024, 1024, 3)
# 图像路径 path of the images
image_path_list = file_utils.get_sub_files(args.image_dir)
# 加载模型 load the model
m = models.ctpn_net(config, 'test')
m.load_weights(config.WEIGHT_PATH, by_name=True)
# 预测 predicting
start_time = datetime.datetime.now()
gen = generator(image_path_list, config.IMAGE_SHAPE)
text_boxes, text_scores, image_metas = m.predict_generator(
generator=gen, steps=len(image_path_list), use_multiprocessing=True)
end_time = datetime.datetime.now()
print("====== Image No.{}:Evaluation completed，time:{} second".format(
len(image_path_list), end_time - start_time))
# 去除padding removing the padding
text_boxes = [np_utils.remove_pad(text_box) for text_box in text_boxes]
text_scores = [
np_utils.remove_pad(text_score)[:, 0] for text_score in text_scores
]
image_metas = image_utils.batch_parse_image_meta(image_metas)
# 文本行检测 text Dector
detector = TextDetector(config)
text_lines = [
detector.detect(boxes, scores, config.IMAGE_SHAPE,
window) for boxes, scores, window in zip(
text_boxes, text_scores, image_metas["window"])
]
# 还原检测文本行边框到原始图像坐标 restore the text boundary to the original image
text_lines = [
image_utils.recover_detect_quad(boxes, window, scale)
for boxes, window, scale in zip(text_lines, image_metas["window"],
image_metas["scale"])
]
# 写入文档中 writing into a txt file
for image_path, boxes in zip(image_path_list, text_lines):
output_filename = os.path.splitext(
'res_' + os.path.basename(image_path))[0] + '.txt'
with open(os.path.join(args.output_dir, output_filename),
mode='w') as f:
for box in boxes.astype(np.int32):
f.write("{},{},{},{},{},{},{},{}\r\n".format(
box[0], box[1], box[2], box[3], box[4], box[5], box[6],
box[7]))
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)