def _load_crnn(state, model_name="model"):
crnn_args = state[model_name]["args"]
crnn_kwargs = state[model_name]["kwargs"]
crnn = CRNN(*crnn_args, **crnn_kwargs)
crnn.load_state_dict(state[model_name]["state_dict"])
crnn.eval()
crnn = to_cuda_if_available(crnn)
logger.info("Model loaded at epoch: {}".format(state["epoch"]))
logger.info(crnn)
return crnn
def main():
ctpn = CTPN(cfg)
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
sess = tf.Session(config=config)
ctpn.load_ckpt(sess) # ctpn load
if torch.cuda.is_available() and cfg.ALL_GPU:
crnn = CRNN(cfg, 32, 1, len(cfg.KEYS) + 1, 256, 1).cuda()
else:
crnn = CRNN(cfg, 32, 1, len(cfg.KEYS) + 1, 256, 1).cpu()
crnn.eval()
crnn.load_state_dict(torch.load(cfg.CRNN_MODEL)) # crnn load
if cfg.ADJUST_ANGLE:
angle_detector = VGG(cfg) # vgg load
angle_detector.load_weights()
data = DataLoader(cfg)
text = TextGenerator(cfg)
# image_path = raw_input("please input your image path and name:") # get image path
image_path = str('/home/jwm/Desktop/OCR-standard/images/xuanye.jpg')
img = data.load_data(image_path)
t = time.time()
if cfg.ADJUST_ANGLE:
img = rotate(img, angle_detector) # rotate image if necessary
# img = cv2.resize(img, (2000,3000),interpolation=cv2.INTER_CUBIC)
text_recs, detected_img, img = detect(img, data, ctpn, sess) # detect text
results = recognize(img, text_recs, crnn, text,
adjust=True) # recognize text
print("It takes time:{}s".format(time.time() - t))
for key in results:
print(results[key][1])
print('the dice length is {}'.format(len(index2char)))
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
train_loader = dataloader.get_train_loader(data_cfg, char2index) ##获取数据loader
if model_cfg['method'] == 'CRNN':
criterion = nn.CTCLoss(zero_infinity=True)
model = CRNN(nclass, model_cfg).to(device)
elif model_cfg['method'] == 'SAR':
criterion = torch.nn.CrossEntropyLoss(ignore_index=char2index['PAD'])
model = SAR(nclass, model_cfg).to(device)
if not model_cfg['load_model_path'] == '':
model.load_state_dict(torch.load(model_cfg['load_model_path']))
optimizer = optim.Adam(model.parameters(), lr=train_cfg['learning_rate'])
def adjust_learning_rate(optimizer,
global_step,
init_lr=train_cfg['learning_rate'],
decay_rate=train_cfg['decay_rate'],
decay_steps=train_cfg['decay_steps'],
min_lr=train_cfg['min_lr']):
iter_num = global_step // decay_steps
new_leatning_rate = max(init_lr * math.pow(decay_rate, iter_num), min_lr)
for param_group in optimizer.param_groups:
param_group['lr'] = new_leatning_rate
print(f'adjust learning rate to {new_leatning_rate}')