yuekun commited on
Commit
f926029
1 Parent(s): 11a6a89

Update README.md

Browse files

update code example

Files changed (1) hide show
  1. README.md +6 -6
README.md CHANGED
@@ -30,21 +30,21 @@ You can use the raw model for optical character recognition (OCR) on text images
30
  Here is how to use this model in PyTorch:
31
 
32
  ```python
33
- from transformers import MGPSTRProcessor, MGPSTRModel
34
  import requests
35
  from PIL import Image
36
 
37
- processor = MGPSTRProcessor.from_pretrained('alibaba-damo/mgp-str-base')
38
- model = MGPSTRModel.from_pretrained('alibaba-damo/mgp-str-base')
39
 
40
  # load image from the IIIT-5k dataset
41
  url = "https://i.postimg.cc/ZKwLg2Gw/367-14.png"
42
  image = Image.open(requests.get(url, stream=True).raw).convert("RGB")
43
 
44
- pixel_values = processor(image, return_tensors="pt").pixel_values
45
- generated_ids, attens = model(pixel_values)
46
 
47
- generated_text = processor.batch_decode(generated_ids)['generated_text']
48
  ```
49
 
50
  ### BibTeX entry and citation info
 
30
  Here is how to use this model in PyTorch:
31
 
32
  ```python
33
+ from transformers import MgpstrProcessor, MgpstrForSceneTextRecognition
34
  import requests
35
  from PIL import Image
36
 
37
+ processor = MgpstrProcessor.from_pretrained('alibaba-damo/mgp-str-base')
38
+ model = MgpstrForSceneTextRecognition.from_pretrained('alibaba-damo/mgp-str-base')
39
 
40
  # load image from the IIIT-5k dataset
41
  url = "https://i.postimg.cc/ZKwLg2Gw/367-14.png"
42
  image = Image.open(requests.get(url, stream=True).raw).convert("RGB")
43
 
44
+ pixel_values = processor(images=image, return_tensors="pt").pixel_values
45
+ outputs = model(pixel_values)
46
 
47
+ generated_text = processor.batch_decode(outputs.logits)['generated_text']
48
  ```
49
 
50
  ### BibTeX entry and citation info