import gradio as gr
from PIL import Image
from transformers import AutoTokenizer, AutoProcessor, AutoModelForImageTextToText
import torch
import spaces
model_path = "nanonets/Nanonets-OCR-s"
# Load model once at startup
print("Loading Nanonets OCR model...")
model = AutoModelForImageTextToText.from_pretrained(
model_path,
torch_dtype="auto",
device_map="auto",
attn_implementation="flash_attention_2"
)
model.eval()
tokenizer = AutoTokenizer.from_pretrained(model_path)
processor = AutoProcessor.from_pretrained(model_path)
print("Model loaded successfully!")
@spaces.GPU()
def ocr_image_gradio(image, max_tokens=4096):
"""Process image through Nanonets OCR model for Gradio interface"""
if image is None:
return "Please upload an image."
try:
prompt = """Extract the text from the above document as if you were reading it naturally. Return the tables in html format. Return the equations in LaTeX representation. If there is an image in the document and image caption is not present, add a small description of the image inside the ![]()
tag; otherwise, add the image caption inside ![]()
. Watermarks should be wrapped in brackets. Ex: OFFICIAL COPY. Page numbers should be wrapped in brackets. Ex: 14 or 9/22. Prefer using ☐ and ☑ for check boxes."""
# Convert PIL image if needed
if not isinstance(image, Image.Image):
image = Image.fromarray(image)
messages = [
{"role": "system", "content": "You are a helpful assistant."},
{"role": "user", "content": [
{"type": "image", "image": image},
{"type": "text", "text": prompt},
]},
]
text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
inputs = processor(text=[text], images=[image], padding=True, return_tensors="pt")
inputs = inputs.to(model.device)
with torch.no_grad():
output_ids = model.generate(**inputs, max_new_tokens=max_tokens, do_sample=False)
generated_ids = [output_ids[len(input_ids):] for input_ids, output_ids in zip(inputs.input_ids, output_ids)]
output_text = processor.batch_decode(generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=True)
return output_text[0]
except Exception as e:
return f"Error processing image: {str(e)}"
# Create Gradio interface
with gr.Blocks(title="Nanonets OCR Demo") as demo:
# Replace simple markdown with styled HTML header that includes resources
gr.HTML("""
🔍 Nanonets OCR - Document Text Extraction
A state-of-the-art image-to-markdown OCR model for intelligent document processing
""")
with gr.Row():
with gr.Column(scale=1):
image_input = gr.Image(
label="Upload Document Image",
type="pil",
height=400
)
max_tokens_slider = gr.Slider(
minimum=1024,
maximum=8192,
value=4096,
step=512,
label="Max Tokens",
info="Maximum number of tokens to generate"
)
extract_btn = gr.Button("Extract Text", variant="primary", size="lg")
with gr.Column(scale=2):
output_text = gr.Textbox(
label="Extracted Text",
lines=20,
show_copy_button=True,
placeholder="Extracted text will appear here..."
)
# Event handlers
extract_btn.click(
fn=ocr_image_gradio,
inputs=[image_input, max_tokens_slider],
outputs=output_text,
show_progress=True
)
image_input.change(
fn=ocr_image_gradio,
inputs=[image_input, max_tokens_slider],
outputs=output_text,
show_progress=True
)
# Add model information section
with gr.Accordion("About Nanonets-OCR-s", open=False):
gr.Markdown("""
## Nanonets-OCR-s
Nanonets-OCR-s is a powerful, state-of-the-art image-to-markdown OCR model that goes far beyond traditional text extraction.
It transforms documents into structured markdown with intelligent content recognition and semantic tagging, making it ideal
for downstream processing by Large Language Models (LLMs).
### Key Features
- **LaTeX Equation Recognition**: Automatically converts mathematical equations and formulas into properly formatted LaTeX syntax.
It distinguishes between inline ($...$) and display ($$...$$) equations.
- **Intelligent Image Description**: Describes images within documents using structured `![]()
` tags, making them digestible
for LLM processing. It can describe various image types, including logos, charts, graphs and so on, detailing their content,
style, and context.
- **Signature Detection & Isolation**: Identifies and isolates signatures from other text, outputting them within a `` tag.
This is crucial for processing legal and business documents.
- **Watermark Extraction**: Detects and extracts watermark text from documents, placing it within a `` tag.
- **Smart Checkbox Handling**: Converts form checkboxes and radio buttons into standardized Unicode symbols (☐, ☑, ☒)
for consistent and reliable processing.
- **Complex Table Extraction**: Accurately extracts complex tables from documents and converts them into both markdown
and HTML table formats.
""")
if __name__ == "__main__":
demo.queue().launch()