Spaces:
Running
on
Zero
Running
on
Zero
ABOUT = """ | |
# TB-OCR Unofficial Demo | |
This is an unofficial demo of [yifeihu/TB-OCR-preview-0.1](https://huggingface.co/yifeihu/TB-OCR-preview-0.1). | |
Overview of TB-OCR: | |
> TB-OCR-preview (Text Block OCR), created by [Yifei Hu](https://x.com/hu_yifei), is an end-to-end OCR model handling text, math latex, and markdown formats all at once. The model takes a block of text as the input and returns clean markdown output. Headers are marked with `##`. Math expressions are guaranteed to be wrapped in brackets `\( inline math \) \[ display math \]` for easier parsing. This model does not require line-detection or math formula detection. | |
(From the [model card](https://huggingface.co/yifeihu/TB-OCR-preview-0.1)) | |
""" | |
# check out https://huggingface.co/microsoft/Phi-3.5-vision-instruct for more details | |
import torch | |
from transformers import AutoModelForCausalLM, AutoProcessor | |
from PIL import Image | |
import requests | |
model_id = "yifeihu/TB-OCR-preview-0.1" | |
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
model = AutoModelForCausalLM.from_pretrained( | |
model_id, | |
device_map=DEVICE, | |
trust_remote_code=True, | |
torch_dtype="auto", | |
# _attn_implementation='flash_attention_2', | |
#load_in_4bit=True # Optional: Load model in 4-bit mode to save memory | |
) | |
processor = AutoProcessor.from_pretrained(model_id, | |
trust_remote_code=True, | |
num_crops=16 | |
) | |
def phi_ocr(image_url): | |
question = "Convert the text to markdown format." | |
image = Image.open(image_url) | |
prompt_message = [{ | |
'role': 'user', | |
'content': f'<|image_1|>\n{question}', | |
}] | |
prompt = processor.tokenizer.apply_chat_template(prompt_message, tokenize=False, add_generation_prompt=True) | |
inputs = processor(prompt, [image], return_tensors="pt").to(DEVICE) | |
generation_args = { | |
"max_new_tokens": 1024, | |
"temperature": 0.1, | |
"do_sample": False | |
} | |
generate_ids = model.generate(**inputs, eos_token_id=processor.tokenizer.eos_token_id, **generation_args | |
) | |
generate_ids = generate_ids[:, inputs['input_ids'].shape[1]:] | |
response = processor.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] | |
response = response.split("<image_end>")[0] # remove the image_end token | |
return response | |
import gradio as gr | |
with gr.Blocks() as demo: | |
gr.Markdown(ABOUT) | |
with gr.Row(): | |
with gr.Column(): | |
img = gr.Image(label="Input image", type="filename") | |
btn = gr.Button("OCR") | |
with gr.Column(): | |
out = gr.Markdown() | |
btn.click(phi_ocr, inputs=img, outputs=out) | |
demo.queue().launch() |