Update README.md
Browse files
README.md
CHANGED
@@ -15,39 +15,35 @@ megfile==3.1.2
|
|
15 |
|
16 |
|
17 |
```python
|
18 |
-
# test.py
|
19 |
-
import torch
|
20 |
-
from PIL import Image
|
21 |
from transformers import AutoModel, AutoTokenizer
|
22 |
|
23 |
-
|
24 |
-
|
25 |
model = model.eval().cuda()
|
26 |
-
tokenizer = AutoTokenizer.from_pretrained('openbmb/MiniCPM-V-2_6', trust_remote_code=True)
|
27 |
|
28 |
-
image = Image.open('xx.jpg').convert('RGB')
|
29 |
-
question = 'What is in the image?'
|
30 |
-
msgs = [{'role': 'user', 'content': [image, question]}]
|
31 |
|
32 |
-
|
33 |
-
|
34 |
-
|
35 |
-
|
36 |
-
)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
37 |
print(res)
|
38 |
|
39 |
-
|
40 |
-
## the model.chat will return a generator
|
41 |
-
res = model.chat(
|
42 |
-
image=None,
|
43 |
-
msgs=msgs,
|
44 |
-
tokenizer=tokenizer,
|
45 |
-
sampling=True,
|
46 |
-
stream=True
|
47 |
-
)
|
48 |
-
|
49 |
-
generated_text = ""
|
50 |
-
for new_text in res:
|
51 |
-
generated_text += new_text
|
52 |
-
print(new_text, flush=True, end='')
|
53 |
```
|
|
|
15 |
|
16 |
|
17 |
```python
|
|
|
|
|
|
|
18 |
from transformers import AutoModel, AutoTokenizer
|
19 |
|
20 |
+
tokenizer = AutoTokenizer.from_pretrained('ucaslcl/GOT-OCR2_0', trust_remote_code=True)
|
21 |
+
model = AutoModel.from_pretrained('ucaslcl/GOT-OCR2_0', trust_remote_code=True, low_cpu_mem_usage=True, device_map='cuda', use_safetensors=True, pad_token_id=tokenizer.eos_token_id)
|
22 |
model = model.eval().cuda()
|
|
|
23 |
|
|
|
|
|
|
|
24 |
|
25 |
+
# input your test image
|
26 |
+
image_file = 'xxx.jpg'
|
27 |
+
|
28 |
+
# plain texts OCR
|
29 |
+
model.chat(tokenizer, image_file, ocr_type='ocr')
|
30 |
+
|
31 |
+
# format texts OCR:
|
32 |
+
model.chat(tokenizer, image_file, ocr_type='format')
|
33 |
+
|
34 |
+
# fine-grained OCR:
|
35 |
+
model.chat(tokenizer, image_file, ocr_type='ocr', ocr_box='')
|
36 |
+
model.chat(tokenizer, image_file, ocr_type='format', ocr_box='')
|
37 |
+
model.chat(tokenizer, image_file, ocr_type='ocr', ocr_color='')
|
38 |
+
model.chat(tokenizer, image_file, ocr_type='format', ocr_color='')
|
39 |
+
|
40 |
+
# multi-crop OCR:
|
41 |
+
res = model.chat_crop(tokenizer, image_file = image_file)
|
42 |
+
|
43 |
+
# render the formatted OCR results:
|
44 |
+
model.chat(tokenizer, image_file, ocr_type='format', ocr_box='', ocr_color='', render=True, save_render_file = '/data/code/a2hf/chat_plus.html')
|
45 |
+
|
46 |
print(res)
|
47 |
|
48 |
+
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
49 |
```
|