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')