def crnnSource():
alphabet = keys.alphabet
converter = util.strLabelConverter(alphabet)
model = crnn.CRNN(32, 1, len(alphabet) + 1, 256, 1).cuda()
path = './crnn/samples/netCRNN63.pth'
model.load_state_dict(torch.load(path))
return model, converter
def batch_test(dirpath):
alphabet = keys_crnn.alphabet
#print(len(alphabet))
#input('\ninput:')
converter = util.strLabelConverter(alphabet)
# model = crnn.CRNN(32, 1, len(alphabet) + 1, 256, 1).cuda()
model = crnn.CRNN(32, 1, len(alphabet) + 1, 256, 1)
path = './samples/model_acc97.pth'
model.load_state_dict(torch.load(path))
#print(model)
paths=glob.glob(os.path.join(dirpath,'*.[jp][pn]g'))
for i in paths:
print(i)
image = Image.open(i).convert('L')
#print(image.size)
scale = image.size[1] * 1.0 / 32
w = image.size[0] / scale
w = int(w)
#print("width:" + str(w))
transformer = dataset.resizeNormalize((w, 32))
# image = transformer(image).cuda()
image = transformer(image)
image = image.view(1, *image.size())
image = Variable(image)
model.eval()
preds = model(image)
#print(preds.shape)
_, preds = preds.max(2)
#print(preds.shape)
preds = preds.squeeze(1)
preds = preds.transpose(-1, 0).contiguous().view(-1)
preds_size = Variable(torch.IntTensor([preds.size(0)]))
raw_pred = converter.decode(preds.data, preds_size.data, raw=True)
sim_pred = converter.decode(preds.data, preds_size.data, raw=False)
print(sim_pred)
def crnnSource():
alphabet = keys_crnn.alphabet
converter = util.strLabelConverter(alphabet)
if torch.cuda.is_available() and GPU:
print("GPU")
model = crnn.CRNN(32, 1, len(alphabet) + 1, 256, 1).cuda()
else:
print("CPU")
model = crnn.CRNN(32, 1, len(alphabet) + 1, 256, 1).cpu()
#model = crnn.CRNN(32, 1, len(alphabet) + 1, 256, 1).cuda()
path = './CRNN/crnn/samples/newmodel.pth'
# 将GUP训练模型的权重转换为单个CPU可用
trainWeights = torch.load(path, map_location=lambda storage, loc: storage)
modelWeights = OrderedDict()
for k, v in trainWeights.items():
name = k.replace('module.', '')
modelWeights[name] = v
model.eval()
model.load_state_dict(modelWeights)
return model, converter
def crnnSource():
alphabet = keys.alphabet
converter = util.strLabelConverter(alphabet)
model = crnn.CRNN(32, 1, len(alphabet)+1, 256, 1).cuda()
path = './crnn/samples/crnn_trained_model.pth'
model.load_state_dict(torch.load(path))
return model,converter
def crnnModel():
alphabet = keys.alphabet
# define a converter for convert rnn results to string
converter = util.strLabelConverter(alphabet)
model = crnn.CRNN(32, 1, len(alphabet) + 1, 256, 1).cuda()
path = './crnn/samples/CRNN.pth'
model.load_state_dict(torch.load(path))
return model, converter
def crnnSource():
alphabet = keys.alphabet
converter = util.strLabelConverter(alphabet)
if torch.cuda.is_available() and GPU:
model = crnn.CRNN(32, 1, len(alphabet) + 1, 256, 1).cuda()
else:
model = crnn.CRNN(32, 1, len(alphabet) + 1, 256, 1).cpu()
path = './crnn/samples/model_acc97.pth'
model.eval()
model.load_state_dict(torch.load(path))
return model, converter
def crnnSource():
alphabet = keys_crnn.alphabet
converter = util.strLabelConverter(alphabet)
if torch.cuda.is_available() and GPU:
print("GPU")
model = crnn.CRNN(32, 1, len(alphabet) + 1, 256, 1).cuda()
else:
print("CPU")
model = crnn.CRNN(32, 1, len(alphabet) + 1, 256, 1).cpu()
print(os.getcwd())
path = '../ModelSet/Chinese-OCR/crnn/samples/model_acc97.pth' # 模型识别种类：5530
# path = '../ModelSet/Chinese-OCR/crnn/samples/mixed_second_finetune_acc97p7.pth' # 对比效果：6736，需要同步修改keys_crnn.py
model.eval()
# 带参数：不使用GPU
model.load_state_dict(torch.load(path, map_location='cpu'))
return model, converter
target = str2label(lexicon, 30)
inputN = input.repeat(lexSize, 1, 1)
logProb = -torch.nn.CTC_forwardBackward(inputN, target, true, inputN.new())
_, idx = torch.max(logProb, 1)
idx = idx[1]
return lexicon[idx]
model = crnn.CRNN(32, 1, 37, 256, 1)
if use_gpu:
model = model.cuda()
print('loading pretrained model from %s' % model_path)
model.load_state_dict(
torch.load(model_path, map_location=lambda storage, loc: storage))
print('end loading')
converter = util.strLabelConverter(alphabet)
transformer = dataset.keep_ratio_normalize(True)
criterion = CTCLoss()
File = open(sv_path, 'w')
im_ls = os.listdir(im_dir)
#im_ls = im_ls[:100]
for nm in im_ls:
if nm.endswith('.png'):
im_path = im_dir + nm
image = Image.open(im_path).convert('L')
image = transformer(image)
if use_gpu:
image = image.cuda()
image = image.view(1, *image.size())
image = Variable(image)
sampler_mode = dataset.randomSequentialSampler(train_dataset,
opt.batchSize)
else:
sampler_mode = None
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=opt.batchSize,
shuffle=False,
sampler=sampler_mode,
num_workers=int(opt.workers),
collate_fn=dataset.alignCollate(
imgH=opt.imgH,
imgW=opt.imgW,
keep_ratio=opt.keep_ratio))
test_dataset = dataset.lmdbDataset(root=opt.valroot)
nc = 1
converter = util.strLabelConverter(opt.alphabet)
criterion = CTCLoss()
# custom weights initialization called on crnn
def weights_init(m):
classname = m.__class__.__name__
if classname.find('Conv') != -1:
m.weight.data.normal_(0.0, 0.02)
elif classname.find('BatchNorm') != -1:
m.weight.data.normal_(1.0, 0.02)
m.bias.data.fill_(0)
crnn = crnn.CRNN(opt.imgH, nc, nclass, opt.nh, 1)
import os
path1=os.path.dirname(__file__)
path2=os.path.dirname(path1)
sys.path.append(path1)
sys.path.append(path2)
#from .keys import alphabetChinese as alphabet
from keys import alphabetChinese as alphabet
import onnxruntime as rt
sys.path.append("..")
#from .util import strLabelConverter, resizeNormalize
from util import strLabelConverter, resizeNormalize
converter = strLabelConverter(''.join(alphabet))
def softmax(x):
x_row_max = x.max(axis=-1)
x_row_max = x_row_max.reshape(list(x.shape)[:-1]+[1])
x = x - x_row_max
x_exp = np.exp(x)
x_exp_row_sum = x_exp.sum(axis=-1).reshape(list(x.shape)[:-1]+[1])
softmax = x_exp / x_exp_row_sum
return softmax
class CRNNHandle:
def __init__(self, model_path):
self.sess = rt.InferenceSession(model_path)
