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)
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 new_model(nclass, preModel):
# 定义你自己的模型
if torch.cuda.is_available() and opt.ngpu:
model = crnn.CRNN(32, 1, nclass + 1, 256, 1).cuda()
else:
model = crnn.CRNN(32, 1, nclass + 1, 256, 1).cpu()
modelDict = model.state_dict() ##
preModelDict = preModel.state_dict() ##
preModelDict = {k: v for k, v in preModelDict.items() if 'rnn.1' not in k}
modelDict.update(preModelDict) ##更新权重
model.load_state_dict(modelDict) ##加载预训练模型权重
return model
def crnnSource(self, dir_model):
alphabet = '0123456789abcdefghijklmnopqrstuvwxyz' # keys.alphabet
converter = utils.strLabelConverter(alphabet)
if self.gpuid == '-1':
print('cpu-version')
model = crnn.CRNN(32, 1, len(alphabet) + 1, 256)
else:
print('gpu-version')
model = crnn.CRNN(32, 1, len(alphabet) + 1, 256).cuda()
path = dir_model + '/bib_number/CRNN/crnn.pth' # './crnn/samples/netCRNN63.pth'
model.load_state_dict(torch.load(path))
print('Loaded network {:s}'.format(path))
return model, converter
def __init__(self, model_path):
#def crnnSource(model_path, use_gpu=True):
alphabet = keys.alphabet # Chinese words
self.converter = crnn_utils.strLabelConverter(alphabet)
# note that in https://github.com/bear63/sceneReco support multi GPU.
# model = crnn.CRNN(32, 1, len(alphabet)+1, 256, 1).cuda()
self.model = crnn.CRNN(32, 1, len(alphabet) + 1, 256)
self.cpu_model = crnn.CRNN(32, 1, len(alphabet) + 1, 256)
if torch.cuda.is_available():
self.model = self.model.cuda()
print('loading pretrained model from %s' % model_path)
#model_path = './crnn/samples/netCRNN63.pth'
model_state_dict = torch.load(model_path)
self.model.load_state_dict(model_state_dict)
self.cpu_model.load_state_dict(model_state_dict)
def __init__(self, model_path, gpu_id=None):
'''
初始化pytorch模型
:param model_path: 模型地址(可以是模型的参数或者参数和计算图一起保存的文件)
:param gpu_id: 在哪一块gpu上运行
'''
self.gpu_id = gpu_id
self.converter = utils.strLabelConverter(alphabet)
if self.gpu_id is not None and isinstance(
self.gpu_id, int) and torch.cuda.is_available():
checkpoint = torch.load(model_path)
self.device = torch.device("cuda:%s" % self.gpu_id)
else:
checkpoint = torch.load(model_path, map_location='cpu')
self.device = torch.device("cpu")
print('text recognition running on device:', self.device)
self.net = crnn.CRNN(CNN=CNN,
RNN=RNN,
nIn=n_in,
n_class=new_class,
nHidden=n_hidden,
nLayer=n_layer,
dropout=dp)
self.net.load_state_dict(checkpoint['state_dict'])
self.net.to(self.device)
self.net.eval()
self.transform = transforms.Compose([transforms.ToTensor()])
def recognize(image_path, alphabet, snapshot, gpu):
model = crnn.CRNN(32, 1, 37, 256)
if torch.cuda.is_available():
model = model.cuda()
print('loading pretrained model from %s' % snapshot)
model.load_state_dict(torch.load(snapshot))
converter = utils.strLabelConverter(alphabet)
transformer = dataset.resizeNormalize((100, 32))
image = Image.open(image_path).convert('L')
image = transformer(image)
if torch.cuda.is_available():
image = image.cuda()
image = image.view(1, *image.size())
image = Variable(image)
model.eval()
preds = model(image)
_, preds = preds.max(2)
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('%-20s => %-20s' % (raw_pred, sim_pred))
return sim_pred
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 load_images_to_predict():
# load model
model_path = './expr/netCRNN_199_423.pth'
alphabet = '0123456789,.:(%$!^&-/);<~|`>?+=_[]{}"\'@#*ABCDEFGHIJKLMNOPQRSTUVWXYZ\ '
imgH = 32 # should be 32
nclass = len(alphabet) + 1
nhiddenstate = 256
model = crnn.CRNN(imgH, 1, nclass, nhiddenstate)
if torch.cuda.is_available():
model = model.cuda()
print('loading pretrained model from %s' % model_path)
model.load_state_dict({k.replace('module.',''):v for k,v in torch.load(model_path).items()})
# load image
#img = [os.path.basename(f) for f in glob.glob(IMAGE_FOLDER)]
#filenames, extns = [os.path.splitext(f)[0:1] for f in glob.glob(IMAGE_FOLDER)]
#exts = [os.path.splitext(f)[1] for f in glob.glob(IMAGE_FOLDER)]
#img_files = [s + e for s , e in (filenames, extns)]
for img in glob.glob(IMAGE_FOLDER):
image = Image.open(img).convert('L')
words_list = []
#with open(BB_FOLDER + img.split('/')[1].split('.')[0]+'.txt', 'r') as boxes:
print('Ïmage name:', os.path.basename(img))
with open(BB_INT_FOLDER + os.path.basename(img).split('.')[0]+'.txt', 'r') as boxes:
for line in csv.reader(boxes):
box = [int(string, 10) for string in line[0:8]]
boxImg = image.crop((box[0], box[1], box[4], box[5]))
words = predict_this_box(boxImg, model, alphabet)
words_list.append(words)
with open( TXT_RESULT_FLD + os.path.basename(img).split('.')[0] +'.txt', 'w+') as resultfile:
for line in words_list:
resultfile.writelines(line+'\n')
def load_images_to_predict():
# load model
model_path = './expr/netCRNN_99_500.pth'
alphabet = '0123456789,.:(%$!^&-/);<~|`>?+=_[]{}"\'@#*ABCDEFGHIJKLMNOPQRSTUVWXYZ\ '
imgH = 32 # should be 32
nclass = len(alphabet) + 1
nhiddenstate = 256
model = crnn.CRNN(imgH, 1, nclass, nhiddenstate)
if torch.cuda.is_available():
model = model.cuda()
print('loading pretrained model from %s' % model_path)
model.load_state_dict({
k.replace('module.', ''): v
for k, v in torch.load(model_path).items()
})
# load image
filenames = [os.path.splitext(f)[0] for f in glob.glob("data_test/*.jpg")]
jpg_files = [s + ".jpg" for s in filenames]
for jpg in jpg_files:
image = Image.open(jpg).convert('L')
words_list = []
with open('boundingbox/' + jpg.split('/')[1].split('.')[0] + '.txt',
'r') as boxes:
for line in csv.reader(boxes):
box = [int(string, 10) for string in line[0:8]]
boxImg = image.crop((box[0], box[1], box[4], box[5]))
words = predict_this_box(boxImg, model, alphabet)
words_list.append(words)
with open('test_result/' + jpg.split('/')[1].split('.')[0] + '.txt',
'w+') as resultfile:
for line in words_list:
resultfile.writelines(line + '\n')
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 __init__(self, weight):
super(Ocr, self).__init__()
# 读取字母表
with open('../data/lol.alphabet', encoding='utf-8') as f:
self.alphabet = f.read().strip()
self.net = crnn.CRNN(32, 1, len(self.alphabet) + 1, 256).to('cuda')
self.net.load_state_dict(torch.load(weight))
self.converter = utils.StrLabelConverter(self.alphabet)
def __init__(self, model_path='./crnn.pytorch-master/data/crnn.pth'):
self.model = crnn.CRNN(32, 1, 37, 256)
if torch.cuda.is_available():
self.model = self.model.cuda()
print('loading pretrained self.model from %s' % model_path)
self.model.load_state_dict(torch.load(model_path))
self.alphabet = '0123456789abcdefghijklmnopqrstuvwxyz'
self.converter = utils.strLabelConverter(self.alphabet)
def __init__(self, model_path):
self.model = crnn.CRNN(32, 1, 37, 256)
if torch.cuda.is_available():
self.model = self.model.cuda()
self.model.load_state_dict(torch.load(model_path))
self.converter = utils.strLabelConverter(
'0123456789abcdefghijklmnopqrstuvwxyz')
self.model.eval()
def main():
if not os.path.exists(opt.output):
os.makedirs(opt.output)
converter = utils.strLabelConverter(opt.alphabet)
collate = dataset.AlignCollate()
train_dataset = dataset.TextLineDataset(text_file=opt.train_list, transform=dataset.ResizeNormalize(100, 32), converter=converter)
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=opt.batchsize, shuffle=True,
num_workers=opt.num_workers, collate_fn=collate)
test_dataset = dataset.TextLineDataset(text_file=opt.train_list, transform=dataset.ResizeNormalize(100, 32), converter=converter)
test_loader = torch.utils.data.DataLoader(test_dataset, shuffle=False, batch_size=opt.batchsize,
num_workers=opt.num_workers, collate_fn=collate)
criterion = nn.CTCLoss()
import models.crnn as crnn
crnn = crnn.CRNN(opt.imgH, opt.nc, opt.num_classes, opt.nh)
crnn.apply(utils.weights_init)
if opt.pretrained != '':
print('loading pretrained model from %s' % opt.pretrained)
crnn.load_state_dict(torch.load(opt.pretrained), strict=False)
print(crnn)
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
crnn = crnn.to(device)
criterion = criterion.to(device)
# setup optimizer
optimizer = optim.Adam(crnn.parameters(), lr=opt.lr)
for epoch in range(opt.num_epochs):
loss_avg = 0.0
i = 0
while i < len(train_loader):
time0 = time.time()
# 训练
train_iter = iter(train_loader)
cost = trainBatch(crnn, train_iter, criterion, optimizer, device) # 一个批次，一个批次训练
loss_avg += cost
i += 1
if i % opt.interval == 0:
print('[%d/%d][%d/%d] Loss: %f Time: %f s' %
(epoch, opt.num_epochs, i, len(train_loader), loss_avg,
time.time() - time0))
loss_avg = 0.0
if (epoch + 1) % opt.valinterval == 0:
val(crnn, test_loader, criterion, converter=converter, device=device, max_iter=100)
def main():
model_path = './data/crnn.pth'
img_path = './data/demo.png'
alphabet = '0123456789abcdefghijklmnopqrstuvwxyz'
model = crnn.CRNN(32, 1, 37, 256)
if torch.cuda.is_available():
model = model.cuda()
print('loading pretrained model from %s' % model_path)
model.load_state_dict(torch.load(model_path))
converter = utils.strLabelConverter(alphabet)
#transformer = dataset.resizeNormalize((100, 32))
dir = Path("images/").glob(
"**/*.jpg") # Modifiable. Takes all .jpg files in the given directory
for image_path in dir:
# Applies a Face - Detect for every image in the above directory
image = cv2.imread(str(image_path))
cv2.imshow('Original Image', image)
cv2.waitKey(0)
plate_contour = plate_detection(image)
if plate_contour is None:
print("No license plates detected")
continue
img = image.copy()
try:
plate = crop_contour(plate_contour, image, img)
except:
print("Plate could not be processed")
continue
cv2.destroyAllWindows()
plate_recognition(plate)
image = plate.convert('L')
#image = transformer(image)
cv2.resize(image, (100, 32))
if torch.cuda.is_available():
image = image.cuda()
image = image.view(1, *image.size())
image = Variable(image)
model.eval()
preds = model(image)
_, preds = preds.max(2)
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('%-20s => %-20s' % (raw_pred, sim_pred))
exit()
def pre_model(nclass, ocrModelPath):
# @@parm nclass:字符总数
# @@预训练模型文件
if torch.cuda.is_available() and opt.ngpu:
model = crnn.CRNN(32, 1, nclass + 1, 256, 1).cuda()
else:
model = crnn.CRNN(32, 1, nclass + 1, 256, 1).cpu()
state_dict = torch.load(ocrModelPath,
map_location=lambda storage, loc: storage)
new_state_dict = OrderedDict()
for k, v in state_dict.items():
name = k.replace('module.', '') # remove `module.`
new_state_dict[name] = v
model.load_state_dict(new_state_dict)
model.eval()
return model
def test():
config = config_args()
transform_val=\
transforms.Compose([
transforms.Resize(size=(config.DATASET.IMAGE_SIZE.H,config.DATASET.IMAGE_SIZE.W)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.5, 0.5, 0.5],std=[0.5, 0.5, 0.5])])
criterion = nn.CTCLoss()
label_tool = LabelTool(char_std_path=config.DATASET.CHAR_FILE)
model = crnn.CRNN(config).to(torch.device("cuda:" + config.CUDNN.GPU))
model.load_state_dict(
torch.load(config.TRAIN.RESUME.MODEL_SAVE)['state_dict'])
images_val, labels_val = label_tool.get_labels(
labels_path=config.DATASET.LABELS_FILE.VAL)
dataset_val = Dataset_OCR(images_root_dir=config.DATASET.IMAGE_ROOT,
images_name=images_val,
transform=transform_val)
dataloader_val = DataLoader(dataset=dataset_val,
batch_size=config.TEST.BATCH,
shuffle=config.TEST.SHUFFLE,
num_workers=config.TEST.WORKERS)
loss_all_val = 0
step_val = 0
nums_all = 0
nums_all_correct = 0
for images_val, indexs_val in dataloader_val:
images_val = images_val.to(torch.device("cuda:" + config.CUDNN.GPU))
with torch.no_grad():
output_val = model(images_val)
sequence_len_val = output_val.shape[0]
target_val, input_lengths_val, target_lengths_val = label_tool.convert_ctcloss_labels(
indexs_val, labels_val, sequence_len_val)
preds_val = output_val.permute(1, 0, 2).argmax(2).cpu().numpy()
preds_str_val, preds_str_val_blank = label_tool.decode_batch(preds_val)
correct_nums = label_tool.cal_correct_nums(preds_str_val,
preds_str_val_blank,
indexs_val, labels_val,
step_val)
nums_all_correct += correct_nums
nums_all += output_val.shape[1]
print('nums_all_correct{},nums_all{}'.format(nums_all_correct,
nums_all))
loss_val = criterion(output_val, target_val, input_lengths_val,
target_lengths_val)
loss_all_val += loss_val
step_val += 1
acc = nums_all_correct / nums_all
print("val_loss_avarage={},val_accuracy={}".format(
loss_all_val / step_val, nums_all_correct / nums_all))
def crnnSource():
# Choose alphabet from arg or key.py(default)
if opt.alphabet is None:
alphabet = keys.alphabet
else:
alphabet = opt.alphabet
converter = utils.strLabelConverter(alphabet)
if torch.cuda.is_available() and opt.gpu:
model = crnn.CRNN(32, 1, len(alphabet) + 1, 256, 1)
model = model.cuda()
else:
model = crnn.CRNN(32, 1, len(alphabet) + 1, 256, 1)
model = model.cpu()
path = opt.model
print('loading pretrained model from %s' % path)
model.eval()
model.load_state_dict(torch.load(path))
return model, converter
def net_init():
nclass = len(params.alphabet) + 1
crnn = net.CRNN(params.imgH, params.nc, nclass, params.nh)
crnn.apply(weights_init)
if params.pretrained != '':
print('loading pretrained model from %s' % params.pretrained)
if params.multi_gpu:
crnn = torch.nn.DataParallel(crnn)
crnn.load_state_dict(torch.load(params.pretrained))
return crnn
def __init__(self, weights, char, height=32, cuda=None):
alphabet = open(char).read().rstrip()
nclass = len(alphabet) + 1
self.height = height
self.device = 'cuda:{}'.format(cuda) if cuda != None else 'cpu'
self.model = crnn.CRNN(32, 3, nclass, 256)
self.model.load_state_dict(
torch.load(weights, map_location=self.device))
if cuda != None:
self.model.cuda(self.device)
self.converter = utils.strLabelConverter(alphabet, ignore_case=False)
self.model.eval()
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
def get_recognition_model():
model_path = './weights/crnn.pth'
#Intial model
model = crnn.CRNN(32, 1, 37, 256)
try:
model = model.cuda()
except:
#there is some issue with model loading,
#and loading it twice solves the problem
model = model.cuda()
print('loading pretrained model from %s' % model_path)
model.load_state_dict(torch.load(model_path))
return model
def load_model(model_path):
# net init
global transformer, model, converter
print('loading pretrained model from %s' % model_path)
nclass = len(params.alphabet) + 1
model = crnn.CRNN(params.imgH, params.nc, nclass, params.nh)
if torch.cuda.is_available():
model = model.cuda()
model = torch.nn.DataParallel(model)
else:
model.load_state_dict(torch.load(model_path, map_location='cpu'))
model.eval()
converter = utils.strLabelConverter(params.alphabet)
transformer = dataset.resizeNormalize((100, 32))
def __init__(self, imgp, labp, bs):
self.crnn_net = None
nclass = len(params.alphabet) + 1
self.crnn_net = crnn.CRNN(32, 1, nclass, 256, n_rnn=2, leakyRelu=False)
crnn_model_path = './crnn.pth'
self.crnn_handle = CRNNHandle.CRNNHandle(crnn_model_path,
self.crnn_net,
gpu_id=1.1)
self.imgp = imgp
self.labp = labp
self.bs = bs
def crnnloader():
#model_path = '/home/ahmed/crnn/data/crnn.pth'
model_path = './CRNN/data/crnn.pth'
#img_path = '/home/ahmed/crnn/data/demo.png'
#img_path='/home/ahmed/Pictures/cogedis/2-total.png'
img_path = './CRNN/data/demo.png'
alphabet = '0123456789abcdefghijklmnopqrstuvwxyz'
model = crnn.CRNN(32, 1, 37, 256, 1)
print('loading pretrained model from %s' % model_path)
model.load_state_dict(torch.load(model_path))
converter = utils.strLabelConverter(alphabet)
return model, converter
def net_init():
nclass = len(params.alphabet) + 1
crnn = net.CRNN(params.imgH, params.nc, nclass, params.nh)
crnn.apply(weights_init)
if params.pretrained != '':
print('loading pretrained model from %s' % params.pretrained)
if params.multi_gpu:
crnn = torch.nn.DataParallel(crnn)
std = torch.load(params.pretrained)
# # Remove the last FC layer
std.popitem(last=True)
std.popitem(last=True)
crnn.load_state_dict(std, strict=False)
return crnn
def load_images_to_predict():
print('load_images_to_predict')
# load model
model_path = './save/netCRNN_225_1250.pth'
alphabet = '0123456789,.:(%$!^&-/);<~|`>?+=_[]{}"\'@#*abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ\ '
imgH = 32 # should be 32
nclass = len(alphabet) + 1
nhiddenstate = 256
model = crnn.CRNN(imgH, 1, nclass, nhiddenstate)
if torch.cuda.is_available():
model = model.cuda()
print('loading pretrained model from %s' % model_path)
model.load_state_dict({
k.replace('module.', ''): v
for k, v in torch.load(model_path, map_location='cpu').items()
})
# load image
main_path = os.path.abspath(os.path.join(os.getcwd(), ".."))
image_path = os.path.join(main_path, 'result/step2/image2/*.jpg')
label_path = os.path.join(main_path, 'result/step2/label/')
image_list = [os.path.splitext(f)[0] for f in glob.glob(image_path)]
jpg_files = [name + ".jpg" for name in image_list]
print('total images: ', len(jpg_files))
count = 1
for jpg in jpg_files:
image = Image.open(jpg).convert('L')
words_list = []
label = label_path + jpg.split('/')[-1][:-3] + 'txt'
txt = sort_lines(label)
for line in txt:
box = line.split(',')
crop_image = image.crop(
(int(box[0]), int(box[1]), int(box[4]), int(box[5])))
words = predict_this_box(crop_image, model, alphabet)
words_list.append(words)
result_path = os.path.join(main_path, 'result/step3/')
save_path = result_path + jpg.split('/')[-1][:-3] + 'txt'
with open(save_path, 'w+') as result:
for line in words_list:
result.writelines(line + '\n')
if count % 1000 == 0:
print(str(count), 'images finished.', 'total images: ',
len(jpg_files))
count += 1
def __init__(self, model_path):
self.model_path = model_path
nclass = len(params.alphabet) + 1
self.model = crnn.CRNN(params.imgH, params.nc, nclass, params.nh)
if torch.cuda.is_available():
self.model = self.model.cuda()
if params.multi_gpu:
self.model = torch.nn.DataParallel(self.model)
self.converter = utils.strLabelConverter(params.alphabet)
self.transformer = dataset.resizeNormalize((100, 32))
print('loading pretrained model from %s' % model_path)
self.model.load_state_dict(torch.load(model_path))
