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 = CTCAsWholeLabelConverter(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
class RecInfer:
def __init__(self, model_path):
ckpt = torch.load(model_path, map_location='cpu')
# with open('crnn_ckpt.txt',"w") as f:
# for k in ckpt['state_dict'].keys():
# f.write(k)
# f.write("---")
# f.write(str(ckpt['state_dict'][k].shape))
# f.write("\n")
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:1' 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'])
print(cfg['dataset']['alphabet'])
def predict(self, img):
# 预处理根据训练来
print("input of rec_infer:" + str(type(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)
print(out.shape)
txt = self.converter.decode(out.softmax(dim=2).detach().cpu().numpy())
return txt
class RecInfer:
def __init__(self, model_path):
ckpt = torch.load(model_path, map_location='cpu')
cfg = ckpt['cfg']
from config.rec_train_config import config
self.model = build_model(config['model'])
self.model = nn.DataParallel(self.model)
self.model.load_state_dict(ckpt['state_dict'])
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
self.model.to(device)
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)
out = self.model(tensor)
txt = self.converter.decode(out.softmax(dim=2).detach().cpu().numpy())
return txt