Update app.py

app.py

@@ -7,20 +7,20 @@ import torch
 import easyocr
 import omegaconf
 
-from vietocr.model.transformerocr import VietOCR
-from vietocr.model.vocab import Vocab
-from vietocr.translate import translate, process_input
 
-config = omegaconf.OmegaConf.load("vgg-seq2seq.yaml")
-config = omegaconf.OmegaConf.to_container(config, resolve=True)
+from vietocr.vietocr.tool.predictor import Predictor
+from vietocr.vietocr.tool.config import Cfg
 
-vocab = Vocab(config['vocab'])
-model = VietOCR(len(vocab),
-        config['backbone'],
-        config['cnn'], 
-        config['transformer'],
-        config['seq_modeling'])
-model.load_state_dict(torch.load('train_old.pth', map_location=torch.device('cpu')))
+# Configure of VietOCR
+config = Cfg.load_config_from_name('vgg_transformer')
+# config = Cfg.load_config_from_file('vietocr/config.yml')
+# config['weights'] = '/Users/bmd1905/Desktop/pretrain_ocr/vi00_vi01_transformer.pth'
+
+config['cnn']['pretrained'] = True
+config['predictor']['beamsearch'] = True
+config['device'] = 'cuda:0' # mps
+
+recognitor = Predictor(config)
 
 torch.hub.download_url_to_file('https://github.com/JaidedAI/EasyOCR/raw/master/examples/english.png', 'english.png')
 torch.hub.download_url_to_file('https://github.com/JaidedAI/EasyOCR/raw/master/examples/thai.jpg', 'thai.jpg')
@@ -51,11 +51,13 @@ def inference(filepath, lang):
         img = Image.open(filepath)
         #img = img[y0:y1, x0:x1]
         width, height =img.size
-        img = img.crop((max(0,x1-5), max(0,y1-5), min(x3+5,width), min(y3+5, height))) # crop the image
-        img = process_input(img, config['dataset']['image_height'], 
-                    config['dataset']['image_min_width'], config['dataset']['image_max_width'])
-        out = translate(img, model)[0].tolist()
-        out = vocab.decode(out)
+        cropped_image = img.crop((max(0,x1-5), max(0,y1-5), min(x3+5,width), min(y3+5, height))) # crop the image
+        try:
+            cropped_image = Image.fromarray(cropped_image)
+        except:
+            continue
+
+        out = recognitor.predict(cropped_image)
         new_bounds.append((bbox,text, out, prob))
     im = PIL.Image.open(filepath)
     draw_boxes(im, bounds)