Spaces:
Running
on
Zero
Running
on
Zero
| import spaces | |
| import json | |
| import os | |
| import traceback | |
| from io import BytesIO | |
| from typing import Any, Dict, List, Optional, Tuple | |
| import re | |
| import time | |
| from threading import Thread | |
| import gradio as gr | |
| import requests | |
| import torch | |
| from PIL import Image | |
| from transformers import ( | |
| Qwen2_5_VLForConditionalGeneration, | |
| AutoProcessor, | |
| TextIteratorStreamer, | |
| ) | |
| from reportlab.lib.pagesizes import A4 | |
| from reportlab.lib.styles import getSampleStyleSheet | |
| from reportlab.lib import colors | |
| from reportlab.platypus import SimpleDocTemplate, Image as RLImage, Paragraph, Spacer | |
| from reportlab.lib.units import inch | |
| import uuid | |
| # --- Constants and Model Setup --- | |
| MAX_INPUT_TOKEN_LENGTH = 4096 | |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| print("CUDA_VISIBLE_DEVICES=", os.environ.get("CUDA_VISIBLE_DEVICES")) | |
| print("torch.__version__ =", torch.__version__) | |
| print("torch.version.cuda =", torch.version.cuda) | |
| print("cuda available:", torch.cuda.is_available()) | |
| print("cuda device count:", torch.cuda.device_count()) | |
| if torch.cuda.is_available(): | |
| print("current device:", torch.cuda.current_device()) | |
| print("device name:", torch.cuda.get_device_name(torch.cuda.current_device())) | |
| print("Using device:", device) | |
| # --- Model Loading --- | |
| MODEL_ID_M = "prithivMLmods/Camel-Doc-OCR-080125" | |
| processor_m = AutoProcessor.from_pretrained(MODEL_ID_M, trust_remote_code=True) | |
| model_m = Qwen2_5_VLForConditionalGeneration.from_pretrained( | |
| MODEL_ID_M, trust_remote_code=True, torch_dtype=torch.float16 | |
| ).to(device).eval() | |
| MODEL_ID_T = "prithivMLmods/Megalodon-OCR-Sync-0713" | |
| processor_t = AutoProcessor.from_pretrained(MODEL_ID_T, trust_remote_code=True) | |
| model_t = Qwen2_5_VLForConditionalGeneration.from_pretrained( | |
| MODEL_ID_T, trust_remote_code=True, torch_dtype=torch.float16 | |
| ).to(device).eval() | |
| MODEL_ID_C = "nanonets/Nanonets-OCR-s" | |
| processor_c = AutoProcessor.from_pretrained(MODEL_ID_C, trust_remote_code=True) | |
| model_c = Qwen2_5_VLForConditionalGeneration.from_pretrained( | |
| MODEL_ID_C, trust_remote_code=True, torch_dtype=torch.float16 | |
| ).to(device).eval() | |
| MODEL_ID_G = "echo840/MonkeyOCR" | |
| SUBFOLDER = "Recognition" | |
| processor_g = AutoProcessor.from_pretrained( | |
| MODEL_ID_G, trust_remote_code=True, subfolder=SUBFOLDER | |
| ) | |
| model_g = Qwen2_5_VLForConditionalGeneration.from_pretrained( | |
| MODEL_ID_G, trust_remote_code=True, subfolder=SUBFOLDER, torch_dtype=torch.float16 | |
| ).to(device).eval() | |
| MODEL_ID_I = "allenai/olmOCR-7B-0725" | |
| processor_i = AutoProcessor.from_pretrained(MODEL_ID_I, trust_remote_code=True) | |
| model_i = Qwen2_5_VLForConditionalGeneration.from_pretrained( | |
| MODEL_ID_I, trust_remote_code=True, torch_dtype=torch.float16 | |
| ).to(device).eval() | |
| # --- Prompts --- | |
| ocr_prompt = "Perform precise OCR on the image. Extract all text content, maintaining the original structure, paragraphs, and tables as formatted markdown." | |
| # --- PDF Generation Functions --- | |
| def generate_pdf(media_path, plain_text, font_size, line_spacing, alignment, image_size): | |
| """Generates a PDF document.""" | |
| filename = f"output_{uuid.uuid4()}.pdf" | |
| doc = SimpleDocTemplate( | |
| filename, | |
| pagesize=A4, | |
| rightMargin=inch, | |
| leftMargin=inch, | |
| topMargin=inch, | |
| bottomMargin=inch | |
| ) | |
| styles = getSampleStyleSheet() | |
| styles["Normal"].fontSize = int(font_size) | |
| styles["Normal"].leading = int(font_size) * line_spacing | |
| styles["Normal"].alignment = { | |
| "Left": 0, | |
| "Center": 1, | |
| "Right": 2, | |
| "Justified": 4 | |
| }[alignment] | |
| story = [] | |
| # Add image with size adjustment | |
| image_sizes = { | |
| "Small": (200, 200), | |
| "Medium": (400, 400), | |
| "Large": (600, 600) | |
| } | |
| img = RLImage(media_path, width=image_sizes[image_size][0], height=image_sizes[image_size][1]) | |
| story.append(img) | |
| story.append(Spacer(1, 12)) | |
| # Add plain text output | |
| text = Paragraph(plain_text, styles["Normal"]) | |
| story.append(text) | |
| doc.build(story) | |
| return filename | |
| # --- Core Application Logic --- | |
| def process_document_stream(model_name: str, image: Image.Image, max_new_tokens: int, font_size: str, line_spacing: float, alignment: str, image_size: str): | |
| """ | |
| Main generator function for OCR task, also generating PDF for preview. | |
| """ | |
| if image is None: | |
| yield "Please upload an image.", "Please upload an image.", None | |
| return | |
| # Select model and processor | |
| if model_name == "Camel-Doc-OCR-080125": processor, model = processor_m, model_m | |
| elif model_name == "Megalodon-OCR-Sync-0713": processor, model = processor_t, model_t | |
| elif model_name == "Nanonets-OCR-s": processor, model = processor_c, model_c | |
| elif model_name == "MonkeyOCR-Recognition": processor, model = processor_g, model_g | |
| elif model_name == "olmOCR-7B-0725": processor, model = processor_i, model_i | |
| else: | |
| yield "Invalid model selected.", "Invalid model selected.", None | |
| return | |
| # Save image temporarily for PDF generation | |
| temp_image_path = f"temp_{uuid.uuid4()}.png" | |
| image.save(temp_image_path) | |
| # Prepare model inputs and streamer | |
| text_prompt = ocr_prompt | |
| messages = [{"role": "user", "content": [{"type": "image", "image": image}, {"type": "text", "text": text_prompt}]}] | |
| prompt_full = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) | |
| inputs = processor(text=[prompt_full], images=[image], return_tensors="pt", padding=True, truncation=True, max_length=MAX_INPUT_TOKEN_LENGTH).to(device) | |
| streamer = TextIteratorStreamer(processor, skip_prompt=True, skip_special_tokens=True) | |
| generation_kwargs = {**inputs, "streamer": streamer, "max_new_tokens": max_new_tokens} | |
| thread = Thread(target=model.generate, kwargs=generation_kwargs) | |
| thread.start() | |
| # Stream raw output to the UI in real-time | |
| buffer = "" | |
| for new_text in streamer: | |
| buffer += new_text | |
| buffer = buffer.replace("<|im_end|>", "") | |
| time.sleep(0.01) | |
| # Generate PDF with current buffer | |
| pdf_file = generate_pdf(temp_image_path, buffer, font_size, line_spacing, alignment, image_size) | |
| yield buffer, buffer, pdf_file | |
| # Final PDF with complete output | |
| pdf_file = generate_pdf(temp_image_path, buffer, font_size, line_spacing, alignment, image_size) | |
| yield buffer, buffer, pdf_file | |
| # Clean up temporary image file | |
| if os.path.exists(temp_image_path): | |
| os.remove(temp_image_path) | |
| # --- Gradio UI Definition --- | |
| def create_gradio_interface(): | |
| """Builds and returns the Gradio web interface.""" | |
| css = """ | |
| .main-container { max-width: 1400px; margin: 0 auto; } | |
| .process-button { border: none !important; color: white !important; font-weight: bold !important; background-color: blue !important;} | |
| .process-button:hover { background-color: darkblue !important; transform: translateY(-2px) !important; box-shadow: 0 4px 8px rgba(0,0,0,0.2) !important; } | |
| .download-btn { background-color: #35a6d6 !important; color: white !important; } | |
| .download-btn:hover { background-color: #22bcff !important; } | |
| """ | |
| with gr.Blocks(theme="bethecloud/storj_theme", css=css) as demo: | |
| gr.HTML(""" | |
| <div class="title" style="text-align: center"> | |
| <h1>Tiny VLMs Lab🧪</h1> | |
| <p style="font-size: 1.1em; color: #6b7280; margin-bottom: 0.6em;"> | |
| Advanced Vision-Language Model for Image Content Extraction and PDF Generation | |
| </p> | |
| </div> | |
| """) | |
| with gr.Row(): | |
| # Left Column (Inputs) | |
| with gr.Column(scale=1): | |
| model_choice = gr.Dropdown( | |
| choices=[ | |
| "Camel-Doc-OCR-080125", | |
| "MonkeyOCR-Recognition", | |
| "olmOCR-7B-0725", | |
| "Nanonets-OCR-s", | |
| "Megalodon-OCR-Sync-0713" | |
| ], | |
| label="Select Model", | |
| value="Nanonets-OCR-s" | |
| ) | |
| image_input = gr.Image(label="Upload Image", type="pil", sources=['upload']) | |
| with gr.Accordion("Advanced Settings", open=False): | |
| max_new_tokens = gr.Slider(minimum=512, maximum=8192, value=4096, step=256, label="Max New Tokens") | |
| font_size = gr.Dropdown( | |
| choices=["8", "10", "12", "14", "16", "18", "20", "22", "24"], | |
| value="16", | |
| label="Font Size" | |
| ) | |
| line_spacing = gr.Dropdown( | |
| choices=[0.5, 1.0, 1.15, 1.5, 2.0, 2.5, 3.0], | |
| value=1.5, | |
| label="Line Spacing" | |
| ) | |
| alignment = gr.Dropdown( | |
| choices=["Left", "Center", "Right", "Justified"], | |
| value="Justified", | |
| label="Text Alignment" | |
| ) | |
| image_size = gr.Dropdown( | |
| choices=["Small", "Medium", "Large"], | |
| value="Medium", | |
| label="Image Size" | |
| ) | |
| process_btn = gr.Button("🚀 Process Document", variant="primary", elem_classes=["process-button"], size="lg") | |
| clear_btn = gr.Button("🗑️ Clear All", variant="secondary") | |
| # Right Column (Outputs) | |
| with gr.Column(scale=2): | |
| with gr.Tabs() as tabs: | |
| with gr.Tab("📝 Extracted Content"): | |
| raw_output_stream = gr.Textbox(label="Raw Model Output Stream", interactive=False, lines=13, show_copy_button=True) | |
| with gr.Row(): | |
| examples = gr.Examples( | |
| examples=["examples/1.png", "examples/2.png", "examples/3.png", "examples/4.png", "examples/5.png"], | |
| inputs=image_input, | |
| label="Examples" | |
| ) | |
| gr.Markdown("[Report-Bug💻](https://huggingface.co/spaces/prithivMLmods/OCR-Comparator/discussions)") | |
| with gr.Tab("📰 README.md"): | |
| with gr.Accordion("(Formatted Result)", open=True): | |
| markdown_output = gr.Markdown(label="Formatted Markdown") | |
| with gr.Tab("📋 PDF Preview"): | |
| pdf_output = gr.File(label="Download PDF", interactive=True) | |
| # Event Handlers | |
| def clear_all_outputs(): | |
| return None, "Raw output will appear here.", "Formatted results will appear here.", None | |
| process_btn.click( | |
| fn=process_document_stream, | |
| inputs=[model_choice, image_input, max_new_tokens, font_size, line_spacing, alignment, image_size], | |
| outputs=[raw_output_stream, markdown_output, pdf_output] | |
| ) | |
| clear_btn.click( | |
| fn=clear_all_outputs, | |
| outputs=[image_input, raw_output_stream, markdown_output, pdf_output] | |
| ) | |
| return demo | |
| if __name__ == "__main__": | |
| demo = create_gradio_interface() | |
| demo.queue(max_size=50).launch(share=True, ssr_mode=False, show_error=True) |