add batch prediction

vietocr/tool/predictor.py

@@ -2,6 +2,7 @@
 from vietocr.tool.utils import download_weights
 
 import torch
+from collections import defaultdict
 
 class Predictor():
     def __init__(self, config):
@@ -21,7 +22,8 @@ def __init__(self, config):
         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'])        
@@ -42,3 +44,42 @@ def predict(self, img, return_prob=False):
             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
+