shuffle=True,
num_workers=int(Config.data_worker),
collate_fn=lib.dataset.alignCollate(imgH=Config.img_height,
imgW=Config.img_width))
n_class = len(Config.alphabet) + 1 # for python3
# n_class = len(Config.alphabet.decode('utf-8')) + 1 # for python2
print("alphabet class num is %s" % n_class)
converter = lib.convert.strLabelConverter(Config.alphabet)
# converter = lib.convert.StrConverter(Config.alphabet)
# print(converter.dict)
criterion = CTCLoss()
net = Net.CRNN(n_class)
print(net)
net.apply(lib.utility.weights_init)
# 加载权重
net.load_state_dict(torch.load('w160_bs64_model/netCRNN_199_200.pth'))
image = torch.FloatTensor(Config.batch_size, 3, Config.img_height,
Config.img_width)
text = torch.IntTensor(Config.batch_size * 5)
length = torch.IntTensor(Config.batch_size)
if cuda:
net.cuda()
image = image.cuda()
_, preds = preds.max(2)
preds = preds.transpose(1, 0).contiguous().view(-1)
preds_size = Variable(torch.IntTensor([preds.size(0)]))
# 预测输出，解码成文字
sim_pred = converter.decode(preds.data, preds_size.data, raw=False)
result += (format(sim_pred) + '\n')
return result
if __name__ == '__main__':
model = Net.CRNN(nclass)
if running_mode == 'gpu' and torch.cuda.is_available():
model = model.cuda()
model.load_state_dict(torch.load(crnn_model_path))
else:
model.load_state_dict(torch.load(crnn_model_path, map_location='cpu'))
print('loading pretrained model from {0}'.format(crnn_model_path))
files = sorted(os.listdir(IMG_ROOT))
for file in files:
started = time.time()
full_path = os.path.join(IMG_ROOT, file)
print("=============================================")
print("ocr image is %s" % full_path)
image = Image.open(full_path)