hantech commited on
Commit
747c9ad
β€’
1 Parent(s): 6f7cb0d

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +19 -17
app.py CHANGED
@@ -7,20 +7,20 @@ import torch
7
  import easyocr
8
  import omegaconf
9
 
10
- from vietocr.model.transformerocr import VietOCR
11
- from vietocr.model.vocab import Vocab
12
- from vietocr.translate import translate, process_input
13
 
14
- config = omegaconf.OmegaConf.load("vgg-seq2seq.yaml")
15
- config = omegaconf.OmegaConf.to_container(config, resolve=True)
16
 
17
- vocab = Vocab(config['vocab'])
18
- model = VietOCR(len(vocab),
19
- config['backbone'],
20
- config['cnn'],
21
- config['transformer'],
22
- config['seq_modeling'])
23
- model.load_state_dict(torch.load('train_old.pth', map_location=torch.device('cpu')))
 
 
 
24
 
25
  torch.hub.download_url_to_file('https://github.com/JaidedAI/EasyOCR/raw/master/examples/english.png', 'english.png')
26
  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):
51
  img = Image.open(filepath)
52
  #img = img[y0:y1, x0:x1]
53
  width, height =img.size
54
- img = img.crop((max(0,x1-5), max(0,y1-5), min(x3+5,width), min(y3+5, height))) # crop the image
55
- img = process_input(img, config['dataset']['image_height'],
56
- config['dataset']['image_min_width'], config['dataset']['image_max_width'])
57
- out = translate(img, model)[0].tolist()
58
- out = vocab.decode(out)
 
 
59
  new_bounds.append((bbox,text, out, prob))
60
  im = PIL.Image.open(filepath)
61
  draw_boxes(im, bounds)
 
7
  import easyocr
8
  import omegaconf
9
 
 
 
 
10
 
11
+ from vietocr.vietocr.tool.predictor import Predictor
12
+ from vietocr.vietocr.tool.config import Cfg
13
 
14
+ # Configure of VietOCR
15
+ config = Cfg.load_config_from_name('vgg_transformer')
16
+ # config = Cfg.load_config_from_file('vietocr/config.yml')
17
+ # config['weights'] = '/Users/bmd1905/Desktop/pretrain_ocr/vi00_vi01_transformer.pth'
18
+
19
+ config['cnn']['pretrained'] = True
20
+ config['predictor']['beamsearch'] = True
21
+ config['device'] = 'cuda:0' # mps
22
+
23
+ recognitor = Predictor(config)
24
 
25
  torch.hub.download_url_to_file('https://github.com/JaidedAI/EasyOCR/raw/master/examples/english.png', 'english.png')
26
  torch.hub.download_url_to_file('https://github.com/JaidedAI/EasyOCR/raw/master/examples/thai.jpg', 'thai.jpg')
 
51
  img = Image.open(filepath)
52
  #img = img[y0:y1, x0:x1]
53
  width, height =img.size
54
+ cropped_image = img.crop((max(0,x1-5), max(0,y1-5), min(x3+5,width), min(y3+5, height))) # crop the image
55
+ try:
56
+ cropped_image = Image.fromarray(cropped_image)
57
+ except:
58
+ continue
59
+
60
+ out = recognitor.predict(cropped_image)
61
  new_bounds.append((bbox,text, out, prob))
62
  im = PIL.Image.open(filepath)
63
  draw_boxes(im, bounds)