def __init__(self,
ctpn_weight_path,
densenet_weight_path,
dict_path,
ctpn_config_path=None,
densenet_config_path=None):
"""
:param ctpn_weight_path: CTPN 模型权重文件路径
:param densenet_weight_path: Densenet 模型权重文件路径
:param dict_path: 字典文件路径
:param ctpn_config_path: CTPN 模型配置文件路径
:param densenet_config_path: Densenet 模型配置文件路径
"""
self.id_to_char = load_dict(dict_path, encoding="utf-8")
# 初始化CTPN模型
if ctpn_config_path is not None:
ctpn_config = CTPN.load_config(ctpn_config_path)
ctpn_config["weight_path"] = ctpn_weight_path
self.ctpn = CTPN(**ctpn_config)
else:
self.ctpn = CTPN()
# 初始化Densenet 模型
if densenet_config_path is not None:
densenet_config = DenseNetOCR.load_config(densenet_config_path)
densenet_config["weight_path"] = densenet_weight_path
self.ocr = DenseNetOCR(**densenet_config)
else:
self.ocr = DenseNetOCR(num_classes=len(self.id_to_char))
parser.add_argument("--epochs", help="迭代数", default=20, type=int)
parser.add_argument("--gpus", help="gpu的数量", default=1, type=int)
parser.add_argument("--images_dir", help="图像位置", default="D:\python_projects\ChineseCalligraphyDetection\data\\train_img")
parser.add_argument("--anno_dir", help="标注文件位置", default="D:\python_projects\ChineseCalligraphyDetection\data\\annotation")
parser.add_argument("--config_file_path", help="模型配置文件位置",
default=default_ctpn_config_path)
parser.add_argument("--weights_file_path", help="模型初始权重文件位置",
default=None)
parser.add_argument("--save_weights_file_path", help="保存模型训练权重文件位置",
default=r'model/cv_weights-ctpnlstm-{epoch:02d}.hdf5')
args = parser.parse_args()
#movefile(args.anno_dir)
K.set_session(get_session(0.8))
config = CTPN.load_config(args.config_file_path)
weights_file_path = args.weights_file_path
if weights_file_path is not None:
config["weight_path"] = weights_file_path
config['num_gpu'] = args.gpus
ctpn = CTPN(**config)
save_weigths_file_path = args.save_weights_file_path
if save_weigths_file_path is None:
try:
if not os.path.exists("model"):
os.makedirs("model")
save_weigths_file_path = "model/weights-ctpnlstm-{epoch:02d}.hdf5"