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()
if torch.cuda.is_available() and not opt.cuda:
print(
"WARNING: You have a CUDA device, so you should probably run with --cuda"
)
loader = DatasetLoader(opt.root, opt.train, opt.val, opt.imgW, opt.imgH)
train_loader = loader.train_loader(opt.batch_size, num_workers=opt.workers)
test_loader = loader.test_loader(opt.batch_size, num_workers=opt.workers)
alphabet = open(os.path.join(opt.root, opt.alphabet)).read().rstrip()
nclass = len(alphabet) + 1
nc = 3
print(len(alphabet), alphabet)
converter = utils.strLabelConverter(alphabet, ignore_case=False)
criterion = CTCLoss()
crnn = crnn.CRNN(opt.imgH, nc, nclass, opt.nh)
if opt.pretrained != '':
print('loading pretrained model from %s' % opt.pretrained)
pretrain = torch.load(opt.pretrained)
crnn.load_state_dict(pretrain, strict=False)
image = torch.FloatTensor(opt.batch_size, 3, opt.imgH, opt.imgH)
text = torch.IntTensor(opt.batch_size * 5)
length = torch.IntTensor(opt.batch_size)
if opt.cuda:
crnn.cuda()
image = image.cuda()
def predict(dir, batch_sz, max_iter=10000):
print('Init CRNN classifier')
image = torch.FloatTensor(batch_sz, 3, imgH, imgH)
model = crnn.CRNN64(imgH, nc, nclass, 256)
#model = crnn128.CRNN128(imgH, nc, nclass, 256)
if gpu != None:
print('Use GPU', gpu)
os.environ['CUDA_VISIBLE_DEVICES'] = gpu
model = model.cuda()
image = image.cuda()
print('loading pretrained model from %s' % pretrained)
model.load_state_dict(torch.load(pretrained, map_location='cpu'))
converter = strLabelConverter(alphabet, ignore_case=False)
val_dataset = ImageFileLoader(dir,
flist=test_list,
label=label,
transform=transform_test)
num_files = len(val_dataset)
val_loader = torch.utils.data.DataLoader(val_dataset,
batch_size=batch_sz,
num_workers=workers,
shuffle=False)
image = Variable(image)
# for p in crnn.parameters():
# p.requires_grad = False
model.eval()
print('Start predict in folder', img_dir)
val_iter = iter(val_loader)
max_iter = min(max_iter, len(val_loader))
print('Number of samples', num_files)
begin = time.time()
with torch.no_grad():
for i in range(max_iter):
data = val_iter.next()
cpu_images, cpu_texts, img_paths = data
batch_size = cpu_images.size(0)
utils.loadData(image, cpu_images)
preds = model(image)
preds = preds.squeeze(1)
raw_pred, sim_pred, sent_prob = decode(converter, preds)
# preds_size = Variable(torch.IntTensor([preds.size(0)] * batch_size))
# _, preds = preds.max(2)
# preds = preds.transpose(1, 0).contiguous().view(-1)
# sim_pred = converter.decode(preds.data, preds_size.data, raw=False)
# #print(sim_pred)
# raw_pred = converter.decode(preds.data, preds_size.data, raw=True)
#print(cpu_texts[0])
if debug:
print('\n', raw_pred)
print('\n', round(sent_prob, 3), sim_pred, img_paths)
inv_tensor = inv_normalize(cpu_images[0])
cv_img = inv_tensor.permute(1, 2, 0).numpy()
cv_img_convert = cv2.cvtColor(cv_img, cv2.COLOR_BGR2RGB)
cv2.imshow('image data', cv_img_convert)
ch = cv2.waitKey(0)
if ch == 27:
break
end = time.time()
processing_time = end - begin
print('Processing time:', processing_time)
print('Speed:', num_files / processing_time, 'fps')
