import math import os import re from pathlib import Path import gradio as gr import spaces import torch from transformers import AutoModel, AutoTokenizer from image_utils import load_image MODEL_PATH = os.getenv("MODEL_PATH", "baidu/Qianfan-OCR") MAX_TILES_PER_IMAGE = int(os.getenv("MAX_TILES_PER_IMAGE", "12")) DEFAULT_MAX_NEW_TOKENS = 2048 MAX_NEW_TOKENS_LIMIT = 8192 ZEROGPU_DURATION_MIN = 60 ZEROGPU_DURATION_MAX = 300 DEFAULT_PROMPT = "Please extract the text from the image." IMAGE_FILE_TYPES = [".png", ".jpg", ".jpeg", ".webp", ".bmp", ".tif", ".tiff"] APP_DIR = Path(__file__).resolve().parent EXAMPLES_DIR = APP_DIR / "examples" LATEX_DELIMITERS = [ {"left": "$$", "right": "$$", "display": True}, {"left": "$", "right": "$", "display": False}, {"left": "\$", "right": "\$", "display": False}, {"left": "\\[", "right": "\\]", "display": True}, ] LATEX_FENCE_RE = re.compile( r"(^|\n)```(?:latex|tex)[ \t]*\n(?P.*?)(?:\n```)(?=\n|$)", re.IGNORECASE | re.DOTALL, ) THINK_BLOCK_RE = re.compile(r".*?", re.DOTALL | re.IGNORECASE) EXAMPLE_ITEMS = [ { "task": "Text Extraction", "prompt": "Please extract the text from the image.", "image": str(EXAMPLES_DIR / "text_block.png"), }, { "task": "Formula Parsing", "prompt": "Please convert the formula in the image to LaTeX.", "image": str(EXAMPLES_DIR / "formula_block.png"), }, { "task": "Table Parsing", "prompt": "Please convert the table in the image to HTML.", "image": str(EXAMPLES_DIR / "table_block.png"), }, { "task": "Document Parsing", "prompt": "Parse this document to Markdown.", "image": str(EXAMPLES_DIR / "document.png"), }, { "task": "Multilingual Scene Text Recognition", "prompt": "Please output the text content from the image.", "image": str(EXAMPLES_DIR / "information_board.jpg"), }, { "task": "Key Information Extraction", "prompt": """请从图片中提取以下信息：“发票号码、开票日期、发票种类、购方名称、购方纳税人识别号、合计金额、合计税额、价税合计(小写)、价税合计(大写)、货物名称、规格型号、单位、数量、单价、金额、税率、税额、销售方名称、销售方纳税人识别号、开票人、备注” 注意： 1. 仅提取清晰可见且可确定的文字内容；模糊、缺失或无法确认的部分请忽略，不要猜测或补全。 2. 值完整提取及格式保持：提取完整的字段值，包括数字、符号、单位等，保持原始格式，包括空格、标点符号等。 3. 使用标准JSON格式输出，字段值保持原始格式。""", "image": str(EXAMPLES_DIR / "invoice.jpg"), }, ] EXAMPLE_ITEMS = [example for example in EXAMPLE_ITEMS if Path(example["image"]).exists()] _MODEL = AutoModel.from_pretrained( MODEL_PATH, torch_dtype="auto", trust_remote_code=True, ).to("cuda") _TOKENIZER = AutoTokenizer.from_pretrained( MODEL_PATH, trust_remote_code=True, ) def load_images(file_paths, device, dtype: torch.dtype): if not file_paths: raise gr.Error("Please upload at least one image.") pixel_values_list = [] for path in file_paths: if not path: continue pixel_values = load_image(str(path), max_num=MAX_TILES_PER_IMAGE) pixel_values_list.append(pixel_values) if not pixel_values_list: raise gr.Error("No valid image files were received.") pixel_values = torch.cat(pixel_values_list, dim=0).to(device=device, dtype=dtype) return pixel_values def validate_prompt(prompt: str) -> str: prompt = (prompt or "").strip() if not prompt: raise gr.Error("Please enter a prompt.") return prompt def build_question(prompt: str, layout_as_thought: bool) -> str: question = validate_prompt(prompt) if layout_as_thought and not question.endswith(""): question = f"{question}" return question def preview_uploaded_images(file_paths): if not file_paths: return [] return [str(path) for path in file_paths if path] def clear_form(): return None, [], DEFAULT_PROMPT, False, DEFAULT_MAX_NEW_TOKENS, "" def load_example(example_index: int): if example_index is None or not (0 <= int(example_index) < len(EXAMPLE_ITEMS)): raise gr.Error(f"Unknown example index: {example_index}") example = EXAMPLE_ITEMS[int(example_index)] image_paths = [example["image"]] return ( image_paths, image_paths, example["prompt"], False, DEFAULT_MAX_NEW_TOKENS, "", ) def normalize_markdown_math(text: str) -> str: if not text: return text def replace_latex_fence(match: re.Match[str]) -> str: body = match.group("body").strip() if not body: return match.group(0) if body.startswith("$$") and body.endswith("$$"): return f"{match.group(1)}{body}" if body.startswith("\\[") and body.endswith("\\]"): return f"{match.group(1)}{body}" return f"{match.group(1)}$$\n{body}\n$$" return LATEX_FENCE_RE.sub(replace_latex_fence, text) def wrap_think_blocks(text: str) -> str: if not text: return text def replace_think_block(match: re.Match[str]) -> str: block = match.group(0).strip() return f"\n```text\n{block}\n```\n" return THINK_BLOCK_RE.sub(replace_think_block, text) def estimate_zerogpu_duration(file_paths, prompt, layout_as_thought, max_new_tokens): del file_paths, prompt, layout_as_thought estimated_duration = math.ceil(int(max_new_tokens) / 25 + 15) final_duration = max( ZEROGPU_DURATION_MIN, min(ZEROGPU_DURATION_MAX, estimated_duration), ) return final_duration @spaces.GPU(duration=estimate_zerogpu_duration) def run_inference(file_paths, prompt, layout_as_thought, max_new_tokens): if not file_paths: raise gr.Error("Please upload at least one image.") pixel_values = load_images(file_paths, _MODEL.device, _MODEL.dtype) question = build_question(prompt, layout_as_thought) generation_config = {"max_new_tokens": int(max_new_tokens)} with torch.no_grad(): response = _MODEL.chat( _TOKENIZER, pixel_values=pixel_values, question=question, generation_config=generation_config, ) return normalize_markdown_math(wrap_think_blocks(response)) def build_demo(): with gr.Blocks(title="Qianfan-OCR Demo") as demo: gr.Markdown( """ # Qianfan-OCR Demo 📄 Technical Report | 🖥️ Qianfan Platform | 💻 GitHub | 🧩 Skill Qianfan-OCR is a 4B-parameter end-to-end document intelligence model developed by the Baidu Qianfan Team. It unifies document parsing, layout analysis, and document understanding within a single vision-language architecture. """, sanitize_html=False, ) with gr.Row(equal_height=True): with gr.Column(scale=1): image_input = gr.File( label="Images", file_count="multiple", file_types=IMAGE_FILE_TYPES, type="filepath", ) image_preview = gr.Gallery( label="Uploaded Images", columns=2, height="auto", show_label=True, ) with gr.Column(scale=1): prompt_input = gr.Textbox( label="Prompt", lines=18, value=DEFAULT_PROMPT, placeholder="Describe the task for the uploaded image(s).", ) layout_as_thought_input = gr.Checkbox( label="Layout-as-Thought", value=False, ) max_new_tokens_input = gr.Slider( label="max_new_tokens", minimum=256, maximum=MAX_NEW_TOKENS_LIMIT, step=256, value=DEFAULT_MAX_NEW_TOKENS, ) with gr.Row(): submit_button = gr.Button("Run", variant="primary") clear_button = gr.Button("Clear") with gr.Column(scale=1): output_box = gr.Markdown( value="", line_breaks=True, sanitize_html=False, latex_delimiters=LATEX_DELIMITERS, height=520, buttons=["copy"], container=True, padding=True, ) if EXAMPLE_ITEMS: example_dataset = gr.Dataset( label="Examples", components=[ gr.Textbox(label="Task", render=False), gr.Image(label="Image", type="filepath", height=120, render=False), gr.Textbox(label="Prompt", render=False), ], samples=[ [example["task"], example["image"], example["prompt"]] for example in EXAMPLE_ITEMS ], type="index", layout="table", ) example_dataset.click( fn=load_example, inputs=example_dataset, outputs=[ image_input, image_preview, prompt_input, layout_as_thought_input, max_new_tokens_input, output_box, ], ) image_input.change( fn=preview_uploaded_images, inputs=image_input, outputs=image_preview, ) submit_button.click( fn=run_inference, inputs=[image_input, prompt_input, layout_as_thought_input, max_new_tokens_input], outputs=output_box, ) clear_button.click( fn=clear_form, outputs=[ image_input, image_preview, prompt_input, layout_as_thought_input, max_new_tokens_input, output_box, ], ) return demo if __name__ == "__main__": app = build_demo() app.queue().launch()