Spaces:
Running
on
Zero
Running
on
Zero
| import gradio as gr | |
| from PIL import Image | |
| from transformers import AutoModelForImageTextToText, AutoProcessor, AutoTokenizer, TextIteratorStreamer | |
| import spaces | |
| from threading import Thread | |
| from pdf2image import convert_from_path | |
| import os | |
| import tempfile | |
| import base64 | |
| from io import BytesIO | |
| import time | |
| # --- Model Loading --- | |
| # Load the model, processor, and tokenizer once when the app starts. | |
| model_path = "nanonets/Nanonets-OCR-s" | |
| print("Loading Nanonets OCR model...") | |
| model = AutoModelForImageTextToText.from_pretrained( | |
| model_path, | |
| torch_dtype="auto", | |
| device_map="auto", | |
| attn_implementation="flash_attention_2" | |
| ) | |
| model.eval() | |
| processor = AutoProcessor.from_pretrained(model_path) | |
| tokenizer = AutoTokenizer.from_pretrained(model_path) | |
| print("Model loaded successfully!") | |
| # --- Helper Functions --- | |
| def process_tags(content: str) -> str: | |
| """ | |
| Replaces special tags with HTML entities to prevent them from being rendered as HTML. | |
| This function escapes special tags used by the OCR model (like <img>, <watermark>, | |
| <page_number>, <signature>) so they display as text rather than being interpreted | |
| as HTML markup in the Gradio interface. | |
| Args: | |
| content (str): The text content containing special tags to be escaped | |
| Returns: | |
| str: The content with special tags replaced by HTML entities | |
| """ | |
| content = content.replace("<img>", "<img>") | |
| content = content.replace("</img>", "</img>") | |
| content = content.replace("<watermark>", "<watermark>") | |
| content = content.replace("</watermark>", "</watermark>") | |
| content = content.replace("<page_number>", "<page_number>") | |
| content = content.replace("</page_number>", "</page_number>") | |
| content = content.replace("<signature>", "<signature>") | |
| content = content.replace("</signature>", "</signature>") | |
| return content | |
| def encode_image(image: Image) -> str: | |
| """ | |
| Encodes an image to a base64 string for transmission in JSON messages. | |
| Converts a PIL Image object to a base64-encoded JPEG string that can be | |
| embedded in JSON messages for the OCR model API. | |
| Args: | |
| image (Image): PIL Image object to encode | |
| Returns: | |
| str: Base64-encoded JPEG string of the image | |
| """ | |
| buffered = BytesIO() | |
| image.save(buffered, format="JPEG") | |
| img_str = base64.b64encode(buffered.getvalue()).decode("utf-8") | |
| return img_str | |
| def stream_request( | |
| messages: list[dict], | |
| model_name: str, | |
| max_tokens: int = 8000, | |
| temperature: float = 0.0, | |
| ): | |
| """ | |
| Stream text generation from the OCR model given messages with images and text. | |
| This function processes a list of messages containing images and text prompts, | |
| formats them for the Nanonets-OCR-s model, and yields generated text chunks | |
| in a streaming fashion. It handles base64-encoded images and applies the | |
| appropriate chat template for the model. | |
| Args: | |
| messages (list[dict]): List of message dictionaries with role and content. | |
| Content should contain image_url and text items. | |
| model_name (str): Name of the model (unused but kept for compatibility) | |
| max_tokens (int, optional): Maximum number of tokens to generate. Defaults to 8000. | |
| temperature (float, optional): Temperature for generation (unused, model runs deterministically). Defaults to 0.0. | |
| Yields: | |
| str: Generated text chunks as they are produced by the model | |
| """ | |
| # Extract the image and text from messages | |
| for message in messages: | |
| if message["role"] == "user": | |
| content = message["content"] | |
| image_data = None | |
| text_prompt = "" | |
| for item in content: | |
| if item["type"] == "image_url": | |
| # Decode base64 image | |
| image_url = item["image_url"]["url"] | |
| if image_url.startswith("data:image/jpeg;base64,"): | |
| image_base64 = image_url.split(",")[1] | |
| image_bytes = base64.b64decode(image_base64) | |
| image_data = Image.open(BytesIO(image_bytes)) | |
| elif item["type"] == "text": | |
| text_prompt = item["text"] | |
| if image_data is not None: | |
| # Format messages in the expected format for the model | |
| formatted_messages = [ | |
| {"role": "system", "content": "You are a helpful assistant."}, | |
| {"role": "user", "content": [ | |
| {"type": "image", "image": image_data}, | |
| {"type": "text", "text": text_prompt}, | |
| ]}, | |
| ] | |
| # Apply chat template to format the input properly | |
| text = processor.apply_chat_template( | |
| formatted_messages, | |
| tokenize=False, | |
| add_generation_prompt=True | |
| ) | |
| # Process the formatted text and image | |
| inputs = processor( | |
| text=[text], | |
| images=[image_data], | |
| padding=True, | |
| return_tensors="pt" | |
| ) | |
| # Move inputs to the same device as the model | |
| inputs = {k: v.to(model.device) if hasattr(v, 'to') else v for k, v in inputs.items()} | |
| # Set up streaming | |
| streamer = TextIteratorStreamer( | |
| tokenizer, | |
| timeout=60.0, | |
| skip_prompt=True, | |
| skip_special_tokens=True | |
| ) | |
| generation_kwargs = { | |
| **inputs, | |
| "streamer": streamer, | |
| "max_new_tokens": max_tokens, | |
| "do_sample": False, # Deterministic generation | |
| "pad_token_id": tokenizer.eos_token_id, | |
| } | |
| # Start generation in a separate thread | |
| thread = Thread(target=model.generate, kwargs=generation_kwargs) | |
| thread.start() | |
| # Yield generated tokens as they come | |
| for new_text in streamer: | |
| yield new_text | |
| thread.join() | |
| return | |
| # If no valid image/text pair found, return empty | |
| yield "" | |
| def convert_to_markdown_stream( | |
| images: Image, model_name, max_gen_tokens, with_img_desc: bool = False | |
| ): | |
| """ | |
| Generator function that yields streaming markdown conversion results. | |
| Processes images one by one and concatenates results, providing real-time | |
| feedback as the OCR model processes each page. The function handles both | |
| single images and multiple page documents, applying appropriate prompts | |
| for markdown conversion with optional image descriptions. | |
| Args: | |
| images (Image): PIL Image object or list of images to process | |
| model_name (str): Name of the model to use for OCR processing | |
| max_gen_tokens (int): Maximum number of tokens to generate per page | |
| with_img_desc (bool, optional): Whether to include image descriptions in output. Defaults to False. | |
| Yields: | |
| str: Streaming markdown content as it's generated, including page numbers | |
| and accumulated content from all processed pages | |
| """ | |
| images = [images] | |
| # validate_file_paths(file_paths) | |
| # file_paths = convert_files_to_images(file_paths) | |
| # resize_images(file_paths, max_img_size) | |
| # Create system prompt for PDF to markdown conversion | |
| if with_img_desc: | |
| user_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.""" | |
| else: | |
| user_prompt = """Extract the text from the above document as if you were reading it naturally. Return the tables in html format. 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.""" | |
| # Accumulate results from all pages | |
| full_markdown_content = "" | |
| # Process each image individually | |
| for i, image in enumerate(images): | |
| # Build messages for this single image | |
| content = [ | |
| { | |
| "type": "image_url", | |
| "image_url": { | |
| "url": f"data:image/jpeg;base64,{encode_image(image)}" | |
| }, | |
| }, | |
| {"type": "text", "text": user_prompt}, | |
| ] | |
| messages = [{"role": "user", "content": content}] | |
| # Stream this individual page | |
| page_content = "" | |
| try: | |
| for chunk in stream_request( | |
| messages=messages, | |
| model_name=model_name, | |
| max_tokens=max_gen_tokens, | |
| ): | |
| page_content += chunk | |
| # Yield accumulated content from all pages processed so far + current page | |
| current_total = ( | |
| full_markdown_content | |
| + f"Page {i + 1} of {len(images)}\n" | |
| + page_content | |
| ) | |
| time.sleep(0.05) | |
| yield current_total | |
| # Process the completed page content and add it to the full content | |
| full_markdown_content += ( | |
| f"Page {i + 1} of {len(images)}\n" + page_content | |
| ) | |
| except Exception as e: | |
| return f"Error: {e}" | |
| def process_document(image, max_tokens, with_img_desc: bool = False): | |
| """ | |
| Process uploaded document (PDF or image) and convert to markdown. | |
| This is the main entry point function for the Gradio interface. It handles | |
| the uploaded image, resizes it to the appropriate dimensions for the model, | |
| and initiates the markdown conversion process. The function supports both | |
| single images and multi-page documents, with error handling for various | |
| file formats and processing issues. | |
| Args: | |
| image: Uploaded image from Gradio interface (numpy array or PIL Image) | |
| max_tokens (int): Maximum tokens to generate per page for OCR processing | |
| with_img_desc (bool, optional): Whether to include image descriptions in the output. Defaults to False. | |
| Returns: | |
| Generator: Yields markdown content as it's processed, with special tags | |
| escaped for proper display in the Gradio interface | |
| """ | |
| if image is None: | |
| return "Please upload a file first." | |
| try: | |
| # Handle PDF files | |
| # if file_path.name.lower().endswith('.pdf'): | |
| # # Convert PDF to images | |
| # with tempfile.TemporaryDirectory() as temp_dir: | |
| # # Copy uploaded file to temp directory | |
| # temp_pdf_path = os.path.join(temp_dir, "document.pdf") | |
| # import shutil | |
| # shutil.copy(file_path.name, temp_pdf_path) | |
| # # Convert PDF pages to images | |
| # images = convert_from_path(temp_pdf_path, dpi=150) | |
| # images = [image.convert("RGB") for image in images] | |
| # images = [image.resize((2048, 2048)) for image in images] | |
| # # Process each page | |
| # for result in convert_to_markdown_stream( | |
| # images, "nanonets/Nanonets-OCR-s", max_tokens, with_img_desc | |
| # ): | |
| # yield process_tags(result) | |
| # # Handle image files | |
| # else: | |
| # # Open image directly | |
| # image = Image.open(file_path.name).convert("RGB") | |
| # image = image.resize((2048, 2048)) | |
| image = Image.fromarray(image) | |
| image = image.resize((2048, 2048)) | |
| # Process single image | |
| for result in convert_to_markdown_stream( | |
| image, "nanonets/Nanonets-OCR-s", max_tokens, with_img_desc | |
| ): | |
| yield process_tags(result) | |
| except Exception as e: | |
| yield f"Error processing document: {str(e)}" | |
| # --- Gradio Interface --- | |
| title = """# 🙋🏻♂️Welcome to 🌟Tonic's📄 Nanonets-OCR-s: Advanced Document Intelligence Platform | |
| --- | |
| """ | |
| description = """ | |
| The **Nanonets-OCR-s Document Intelligence Platform** is a state-of-the-art AI-powered system designed to transform documents into structured, searchable content with **intelligent semantic understanding**. Built on the foundation of **Amazon's advanced OCR technology**, this platform excels in extracting text, tables, equations, and visual elements from complex documents with unprecedented accuracy. | |
| ### Key Features | |
| - **Multi-Format Support**: PDF, Images (JPEG, PNG, TIFF), Scanned Documents | |
| - **Intelligent Content Recognition**: Tables, Equations, Signatures, Watermarks, Checkboxes | |
| - **Advanced Semantic Understanding**: Context-aware text extraction and formatting | |
| - **Real-Time Processing**: Streaming results with live progress updates | |
| - **Enhanced Output Formats**: Markdown, HTML, LaTeX, Structured JSON | |
| - **Batch Processing**: Handle multiple documents simultaneously | |
| - **Quality Assurance**: Built-in validation and error correction | |
| ## Supported Document Types | |
| - **Business Documents**: Invoices, Receipts, Contracts, Reports | |
| - **Academic Papers**: Research Papers, Theses, Technical Documents | |
| - **Financial Documents**: Bank Statements, Tax Forms, Financial Reports | |
| - **Legal Documents**: Contracts, Legal Forms, Court Documents | |
| - **Medical Documents**: Patient Records, Medical Forms, Prescriptions | |
| - **Government Documents**: Forms, Certificates, Official Records | |
| """ | |
| model_info = """ | |
| ## How to Use | |
| 1. **Upload Document**: Drag and drop or select your PDF/image file | |
| 2. **Configure Settings**: Adjust max tokens and image description options | |
| 3. **Select Processing Mode**: Choose between basic extraction or enhanced analysis | |
| 4. **Click Convert**: Watch real-time processing with streaming results | |
| 5. **Download Results**: Get formatted markdown with preserved structure | |
| ## Model Information | |
| - **Core Model**: Nanonets-OCR-s Foundation Model | |
| - **Architecture**: Advanced Vision-Language Transformer | |
| - **Training Data**: 10M+ documents across multiple domains | |
| - **Accuracy**: 99.2% text recognition accuracy | |
| - **Languages**: Multi-language support (English, Spanish, French, German, etc.) | |
| - **Processing Speed**: Real-time streaming with GPU acceleration | |
| """ | |
| join_us = """ | |
| ## Join the Community | |
| 🌟 **Advanced Stock Prediction** is continuously evolving! Join our active builder's community 👻 | |
| [](https://discord.gg/qdfnvSPcqP) | |
| [](https://huggingface.co/TeamTonic) | |
| [](https://github.com/Tonic-AI) | |
| 🤗Big thanks to Yuvi Sharma and all the folks at huggingface for the community grant 🤗 | |
| """ | |
| with gr.Blocks(title="Nanonets-OCR-s: Advanced Document Intelligence", theme=gr.themes.Soft()) as demo: | |
| with gr.Row(): | |
| gr.Markdown(title) | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| with gr.Group(): | |
| gr.Markdown(description) | |
| with gr.Column(scale=1): | |
| with gr.Group(): | |
| gr.Markdown(model_info) | |
| gr.Markdown(join_us) | |
| gr.Markdown("---") # Add a separator | |
| # Main Processing Interface | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| with gr.Group(): | |
| gr.Markdown("### 📤 Document Upload & Configuration") | |
| file_input = gr.Image( | |
| label="Upload Document (Supported formats: PDF, JPEG, PNG, TIFF)", | |
| height=200 | |
| ) | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| max_tokens_slider = gr.Slider( | |
| minimum=1024, | |
| maximum=8192, | |
| value=4096, | |
| step=512, | |
| label="Max Tokens per Page (Higher values = more detailed extraction)" | |
| ) | |
| with gr.Column(scale=1): | |
| with_img_desc_checkbox = gr.Checkbox( | |
| label="Include Image Descriptions (Add AI-generated descriptions for images)", | |
| value=False | |
| ) | |
| extract_btn = gr.Button( | |
| "🚀 Convert to Markdown", | |
| variant="primary", | |
| size="lg", | |
| scale=2 | |
| ) | |
| with gr.Column(scale=2): | |
| with gr.Group(): | |
| gr.Markdown("### 📄 Processing Results") | |
| output_text = gr.Markdown( | |
| label="Extracted Content", | |
| latex_delimiters=[{"left": "$$", "right": "$$", "display": True}, {"left": "$", "right": "$", "display": False}], | |
| line_breaks=True, | |
| show_copy_button=True, | |
| height=600, | |
| ) | |
| # Connect the processing function | |
| extract_btn.click( | |
| fn=process_document, | |
| inputs=[file_input, max_tokens_slider, with_img_desc_checkbox], | |
| concurrency_limit=4, | |
| outputs=output_text | |
| ) | |
| with gr.Accordion("About the Model (Nanonets-OCR-s)", open=False): | |
| gr.Markdown(""" | |
| ### Key Features | |
| - **LaTeX Equation Recognition**: Converts mathematical equations into properly formatted LaTeX. | |
| - **Intelligent Image Description**: Describes images within documents using structured `<img>` tags. | |
| - **Signature & Watermark Detection**: Identifies and isolates signatures and watermarks within `<signature>` and `<watermark>` tags. | |
| - **Smart Checkbox Handling**: Converts form checkboxes into standardized Unicode symbols (☐, ☑). | |
| - **Complex Table Extraction**: Accurately converts tables into HTML format. | |
| """) | |
| if __name__ == "__main__": | |
| demo.queue().launch(debug=True, ssr_mode=False, mcp_server=True) |