def val_on_image(img, model, gpu):
imgH = config.imgH
h, w = img.shape[:2]
imgW = imgH * w // h
transformer = mydataset.resizeNormalize((imgW, imgH), is_test=True)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
image = Image.fromarray(np.uint8(img)).convert('L')
image = transformer(image)
if gpu:
image = image.cuda()
image = image.view(1, *image.size())
image = Variable(image)
model.eval()
preds = model(image)
preds = F.log_softmax(preds, 2)
conf, 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)
return sim_pred
def val_on_image(img, model, gpu):
imgH = config.imgH
imgW = config.imgW
transformer = mydataset.resizeNormalize((imgW, imgH), is_test=True)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
image = Image.fromarray(np.uint8(img)).convert('L')
image = transformer(image)
if gpu:
image = image.cuda()
image = image.view(1, *image.size())
image = Variable(image)
model.eval()
preds = model(image)
_, H, W, _ = preds.size()
T = H * W
preds = preds.view((T, -1))
preds = torch.sum(preds, 0)
preds, idx = torch.max(preds, 0)
preds = alphabet[idx]
# raw_pred = converter.decode( preds.data, preds_size.data, raw=True )
sim_pred = converter.decode(idx, 1, raw=False)
return sim_pred
config.batchSize)
else:
sampler = None
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=config.batchSize,
shuffle=True,
sampler=sampler,
num_workers=int(config.workers),
collate_fn=mydataset.alignCollate(
imgH=config.imgH,
imgW=config.imgW,
keep_ratio=config.keep_ratio))
test_dataset = mydataset.MyDataset(info_filename=config.val_infofile,
transform=mydataset.resizeNormalize(
(config.imgW, config.imgH),
is_test=True))
converter = utils.strLabelConverter(config.alphabet)
# criterion = CTCLoss(reduction='sum',zero_infinity=True)
criterion = CTCLoss()
best_acc = 0.9
# 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)
# cudnn.benchmark = True
train_dataset = mydataset.MyDataset(info_filename=config.train_infofile)
assert train_dataset
if not config.random_sample:
sampler = mydataset.randomSequentialSampler(train_dataset, config.batchSize)
else:
sampler = None
train_loader = torch.utils.data.DataLoader(
train_dataset, batch_size=config.batchSize,
shuffle=True, sampler=sampler,
num_workers=int(config.workers),
collate_fn=mydataset.alignCollate(imgH=config.imgH, imgW=config.imgW, keep_ratio=config.keep_ratio))
test_dataset = mydataset.MyDataset(
info_filename=config.val_infofile, transform=mydataset.resizeNormalize((config.imgW, config.imgH), is_test=True))
converter = utils.strLabelConverter(config.alphabet)
criterion = CTCLoss(reduction='sum',zero_infinity=True)
# criterion = CTCLoss()
best_acc = 0.9
# 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)
alphabet = keys.alphabet
print(len(alphabet))
imgH = config.imgH
imgW = config.imgW
model = crnn.CRNN(imgH, 1, len(alphabet) + 1, 256)
if gpu:
model = model.cuda()
print('loading pretrained model from %s' % model_path)
if gpu:
model.load_state_dict(torch.load(model_path))
else:
model.load_state_dict(
torch.load(model_path, map_location=lambda storage, loc: storage))
converter = utils.strLabelConverter(alphabet)
transformer = mydataset.resizeNormalize((imgW, imgH), is_test=True)
def recognize_downline(img, crnn_model=model):
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
image = Image.fromarray(np.uint8(img)).convert('L')
image = transformer(image)
if gpu:
image = image.cuda()
image = image.view(1, *image.size())
image = Variable(image)
model.eval()
preds = model(image)
preds = F.log_softmax(preds, 2)