Update app.py

app.py

@@ -1,259 +1,229 @@
-import spaces
-import json
-import math
-import os
-import traceback
-from io import BytesIO
-from typing import Any, Dict, List, Optional, Tuple
-import re
-
-import fitz  # PyMuPDF
 import gradio as gr
-import requests
+import spaces
 import torch
-from huggingface_hub import snapshot_download
-from PIL import Image, ImageDraw, ImageFont
-from qwen_vl_utils import process_vision_info
-from transformers import AutoModelForCausalLM, AutoProcessor
-
-# Constants
-MIN_PIXELS = 3136
-MAX_PIXELS = 11289600
-IMAGE_FACTOR = 28
-
-# Prompts
-prompt = """Please output the layout information from the PDF image, including each layout element's bbox, its category, and the corresponding text content within the bbox.
-
-1. Bbox format: [x1, y1, x2, y2]
-2. Layout Categories: ['Caption', 'Footnote', 'Formula', 'List-item', 'Page-footer', 'Page-header', 'Picture', 'Section-header', 'Table', 'Text', 'Title'].
-3. Text Extraction & Formatting Rules:
-    - Picture: Omit text.
-    - Formula: Format as LaTeX.
-    - Table: Format as HTML.
-    - Others: Format as Markdown.
-4. Output must be the original text with no translation, sorted in human reading order.
-5. Final output: single JSON object.
-"""
-
-# Utility functions
-def round_by_factor(number: int, factor: int) -> int:
-    return round(number / factor) * factor
-
-def smart_resize(height: int, width: int, factor: int = 28,
-                 min_pixels: int = 3136, max_pixels: int = 11289600):
-    if max(height, width) / min(height, width) > 200:
-        raise ValueError("absolute aspect ratio must be smaller than 200")
-    h_bar = max(factor, round_by_factor(height, factor))
-    w_bar = max(factor, round_by_factor(width, factor))
-    if h_bar * w_bar > max_pixels:
-        beta = math.sqrt((height * width) / max_pixels)
-        h_bar = round_by_factor(height / beta, factor)
-        w_bar = round_by_factor(width / beta, factor)
-    elif h_bar * w_bar < min_pixels:
-        beta = math.sqrt(min_pixels / (height * width))
-        h_bar = round_by_factor(height * beta, factor)
-        w_bar = round_by_factor(width * beta, factor)
-    return h_bar, w_bar
+from gradio_pdf import PDF
+from pdf2image import convert_from_path
+from PIL import Image
+from transformers import AutoModelForImageTextToText, AutoProcessor, AutoTokenizer
 
-def fetch_image(image_input, min_pixels=None, max_pixels=None):
-    if isinstance(image_input, str):
-        if image_input.startswith(("http://", "https://")):
-            response = requests.get(image_input)
-            image = Image.open(BytesIO(response.content)).convert('RGB')
-        else:
-            image = Image.open(image_input).convert('RGB')
-    elif isinstance(image_input, Image.Image):
-        image = image_input.convert('RGB')
-    else:
-        raise ValueError(f"Invalid image input type: {type(image_input)}")
-    if min_pixels is not None or max_pixels is not None:
-        min_pixels = min_pixels or MIN_PIXELS
-        max_pixels = max_pixels or MAX_PIXELS
-        height, width = smart_resize(image.height, image.width, factor=IMAGE_FACTOR,
-                                     min_pixels=min_pixels, max_pixels=max_pixels)
-        image = image.resize((width, height), Image.LANCZOS)
-    return image
+model_path = "nanonets/Nanonets-OCR-s"
 
-def load_images_from_pdf(pdf_path: str) -> List[Image.Image]:
-    images = []
-    try:
-        pdf_document = fitz.open(pdf_path)
-        for page_num in range(len(pdf_document)):
-            page = pdf_document.load_page(page_num)
-            mat = fitz.Matrix(2.0, 2.0)
-            pix = page.get_pixmap(matrix=mat)
-            img_data = pix.tobytes("ppm")
-            image = Image.open(BytesIO(img_data)).convert('RGB')
-            images.append(image)
-        pdf_document.close()
-    except Exception as e:
-        print(f"Error loading PDF: {e}")
-    return images
+# 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()
 
-def is_arabic_text(text: str) -> bool:
-    if not text:
-        return False
-    header_pattern = r'^#{1,6}\s+(.+)$'
-    paragraph_pattern = r'^(?!#{1,6}\s|!\[|```|\||\s*[-*+]\s|\s*\d+\.\s)(.+)$'
-    content_text = []
-    for line in text.split('\n'):
-        line = line.strip()
-        if not line:
-            continue
-        header_match = re.match(header_pattern, line, re.MULTILINE)
-        if header_match:
-            content_text.append(header_match.group(1))
-            continue
-        if re.match(paragraph_pattern, line, re.MULTILINE):
-            content_text.append(line)
-    if not content_text:
-        return False
-    combined_text = ' '.join(content_text)
-    arabic_chars = sum(1 for c in combined_text if '\u0600' <= c <= '\u06FF' or '\u0750' <= c <= '\u077F' or '\u08A0' <= c <= '\u08FF')
-    total_chars = sum(1 for c in combined_text if c.isalpha())
-    return total_chars > 0 and (arabic_chars / total_chars) > 0.5
+tokenizer = AutoTokenizer.from_pretrained(model_path)
+processor = AutoProcessor.from_pretrained(model_path)
+print("Model loaded successfully!")
 
-def layoutjson2md(image: Image.Image, layout_data: List[Dict], text_key='text') -> str:
-    import base64
-    markdown_lines = []
-    try:
-        sorted_items = sorted(layout_data, key=lambda x: (x.get('bbox', [0, 0, 0, 0])[1], x.get('bbox', [0, 0, 0, 0])[0]))
-        for item in sorted_items:
-            category = item.get('category', '')
-            text = item.get(text_key, '')
-            if category == 'Picture':
-                markdown_lines.append("![Image](Image detected)\n")
-            elif not text:
-                continue
-            elif category == 'Title':
-                markdown_lines.append(f"# {text}\n")
-            elif category == 'Section-header':
-                markdown_lines.append(f"## {text}\n")
-            elif category == 'Text':
-                markdown_lines.append(f"{text}\n")
-            elif category == 'List-item':
-                markdown_lines.append(f"- {text}\n")
-            elif category == 'Table':
-                markdown_lines.append(f"{text}\n")
-            elif category == 'Formula':
-                markdown_lines.append(f"$$\n{text}\n$$\n")
-            elif category == 'Caption':
-                markdown_lines.append(f"*{text}*\n")
-            elif category == 'Footnote':
-                markdown_lines.append(f"^{text}^\n")
-    except Exception as e:
-        print(f"Error converting to markdown: {e}")
-        return str(layout_data)
-    return "\n".join(markdown_lines)
-
-# Model
-model_id = "rednote-hilab/dots.ocr"
-model_path = "./models/dots-ocr-local"
-snapshot_download(repo_id=model_id, local_dir=model_path, local_dir_use_symlinks=False)
-model = AutoModelForCausalLM.from_pretrained(model_path, attn_implementation="flash_attention_2",
-                                             torch_dtype=torch.bfloat16, device_map="auto", trust_remote_code=True)
-processor = AutoProcessor.from_pretrained(model_path, trust_remote_code=True)
-device = "cuda" if torch.cuda.is_available() else "cpu"
-
-# State
-pdf_cache = {"images": [], "current_page": 0, "total_pages": 0, "file_type": None, "is_parsed": False, "results": []}
 
 @spaces.GPU()
-def inference(image: Image.Image, prompt: str, max_new_tokens=24000) -> str:
-    messages = [{"role": "user", "content": [{"type": "image", "image": image}, {"type": "text", "text": prompt}]}]
-    text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-    image_inputs, video_inputs = process_vision_info(messages)
-    inputs = processor(text=[text], images=image_inputs, videos=video_inputs, padding=True, return_tensors="pt").to(device)
-    with torch.no_grad():
-        generated_ids = model.generate(**inputs, max_new_tokens=max_new_tokens, do_sample=False, temperature=0.1)
-    generated_ids_trimmed = [out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)]
-    output_text = processor.batch_decode(generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False)
-    return output_text[0] if output_text else ""
-
-def process_image(image: Image.Image, min_pixels=None, max_pixels=None):
-    if min_pixels is not None or max_pixels is not None:
-        image = fetch_image(image, min_pixels=min_pixels, max_pixels=max_pixels)
-    raw_output = inference(image, prompt)
-    try:
-        layout_data = json.loads(raw_output)
-        return layoutjson2md(image, layout_data), layout_data
-    except json.JSONDecodeError:
-        return raw_output, None
+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)
 
-def load_file_for_preview(file_path: str):
-    global pdf_cache
-    if not file_path or not os.path.exists(file_path):
-        return None, "No file selected"
-    ext = os.path.splitext(file_path)[1].lower()
-    if ext == '.pdf':
-        images = load_images_from_pdf(file_path)
-        pdf_cache.update({"images": images, "current_page": 0, "total_pages": len(images),
-                          "file_type": "pdf", "is_parsed": False, "results": []})
-        return images[0], f"Page 1 / {len(images)}"
-    else:
-        img = Image.open(file_path).convert('RGB')
-        pdf_cache.update({"images": [img], "current_page": 0, "total_pages": 1,
-                          "file_type": "image", "is_parsed": False, "results": []})
-        return img, "Page 1 / 1"
-
-def turn_page(direction: str):
-    global pdf_cache
-    if not pdf_cache["images"]:
-        return None, '<div class="page-info">No file loaded</div>', "No results yet"
-    if direction == "prev":
-        pdf_cache["current_page"] = max(0, pdf_cache["current_page"] - 1)
-    elif direction == "next":
-        pdf_cache["current_page"] = min(pdf_cache["total_pages"] - 1, pdf_cache["current_page"] + 1)
-    idx = pdf_cache["current_page"]
-    img = pdf_cache["images"][idx]
-    page_info_html = f'<div class="page-info">Page {idx + 1} / {pdf_cache["total_pages"]}</div>'
-    markdown_content = "Page not processed yet"
-    if pdf_cache["is_parsed"] and idx < len(pdf_cache["results"]):
-        markdown_content = pdf_cache["results"][idx]
-    if is_arabic_text(markdown_content):
-        markdown_content = gr.update(value=markdown_content, rtl=True)
-    return img, page_info_html, markdown_content
+    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]
 
-def create_gradio_interface():
-    css = ".page-info {text-align: center;padding: 8px 16px;border-radius: 20px;font-weight: bold;margin: 10px 0;}"
-    with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
-        gr.HTML("<h1 style='text-align:center'>🔍 Dot-OCR - Extracted Content Only</h1>")
-        with gr.Row():
-            with gr.Column(scale=1):
-                file_input = gr.File(label="Upload Image or PDF", file_types=[".jpg", ".jpeg", ".png", ".bmp", ".tiff", ".pdf"], type="filepath")
-                image_preview = gr.Image(label="Preview", type="pil", interactive=False, height=300)
-                with gr.Row():
-                    prev_page_btn = gr.Button("◀ Previous")
-                    page_info = gr.HTML('<div class="page-info">No file loaded</div>')
-                    next_page_btn = gr.Button("Next ▶")
-                process_btn = gr.Button("🚀 Process Document", variant="primary")
-                clear_btn = gr.Button("🗑️ Clear All", variant="secondary")
-            with gr.Column(scale=2):
-                markdown_output = gr.Markdown(value="Click 'Process Document' to see extracted content...", height=500)
-        file_input.change(load_file_for_preview, inputs=file_input, outputs=[image_preview, page_info])
-        prev_page_btn.click(lambda: turn_page("prev"), outputs=[image_preview, page_info, markdown_output])
-        next_page_btn.click(lambda: turn_page("next"), outputs=[image_preview, page_info, markdown_output])
-        process_btn.click(lambda f: _process_document(f), inputs=file_input, outputs=[markdown_output])
-        clear_btn.click(lambda: (None, None, '<div class="page-info">No file loaded</div>', "Click 'Process Document' to see extracted content..."),
-                        outputs=[file_input, image_preview, page_info, markdown_output])
-    return demo
 
-def _process_document(file_path):
-    global pdf_cache
-    if not file_path:
-        return "Please upload a file first."
-    img, _ = load_file_for_preview(file_path)
-    results = []
-    for page_img in pdf_cache["images"]:
-        md_content, _ = process_image(page_img)
-        results.append(md_content)
-    pdf_cache["results"] = results
-    pdf_cache["is_parsed"] = True
-    combined_md = "\n\n---\n\n".join(results)
-    if is_arabic_text(combined_md):
-        return gr.update(value=combined_md, rtl=True)
-    return combined_md
+@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 = create_gradio_interface()
-    demo.queue(max_size=10).launch(server_name="0.0.0.0", server_port=7860)
+    demo.queue().launch(ssr_mode=False)