class RecInfer:
def __init__(self, model_path):
ckpt = torch.load(model_path, map_location='cpu')
cfg = ckpt['cfg']
self.model = build_model(cfg['model'])
state_dict = {}
for k, v in ckpt['state_dict'].items():
state_dict[k.replace('module.', '')] = v
self.model.load_state_dict(state_dict)
self.device = torch.device(
'cuda:0' if torch.cuda.is_available() else 'cpu')
self.model.to(self.device)
self.model.eval()
self.process = RecDataProcess(cfg['dataset']['train']['dataset'])
self.converter = CTCLabelConverter(cfg['dataset']['alphabet'])
def predict(self, img):
# 预处理根据训练来
img = self.process.resize_with_specific_height(img)
# img = self.process.width_pad_img(img, 120)
img = self.process.normalize_img(img)
tensor = torch.from_numpy(img.transpose([2, 0, 1])).float()
tensor = tensor.unsqueeze(dim=0)
tensor = tensor.to(self.device)
out = self.model(tensor)
txt = self.converter.decode(out.softmax(dim=2).detach().cpu().numpy())
return txt
class RecInfer:
def __init__(self, model_path, batch_size=16):
ckpt = torch.load(model_path, map_location='cpu')
cfg = ckpt['cfg']
self.model = build_model(cfg['model'])
state_dict = {}
for k, v in ckpt['state_dict'].items():
state_dict[k.replace('module.', '')] = v
self.model.load_state_dict(state_dict)
self.device = torch.device(
'cuda:0' if torch.cuda.is_available() else 'cpu')
self.model.to(self.device)
self.model.eval()
self.process = RecDataProcess(cfg['dataset']['train']['dataset'])
self.converter = CTCLabelConverter(cfg['dataset']['alphabet'])
self.batch_size = batch_size
def predict(self, imgs):
# 预处理根据训练来
if not isinstance(imgs, list):
imgs = [imgs]
imgs = [
self.process.normalize_img(
self.process.resize_with_specific_height(img)) for img in imgs
]
widths = np.array([img.shape[1] for img in imgs])
idxs = np.argsort(widths)
txts = []
for idx in range(0, len(imgs), self.batch_size):
batch_idxs = idxs[idx:min(len(imgs), idx + self.batch_size)]
batch_imgs = [
self.process.width_pad_img(imgs[idx],
imgs[batch_idxs[-1]].shape[1])
for idx in batch_idxs
]
batch_imgs = np.stack(batch_imgs)
tensor = torch.from_numpy(batch_imgs.transpose([0, 3, 1,
2])).float()
tensor = tensor.to(self.device)
with torch.no_grad():
out = self.model(tensor)
out = out.softmax(dim=2)
out = out.cpu().numpy()
txts.extend(
[self.converter.decode(np.expand_dims(txt, 0)) for txt in out])
#按输入图像的顺序排序
idxs = np.argsort(idxs)
out_txts = [txts[idx] for idx in idxs]
return out_txts