Upload 3 files
Browse files- main.py +30 -0
- radio.jpeg +0 -0
- requirements.txt +7 -0
main.py
ADDED
|
@@ -0,0 +1,30 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
from transformers import AutoModel, AutoTokenizer
|
| 2 |
+
|
| 3 |
+
tokenizer = AutoTokenizer.from_pretrained('ucaslcl/GOT-OCR2_0', trust_remote_code=True)
|
| 4 |
+
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)
|
| 5 |
+
model = model.eval().cuda()
|
| 6 |
+
|
| 7 |
+
|
| 8 |
+
# input your test image
|
| 9 |
+
image_file = 'radio.jpeg'
|
| 10 |
+
|
| 11 |
+
# plain texts OCR
|
| 12 |
+
res = model.chat(tokenizer, image_file, ocr_type='ocr')
|
| 13 |
+
|
| 14 |
+
# format texts OCR:
|
| 15 |
+
# res = model.chat(tokenizer, image_file, ocr_type='format')
|
| 16 |
+
|
| 17 |
+
# fine-grained OCR:
|
| 18 |
+
# res = model.chat(tokenizer, image_file, ocr_type='ocr', ocr_box='')
|
| 19 |
+
# res = model.chat(tokenizer, image_file, ocr_type='format', ocr_box='')
|
| 20 |
+
# res = model.chat(tokenizer, image_file, ocr_type='ocr', ocr_color='')
|
| 21 |
+
# res = model.chat(tokenizer, image_file, ocr_type='format', ocr_color='')
|
| 22 |
+
|
| 23 |
+
# multi-crop OCR:
|
| 24 |
+
# res = model.chat_crop(tokenizer, image_file, ocr_type='ocr')
|
| 25 |
+
# res = model.chat_crop(tokenizer, image_file, ocr_type='format')
|
| 26 |
+
|
| 27 |
+
# render the formatted OCR results:
|
| 28 |
+
# res = model.chat(tokenizer, image_file, ocr_type='format', render=True, save_render_file = './demo.html')
|
| 29 |
+
|
| 30 |
+
print(res)
|
radio.jpeg
ADDED
|
requirements.txt
ADDED
|
@@ -0,0 +1,7 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
torch==2.0.1
|
| 2 |
+
torchvision==0.15.2
|
| 3 |
+
transformers==4.37.2
|
| 4 |
+
tiktoken==0.6.0
|
| 5 |
+
verovio==4.3.1
|
| 6 |
+
accelerate==0.28.0
|
| 7 |
+
|