prithivMLmods/POINTS-Reader-OCR
18
1import spaces2import json3import math4import os5import traceback6from io import BytesIO7from typing import Any, Dict, List, Optional, Tuple8from typing import Iterable9import re10import time11from threading import Thread12from io import BytesIO13import subprocess14import uuid15import tempfile16 17import gradio as gr18import requests19import torch20from PIL import Image21import fitz22import numpy as np23 24from transformers import AutoModelForCausalLM, AutoTokenizer, Qwen2VLImageProcessor25 26from reportlab.lib.pagesizes import A427from reportlab.lib.styles import getSampleStyleSheet28from reportlab.platypus import SimpleDocTemplate, Image as RLImage, Paragraph, Spacer29from reportlab.lib.units import inch30 31from transformers.image_utils import load_image32from gradio.themes import Soft33from gradio.themes.utils import colors, fonts, sizes34 35colors.steel_blue = colors.Color(36 name="steel_blue",37 c50="#EBF3F8",38 c100="#D3E5F0",39 c200="#A8CCE1",40 c300="#7DB3D2",41 c400="#529AC3",42 c500="#4682B4",43 c600="#3E72A0",44 c700="#36638C",45 c800="#2E5378",46 c900="#264364",47 c950="#1E3450",48)49 50class SteelBlueTheme(Soft):51 def __init__(52 self,53 *,54 primary_hue: colors.Color | str = colors.gray,55 secondary_hue: colors.Color | str = colors.steel_blue,56 neutral_hue: colors.Color | str = colors.slate,57 text_size: sizes.Size | str = sizes.text_lg,58 font: fonts.Font | str | Iterable[fonts.Font | str] = (59 fonts.GoogleFont("Outfit"), "Arial", "sans-serif",60 ),61 font_mono: fonts.Font | str | Iterable[fonts.Font | str] = (62 fonts.GoogleFont("IBM Plex Mono"), "ui-monospace", "monospace",63 ),64 ):65 super().__init__(66 primary_hue=primary_hue,67 secondary_hue=secondary_hue,68 neutral_hue=neutral_hue,69 text_size=text_size,70 font=font,71 font_mono=font_mono,72 )73 super().set(74 background_fill_primary="*primary_50",75 background_fill_primary_dark="*primary_900",76 body_background_fill="linear-gradient(135deg, *primary_200, *primary_100)",77 body_background_fill_dark="linear-gradient(135deg, *primary_900, *primary_800)",78 button_primary_text_color="white",79 button_primary_text_color_hover="white",80 button_primary_background_fill="linear-gradient(90deg, *secondary_500, *secondary_600)",81 button_primary_background_fill_hover="linear-gradient(90deg, *secondary_600, *secondary_700)",82 button_primary_background_fill_dark="linear-gradient(90deg, *secondary_600, *secondary_800)",83 button_primary_background_fill_hover_dark="linear-gradient(90deg, *secondary_500, *secondary_500)",84 button_secondary_text_color="black",85 button_secondary_text_color_hover="white",86 button_secondary_background_fill="linear-gradient(90deg, *primary_300, *primary_300)",87 button_secondary_background_fill_hover="linear-gradient(90deg, *primary_400, *primary_400)",88 button_secondary_background_fill_dark="linear-gradient(90deg, *primary_500, *primary_600)",89 button_secondary_background_fill_hover_dark="linear-gradient(90deg, *primary_500, *primary_500)",90 slider_color="*secondary_500",91 slider_color_dark="*secondary_600",92 block_title_text_weight="600",93 block_border_width="3px",94 block_shadow="*shadow_drop_lg",95 button_primary_shadow="*shadow_drop_lg",96 button_large_padding="11px",97 color_accent_soft="*primary_100",98 block_label_background_fill="*primary_200",99 )100 101steel_blue_theme = SteelBlueTheme()102 103# --- Constants and Model Setup ---104MAX_INPUT_TOKEN_LENGTH = 4096105device = torch.device("cuda" if torch.cuda.is_available() else "cpu")106 107print("CUDA_VISIBLE_DEVICES=", os.environ.get("CUDA_VISIBLE_DEVICES"))108print("torch.__version__ =", torch.__version__)109print("torch.version.cuda =", torch.version.cuda)110print("cuda available:", torch.cuda.is_available())111print("cuda device count:", torch.cuda.device_count())112if torch.cuda.is_available():113 print("current device:", torch.cuda.current_device())114 print("device name:", torch.cuda.get_device_name(torch.cuda.current_device()))115 116print("Using device:", device)117 118 119# --- Model Loading: tencent/POINTS-Reader ---120MODEL_PATH = 'tencent/POINTS-Reader'121 122print(f"Loading model: {MODEL_PATH}")123model = AutoModelForCausalLM.from_pretrained(124 MODEL_PATH,125 trust_remote_code=True,126 torch_dtype=torch.bfloat16,127 device_map='auto'128)129tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, trust_remote_code=True)130image_processor = Qwen2VLImageProcessor.from_pretrained(MODEL_PATH)131print("Model loaded successfully.")132 133 134# --- PDF Generation and Preview Utility Function ---135def generate_and_preview_pdf(image: Image.Image, text_content: str, font_size: int, line_spacing: float, alignment: str, image_size: str):136 """137 Generates a PDF, saves it, and then creates image previews of its pages.138 Returns the path to the PDF and a list of paths to the preview images.139 """140 if image is None or not text_content or not text_content.strip():141 raise gr.Error("Cannot generate PDF. Image or text content is missing.")142 143 # --- 1. Generate the PDF ---144 temp_dir = tempfile.gettempdir()145 pdf_filename = os.path.join(temp_dir, f"output_{uuid.uuid4()}.pdf")146 doc = SimpleDocTemplate(147 pdf_filename,148 pagesize=A4,149 rightMargin=inch, leftMargin=inch,150 topMargin=inch, bottomMargin=inch151 )152 styles = getSampleStyleSheet()153 style_normal = styles["Normal"]154 style_normal.fontSize = int(font_size)155 style_normal.leading = int(font_size) * line_spacing156 style_normal.alignment = {"Left": 0, "Center": 1, "Right": 2, "Justified": 4}[alignment]157 158 story = []159 160 img_buffer = BytesIO()161 image.save(img_buffer, format='PNG')162 img_buffer.seek(0)163 164 page_width, _ = A4165 available_width = page_width - 2 * inch166 image_widths = {167 "Small": available_width * 0.3,168 "Medium": available_width * 0.6,169 "Large": available_width * 0.9,170 }171 img_width = image_widths[image_size]172 img = RLImage(img_buffer, width=img_width, height=image.height * (img_width / image.width))173 story.append(img)174 story.append(Spacer(1, 12))175 176 cleaned_text = re.sub(r'#+\s*', '', text_content).replace("*", "")177 text_paragraphs = cleaned_text.split('\n')178 179 for para in text_paragraphs:180 if para.strip():181 story.append(Paragraph(para, style_normal))182 183 doc.build(story)184 185 # --- 2. Render PDF pages as images for preview ---186 preview_images = []187 try:188 pdf_doc = fitz.open(pdf_filename)189 for page_num in range(len(pdf_doc)):190 page = pdf_doc.load_page(page_num)191 pix = page.get_pixmap(dpi=150)192 preview_img_path = os.path.join(temp_dir, f"preview_{uuid.uuid4()}_p{page_num}.png")193 pix.save(preview_img_path)194 preview_images.append(preview_img_path)195 pdf_doc.close()196 except Exception as e:197 print(f"Error generating PDF preview: {e}")198 199 return pdf_filename, preview_images200 201 202# --- Core Application Logic ---203@spaces.GPU204def process_document_stream(205 image: Image.Image, 206 prompt_input: str,207 image_scale_factor: float, # New parameter for image scaling208 max_new_tokens: int,209 temperature: float,210 top_p: float,211 top_k: int,212 repetition_penalty: float213):214 """215 Main function that handles model inference using tencent/POINTS-Reader.216 """217 if image is None:218 yield "Please upload an image.", ""219 return220 if not prompt_input or not prompt_input.strip():221 yield "Please enter a prompt.", ""222 return223 224 # --- IMPLEMENTATION: Image Scaling based on user input ---225 if image_scale_factor > 1.0:226 try:227 original_width, original_height = image.size228 new_width = int(original_width * image_scale_factor)229 new_height = int(original_height * image_scale_factor)230 print(f"Scaling image from {image.size} to ({new_width}, {new_height}) with factor {image_scale_factor}.")231 # Use a high-quality resampling filter for better results232 image = image.resize((new_width, new_height), Image.Resampling.LANCZOS)233 except Exception as e:234 print(f"Error during image scaling: {e}")235 # Continue with the original image if scaling fails236 pass237 # --- END IMPLEMENTATION ---238 239 temp_image_path = None240 try:241 # --- FIX: Save the PIL Image to a temporary file ---242 # The model expects a file path, not a PIL object.243 temp_dir = tempfile.gettempdir()244 temp_image_path = os.path.join(temp_dir, f"temp_image_{uuid.uuid4()}.png")245 image.save(temp_image_path)246 247 # Prepare content for the model using the temporary file path248 content = [249 dict(type='image', image=temp_image_path),250 dict(type='text', text=prompt_input)251 ]252 messages = [253 {254 'role': 'user',255 'content': content256 }257 ]258 259 # Prepare generation configuration from UI inputs260 generation_config = {261 'max_new_tokens': max_new_tokens,262 'repetition_penalty': repetition_penalty,263 'temperature': temperature,264 'top_p': top_p,265 'top_k': top_k,266 'do_sample': True if temperature > 0 else False267 }268 269 # Run inference270 response = model.chat(271 messages,272 tokenizer,273 image_processor,274 generation_config275 )276 # Yield the full response at once277 yield response, response278 279 except Exception as e:280 traceback.print_exc()281 yield f"An error occurred during processing: {str(e)}", ""282 finally:283 # --- Clean up the temporary image file ---284 if temp_image_path and os.path.exists(temp_image_path):285 os.remove(temp_image_path)286 287css = """288.main-container { max-width: 1400px; margin: 0 auto; }289.process-button { border: none !important; color: white !important; font-weight: bold !important; background-color: blue !important;}290.process-button:hover { background-color: darkblue !important; transform: translateY(-2px) !important; box-shadow: 0 4px 8px rgba(0,0,0,0.2) !important; }291#gallery { min-height: 400px; }292"""293 294def create_gradio_interface():295 with gr.Blocks() as demo:296 gr.HTML(f"""297 <div class="title" style="text-align: center">298 <h1>Document Conversion with POINTS Reader 📖</h1>299 <p style="font-size: 1.1em; color: #6b7280; margin-bottom: 0.6em;">300 Using tencent/POINTS-Reader Multimodal for Image Content Extraction301 </p>302 </div>303 """)304 305 with gr.Row():306 # Left Column (Inputs)307 with gr.Column(scale=1):308 gr.Textbox(309 label="Model in Use",310 value="tencent/POINTS-Reader",311 interactive=False312 )313 prompt_input = gr.Textbox(314 label="Query Input",315 placeholder="✦︎ Enter the prompt",316 value="Perform OCR on the image precisely.",317 )318 image_input = gr.Image(label="Upload Image", type="pil", sources=['upload'])319 320 with gr.Accordion("Advanced Settings", open=False):321 # --- NEW UI ELEMENT: Image Scaling Slider ---322 image_scale_factor = gr.Slider(323 minimum=1.0, 324 maximum=3.0, 325 value=1.0, 326 step=0.1, 327 label="Image Upscale Factor",328 info="Increases image size before processing. Can improve OCR on small text. Default: 1.0 (no change)."329 )330 # --- END NEW UI ELEMENT ---331 max_new_tokens = gr.Slider(minimum=512, maximum=8192, value=2048, step=256, label="Max New Tokens")332 temperature = gr.Slider(label="Temperature", minimum=0.1, maximum=1.0, step=0.05, value=0.7)333 top_p = gr.Slider(label="Top-p (nucleus sampling)", minimum=0.05, maximum=1.0, step=0.05, value=0.8)334 top_k = gr.Slider(label="Top-k", minimum=1, maximum=100, step=1, value=20)335 repetition_penalty = gr.Slider(label="Repetition penalty", minimum=1.0, maximum=2.0, step=0.05, value=1.05)336 337 gr.Markdown("### PDF Export Settings")338 font_size = gr.Dropdown(choices=["8", "10", "12", "14", "16", "18"], value="12", label="Font Size")339 line_spacing = gr.Dropdown(choices=[1.0, 1.15, 1.5, 2.0], value=1.15, label="Line Spacing")340 alignment = gr.Dropdown(choices=["Left", "Center", "Right", "Justified"], value="Justified", label="Text Alignment")341 image_size = gr.Dropdown(choices=["Small", "Medium", "Large"], value="Medium", label="Image Size in PDF")342 343 process_btn = gr.Button("🚀 Process Image", variant="primary", elem_classes=["process-button"], size="lg")344 clear_btn = gr.Button("🗑️ Clear All", variant="secondary")345 346 # Right Column (Outputs)347 with gr.Column(scale=2):348 with gr.Tabs() as tabs:349 with gr.Tab("📝 Extracted Content"):350 raw_output_stream = gr.Textbox(label="Raw Model Output (max T ≤ 120s)", interactive=True, lines=15)351 with gr.Row():352 examples = gr.Examples(353 examples=["examples/1.jpeg", 354 "examples/2.jpeg", 355 "examples/3.jpeg",356 "examples/4.jpeg", 357 "examples/5.jpeg"],358 inputs=image_input, label="Examples"359 )360 gr.Markdown("[Report-Bug💻](https://huggingface.co/spaces/prithivMLmods/POINTS-Reader-OCR/discussions) | [prithivMLmods🤗](https://huggingface.co/prithivMLmods)")361 362 with gr.Tab("📰 README.md"):363 with gr.Accordion("(Result.md)", open=True): 364 # --- FIX: Added latex_delimiters to enable LaTeX rendering ---365 markdown_output = gr.Markdown(latex_delimiters=[366 {"left": "$$", "right": "$$", "display": True},367 {"left": "$", "right": "$", "display": False}368 ])369 370 with gr.Tab("📋 PDF Preview"):371 generate_pdf_btn = gr.Button("📄 Generate PDF & Render", variant="primary")372 pdf_output_file = gr.File(label="Download Generated PDF", interactive=False)373 pdf_preview_gallery = gr.Gallery(label="PDF Page Preview", show_label=True, elem_id="gallery", columns=2, object_fit="contain", height="auto")374 375 # Event Handlers376 def clear_all_outputs():377 return None, "", "Raw output will appear here.", "", None, None378 379 process_btn.click(380 fn=process_document_stream,381 # --- UPDATE: Add the new slider to the inputs list ---382 inputs=[image_input, prompt_input, image_scale_factor, max_new_tokens, temperature, top_p, top_k, repetition_penalty],383 outputs=[raw_output_stream, markdown_output]384 )385 386 generate_pdf_btn.click(387 fn=generate_and_preview_pdf,388 inputs=[image_input, raw_output_stream, font_size, line_spacing, alignment, image_size],389 outputs=[pdf_output_file, pdf_preview_gallery]390 )391 392 clear_btn.click(393 clear_all_outputs,394 outputs=[image_input, prompt_input, raw_output_stream, markdown_output, pdf_output_file, pdf_preview_gallery]395 )396 return demo397 398if __name__ == "__main__":399 demo = create_gradio_interface()400 demo.queue(max_size=50).launch(theme=steel_blue_theme, css=css, mcp_server=True, ssr_mode=False, show_error=True)