File size: 2,326 Bytes
6498fe6 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 | from Vocr.tool.translate import build_model, translate, translate_beam_search, process_input
import torch
from collections import defaultdict
class Predictor:
def __init__(self, config):
device = config['device']
model, vocab = build_model(config)
weights = config['weights']
model.load_state_dict(torch.load(weights, map_location=torch.device(device)))
self.config = config
self.model = model
self.vocab = vocab
self.device = device
def predict(self, img, return_prob=False):
img = process_input(img, self.config['dataset']['image_height'],
self.config['dataset']['image_min_width'], self.config['dataset']['image_max_width'])
img = img.to(self.config['device'])
if self.config['predictor']['beamsearch']:
sent = translate_beam_search(img, self.model)
s = sent
prob = None
else:
s, prob = translate(img, self.model)
s = s[0].tolist()
prob = prob[0]
s = self.vocab.decode(s)
if return_prob:
return s, prob
else:
return s
def predict_batch(self, imgs, return_prob=False):
bucket = defaultdict(list)
bucket_idx = defaultdict(list)
bucket_pred = {}
sents, probs = [0] * len(imgs), [0] * len(imgs)
for i, img in enumerate(imgs):
img = process_input(img, self.config['dataset']['image_height'],
self.config['dataset']['image_min_width'], self.config['dataset']['image_max_width'])
bucket[img.shape[-1]].append(img)
bucket_idx[img.shape[-1]].append(i)
for k, batch in bucket.items():
batch = torch.cat(batch, 0).to(self.device)
s, prob = translate(batch, self.model)
prob = prob.tolist()
s = s.tolist()
s = self.vocab.batch_decode(s)
bucket_pred[k] = (s, prob)
for k in bucket_pred:
idx = bucket_idx[k]
sent, prob = bucket_pred[k]
for i, j in enumerate(idx):
sents[j] = sent[i]
probs[j] = prob[i]
if return_prob:
return sents, probs
else:
return sents
|