app.py

import gradio as gr
import spaces
import torch
from gradio_pdf import PDF
from pdf2image import convert_from_path
from PIL import Image
from transformers import AutoModelForImageTextToText, AutoProcessor, AutoTokenizer

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."

    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 <img></img> tag; otherwise, add the image caption inside <img></img>. Watermarks should be wrapped in brackets. Ex: <watermark>OFFICIAL COPY</watermark>. Page numbers should be wrapped in brackets. Ex: <page_number>14</page_number> or <page_number>9/22</page_number>. 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,
            repetition_penalty=1.25,
        )
        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]


@spaces.GPU()
def ocr_pdf_gradio(pdf_path, max_tokens=4096, progress=gr.Progress()):
    """Process each page of a PDF through Nanonets OCR model"""
    if pdf_path is None:
        return "Please upload a PDF file."

    # Convert PDF to images
    progress(0, desc="Converting PDF to images...")
    pdf_images = convert_from_path(pdf_path)

    # Process each page
    all_text = []
    total_pages = len(pdf_images)

    for i, image in enumerate(pdf_images):
        progress(
            (i + 1) / total_pages, desc=f"Processing page {i + 1}/{total_pages}..."
        )
        page_text = ocr_image_gradio(image, max_tokens)
        all_text.append(f"--- PAGE {i + 1} ---\n{page_text}\n")

    # Combine results
    combined_text = "\n".join(all_text)
    return combined_text


# Create Gradio interface
with gr.Blocks(title="Nanonets OCR Demo") as demo:
    # Replace simple markdown with styled HTML header that includes resources
    gr.HTML("""
    <div class="title" style="text-align: center">
        <h1>🔍 Nanonets OCR - Document Text Extraction</h1>
        <p style="font-size: 1.1em; color: #6b7280; margin-bottom: 0.6em;">
            A state-of-the-art image-to-markdown OCR model for intelligent document processing
        </p>
        <div style="display: flex; justify-content: center; gap: 20px; margin: 15px 0;">
            <a href="https://huggingface.co/nanonets/Nanonets-OCR-s" target="_blank" style="text-decoration: none; color: #2563eb; font-weight: 500;">
                📚 Hugging Face Model
            </a>
            <a href="https://nanonets.com/research/nanonets-ocr-s/" target="_blank" style="text-decoration: none; color: #2563eb; font-weight: 500;">
                📝 Release Blog
            </a>
            <a href="https://github.com/NanoNets/docext" target="_blank" style="text-decoration: none; color: #2563eb; font-weight: 500;">
                💻 GitHub Repository
            </a>
        </div>
    </div>
    """)

    with gr.Tabs() as tabs:
        # Image tab
        with gr.TabItem("Image OCR"):
            with gr.Row():
                with gr.Column(scale=1):
                    image_input = gr.Image(
                        label="Upload Document Image", type="pil", height=400
                    )
                    image_max_tokens = gr.Slider(
                        minimum=1024,
                        maximum=8192,
                        value=4096,
                        step=512,
                        label="Max Tokens",
                        info="Maximum number of tokens to generate",
                    )
                    image_extract_btn = gr.Button(
                        "Extract Text", variant="primary", size="lg"
                    )

                with gr.Column(scale=2):
                    image_output_text = gr.Textbox(
                        label="Extracted Text",
                        lines=20,
                        show_copy_button=True,
                        placeholder="Extracted text will appear here...",
                    )

        # PDF tab
        with gr.TabItem("PDF OCR"):
            with gr.Row():
                with gr.Column(scale=1):
                    pdf_input = PDF(label="Upload PDF Document", height=400)
                    pdf_max_tokens = gr.Slider(
                        minimum=1024,
                        maximum=8192,
                        value=4096,
                        step=512,
                        label="Max Tokens per Page",
                        info="Maximum number of tokens to generate for each page",
                    )
                    pdf_extract_btn = gr.Button(
                        "Extract PDF Text", variant="primary", size="lg"
                    )

                with gr.Column(scale=2):
                    pdf_output_text = gr.Textbox(
                        label="Extracted Text (All Pages)",
                        lines=20,
                        show_copy_button=True,
                        placeholder="Extracted text will appear here...",
                    )

    # Event handlers for Image tab
    image_extract_btn.click(
        fn=ocr_image_gradio,
        inputs=[image_input, image_max_tokens],
        outputs=image_output_text,
        show_progress=True,
    )

    image_input.change(
        fn=ocr_image_gradio,
        inputs=[image_input, image_max_tokens],
        outputs=image_output_text,
        show_progress=True,
    )

    # Event handlers for PDF tab
    pdf_extract_btn.click(
        fn=ocr_pdf_gradio,
        inputs=[pdf_input, pdf_max_tokens],
        outputs=pdf_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 `<img>` 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 `<signature>` tag. 
          This is crucial for processing legal and business documents.
        
        - **Watermark Extraction**: Detects and extracts watermark text from documents, placing it within a `<watermark>` 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()