import base64 import io import json import os from typing import Dict, List, Tuple, Any, Optional import time import requests from PIL import Image import gradio as gr import re import tempfile from urllib.parse import urlparse # ========================= # Config # ========================= DEFAULT_API_URL = os.environ.get("API_URL") TOKEN = os.environ.get("TOKEN") LOGO_IMAGE_PATH = "./assets/logo.jpg" GOOGLE_FONTS_URL = "" LATEX_DELIMS = [ {"left": "$$", "right": "$$", "display": True}, {"left": "$", "right": "$", "display": False}, {"left": "\\(", "right": "\\)", "display": False}, {"left": "\\[", "right": "\\]", "display": True}, ] AUTH_HEADER = {"Authorization": f"bearer {TOKEN}"} if TOKEN else {} JSON_HEADERS = {**AUTH_HEADER, "Content-Type": "application/json"} if AUTH_HEADER else {"Content-Type": "application/json"} # ========================= # Base64 & Examples (URL直链渲染) # ========================= def image_to_base64_data_url(filepath: str) -> str: """仅用于本地上传预览的兼容方案；URL 预览不会用到它。""" try: ext = os.path.splitext(filepath)[1].lower() mime_types = {".jpg": "image/jpeg", ".jpeg": "image/jpeg", ".png": "image/png", ".gif": "image/gif", ".webp": "image/webp", ".bmp": "image/bmp"} mime_type = mime_types.get(ext, "image/jpeg") with open(filepath, "rb") as image_file: encoded_string = base64.b64encode(image_file.read()).decode("utf-8") return f"data:{mime_type};base64,{encoded_string}" except Exception as e: print(f"Error encoding image to Base64: {e}") return "" def _escape_inequalities_in_math(md: str) -> str: """把数学块中的 < > 替换为 \\lt \\gt，避免被 Markdown 误解析。""" _MATH_PATTERNS = [ re.compile(r"\$\$([\s\S]+?)\$\$"), re.compile(r"\$([^\$]+?)\$"), re.compile(r"\\\[([\s\S]+?)\\\]"), re.compile(r"\\\(([\s\S]+?)\\\)"), ] def fix(s: str) -> str: s = s.replace("<=", r" \le ").replace(">=", r" \ge ") s = s.replace("≤", r" \le ").replace("≥", r" \ge ") s = s.replace("<", r" \lt ").replace(">", r" \gt ") return s for pat in _MATH_PATTERNS: md = pat.sub(lambda m: m.group(0).replace(m.group(1), fix(m.group(1))), md) return md def _get_examples_from_dir(dir_path: str) -> List[List[str]]: """ 从本地目录读取文件名，拼出远程直链 URL（不下载、不转码），用于 直接渲染。 你原来使用的 BOS 基础路径保留。 """ BASE_URL = "https://paddle-model-ecology.bj.bcebos.com/PPOCRVL/dataset/examples" supported_exts = {".png", ".jpg", ".jpeg", ".bmp", ".webp"} examples = [] if not os.path.exists(dir_path): print(f"Warning: example dir {dir_path} not found.") return [] for filename in sorted(os.listdir(dir_path)): ext = os.path.splitext(filename)[1].lower() if ext in supported_exts: subdir = os.path.basename(dir_path.rstrip("/")) img_url = f"{BASE_URL}/{subdir}/{filename}" examples.append([img_url]) return examples def _on_gallery_select(example_paths: List[str], evt: gr.SelectData): """ 与原版不同：直接返回 URL，不再下载到本地临时文件。 """ idx = evt.index selected = example_paths[idx] if isinstance(selected, list): selected = selected[0] return selected # 直接是 https://... URL TARGETED_EXAMPLES_DIR = "examples/targeted" COMPLEX_EXAMPLES_DIR = "examples/complex" targeted_recognition_examples = _get_examples_from_dir(TARGETED_EXAMPLES_DIR) complex_document_examples = _get_examples_from_dir(COMPLEX_EXAMPLES_DIR) # ========================= # UI Helpers（URL直链渲染） # ========================= def render_uploaded_image_div(path_or_url: str) -> str: """ 支持两种输入： - 远程 URL：直接用 渲染 - 本地文件：为兼容旧逻辑，依然转 data: URL 预览（也可以改为 File 组件，这里先保持一致） """ if not path_or_url: return "" is_url = isinstance(path_or_url, str) and path_or_url.startswith(("http://", "https://")) if is_url: src = path_or_url # 直接远程URL else: src = image_to_base64_data_url(path_or_url) # 本地上传时的兼容 return f"""
Preview image
""" def update_preview_visibility(path_or_url: Optional[str]) -> Dict: if path_or_url: html_content = render_uploaded_image_div(path_or_url) return gr.update(value=html_content, visible=True) else: return gr.update(value="", visible=False) # ========================= # API 调用逻辑（支持URL或本地文件） # ========================= def _file_to_b64_image_only(path_or_url: str) -> Tuple[str, int]: """ 输入可以是本地文件路径或远程URL。 - URL：仅在发请求给后端时下载字节转Base64（不影响前端渲染）。 - 本地：读取文件字节。 """ if not path_or_url: raise ValueError("Please upload an image first.") is_url = isinstance(path_or_url, str) and path_or_url.startswith(("http://", "https://")) content: bytes if is_url: r = requests.get(path_or_url, timeout=600) r.raise_for_status() content = r.content ext = os.path.splitext(urlparse(path_or_url).path)[1].lower() else: ext = os.path.splitext(path_or_url)[1].lower() with open(path_or_url, "rb") as f: content = f.read() # 放宽后缀限制：有些URL可能没有后缀，这里仅在极端情况下提示 supported = {".png", ".jpg", ".jpeg", ".bmp", ".webp"} if ext and (ext not in supported): print(f"Warning: file extension {ext} not in supported set {supported}, continue anyway.") return base64.b64encode(content).decode("utf-8"), 1 # 1 = image 类型 def _call_api(api_url: str, path_or_url: str, use_layout_detection: bool, prompt_label: Optional[str], use_chart_recognition: bool = False, use_doc_unwarping: bool = True, use_doc_orientation_classify: bool = True) -> Dict[str, Any]: b64, file_type = _file_to_b64_image_only(path_or_url) payload = { "file": b64, "useLayoutDetection": bool(use_layout_detection), "fileType": file_type, "useDocUnwarping": use_doc_unwarping, "useDocOrientationClassify": use_doc_orientation_classify } if not use_layout_detection: if not prompt_label: raise ValueError("Please select a recognition type.") payload["promptLabel"] = prompt_label.strip().lower() if use_layout_detection and use_chart_recognition: payload["useChartRecognition"] = True try: print(f"Sending API request to {api_url}...") start_time = time.time() resp = requests.post(api_url, json=payload, headers=JSON_HEADERS, timeout=600) end_time = time.time() print(f"Received API response in {end_time - start_time:.2f} seconds.") resp.raise_for_status() data = resp.json() except requests.exceptions.RequestException as e: raise gr.Error(f"API request failed: {e}") except json.JSONDecodeError: raise gr.Error(f"Invalid JSON response from server:\n{getattr(resp, 'text', '')}") if data.get("errorCode", -1) != 0: raise gr.Error("API returned an error:") return data def _process_api_response_page(result: Dict[str, Any]) -> Tuple[str, str, str]: """ 处理后端返回结果： 1) 把 markdown 里的占位图路径替换为真实URL 2) 构造一个可视化（如果有） """ layout_results = (result or {}).get("layoutParsingResults", []) if not layout_results: return "No content was recognized.", "

No visualization available.

", "" page0 = layout_results[0] or {} md_data = page0.get("markdown") or {} md_text = md_data.get("text", "") or "" md_images_map = md_data.get("images", {}) if md_images_map: for placeholder_path, image_url in md_images_map.items(): md_text = md_text.replace(f'src="{placeholder_path}"', f'src="{image_url}"') \ .replace(f']({placeholder_path})', f']({image_url})') output_html = "

No visualization image available.

" out_imgs = page0.get("outputImages") or {} sorted_urls = [img_url for _, img_url in sorted(out_imgs.items()) if img_url] output_image_url: Optional[str] = None if len(sorted_urls) >= 2: output_image_url = sorted_urls[1] elif sorted_urls: output_image_url = sorted_urls[0] if output_image_url: print(f"Found visualization image URL: {output_image_url}") output_html = f'Detection Visualization' md_text = _escape_inequalities_in_math(md_text) return md_text or "(Empty result)", output_html, md_text def handle_complex_doc(path_or_url: str, use_chart_recognition: bool, use_doc_unwarping: bool, use_doc_orientation_classify: bool) -> Tuple[str, str, str]: if not path_or_url: raise gr.Error("Please upload an image first.") data = _call_api(DEFAULT_API_URL, path_or_url, use_layout_detection=True, prompt_label=None, use_chart_recognition=use_chart_recognition, use_doc_unwarping=use_doc_unwarping, use_doc_orientation_classify=use_doc_orientation_classify) result = data.get("result", {}) return _process_api_response_page(result) def handle_targeted_recognition(path_or_url: str, prompt_choice: str) -> Tuple[str, str]: if not path_or_url: raise gr.Error("Please upload an image first.") mapping = { "Text Recognition": "ocr", "Formula Recognition": "formula", "Table Recognition": "table", "Chart Recognition": "chart", } label = mapping.get(prompt_choice, "ocr") data = _call_api(DEFAULT_API_URL, path_or_url, use_layout_detection=False, prompt_label=label, use_doc_unwarping=False, use_doc_orientation_classify=False) result = data.get("result", {}) md_preview, _, md_raw = _process_api_response_page(result) return md_preview, md_raw # ========================= # CSS & UI # ========================= custom_css = """ body, .gradio-container { font-family: "Noto Sans SC", "Microsoft YaHei", "PingFang SC", sans-serif; } .app-header { text-align: center; max-width: 900px; margin: 0 auto 8px !important; } .gradio-container { padding: 4px 0 !important; } .gradio-container [data-testid="tabs"], .gradio-container .tabs { margin-top: 0 !important; } .gradio-container [data-testid="tabitem"], .gradio-container .tabitem { padding-top: 4px !important; } .quick-links { text-align: center; padding: 8px 0; border: 1px solid #e5e7eb; border-radius: 8px; margin: 8px auto; max-width: 900px; } .quick-links a { margin: 0 12px; font-size: 14px; font-weight: 600; color: #3b82f6; text-decoration: none; } .quick-links a:hover { text-decoration: underline; } .prompt-grid { display: flex; flex-wrap: wrap; gap: 8px; margin-top: 6px; } .prompt-grid button { height: 40px !important; padding: 0 12px !important; border-radius: 8px !important; font-weight: 600 !important; font-size: 13px !important; letter-spacing: 0.2px; } #image_preview_vl, #image_preview_doc { height: 400px !important; overflow: auto; } #image_preview_vl img, #image_preview_doc img, #vis_image_doc img { width: 100% !important; height: auto !important; object-fit: contain !important; display: block; } #md_preview_vl, #md_preview_doc { max-height: 540px; min-height: 180px; overflow: auto; scrollbar-gutter: stable both-edges; } #md_preview_vl .prose, #md_preview_doc .prose { line-height: 1.7 !important; } #md_preview_vl .prose img, #md_preview_doc .prose img { display: block; margin: 0 auto; max-width: 100%; height: auto; } .notice { margin: 8px auto 0; max-width: 900px; padding: 10px 12px; border: 1px solid #e5e7eb; border-radius: 8px; background: #f8fafc; font-size: 14px; line-height: 1.6; } .notice strong { font-weight: 700; } .notice a { color: #3b82f6; text-decoration: none; } .notice a:hover { text-decoration: underline; } .checkbox-row .gradio-checkbox { flex-grow: 1; text-align: center; } """ with gr.Blocks(head=GOOGLE_FONTS_URL, css=custom_css, theme=gr.themes.Soft()) as demo: logo_data_url = image_to_base64_data_url(LOGO_IMAGE_PATH) if os.path.exists(LOGO_IMAGE_PATH) else "" gr.HTML(f"""
App Logo
""") gr.HTML("""
Heads up: The Hugging Face demo can be slow at times. For a faster experience, please try Baidu AI Studio or ModelScope.
""") gr.HTML("""
GitHub | Technical Report | Model | Official Website
""") with gr.Tabs(): # ===================== Document Parsing ===================== with gr.Tab("Document Parsing"): with gr.Row(): with gr.Column(scale=5): file_doc = gr.File(label="Upload Image", file_count="single", type="filepath", file_types=["image"]) preview_doc_html = gr.HTML(value="", elem_id="image_preview_doc", visible=False) gr.Markdown("_( Use this mode for recognizing full-page documents with structured layouts, such as reports, papers, or magazines.)_") gr.Markdown("💡 *To recognize a single, pre-cropped element (e.g., a table or formula), switch to the 'Element-level Recognition' tab for better results.*") example_url_doc = gr.State(value=None) with gr.Row(variant="panel"): with gr.Column(scale=2): btn_parse = gr.Button("Parse Document", variant="primary") with gr.Column(scale=3): with gr.Row(elem_classes=["checkbox-row"]): chart_parsing_switch = gr.Checkbox(label="Enable chart parsing", value=False, min_width=10) # ############### 修改点在这里 ############### doc_unwarping_switch = gr.Checkbox(label="Enable document unwarping", value=False, min_width=10) doc_orientation_switch = gr.Checkbox(label="Enable orientation classification", value=False, min_width=10) if complex_document_examples: complex_paths = [e[0] for e in complex_document_examples] complex_state = gr.State(complex_paths) gallery_complex = gr.Gallery( value=complex_paths, columns=4, height=400, preview=False, label=None, allow_preview=False ) def on_gallery_select_for_doc(paths, evt: gr.SelectData): idx = evt.index if isinstance(idx, (list, tuple)): idx = idx[0] try: url = paths[int(idx)] except Exception: raise gr.Error(f"Invalid index from gallery: {evt.index}") return url, update_preview_visibility(url) gallery_complex.select( fn=on_gallery_select_for_doc, inputs=[complex_state], outputs=[example_url_doc, preview_doc_html], ) # ===================== 更新日志模块 ===================== gr.Markdown("""

History Updates

""") with gr.Column(scale=7): with gr.Tabs(): with gr.Tab("Markdown Preview"): md_preview_doc = gr.Markdown("Please upload an image and click 'Parse Document'.", latex_delimiters=LATEX_DELIMS, elem_id="md_preview_doc") with gr.Tab("Visualization"): vis_image_doc = gr.HTML(label="Detection Visualization", elem_id="vis_image_doc") with gr.Tab("Markdown Source"): md_raw_doc = gr.Code(label="Markdown Source Code", language="markdown") def on_file_doc_change(fp): return None, update_preview_visibility(fp) file_doc.change(fn=on_file_doc_change, inputs=[file_doc], outputs=[example_url_doc, preview_doc_html]) def parse_doc_router(fp, example_url, use_chart, use_unwarping, use_orientation): src = fp if fp else example_url if not src: raise gr.Error("Please upload an image or pick an example first.") return handle_complex_doc(src, use_chart, use_unwarping, use_orientation) btn_parse.click(fn=parse_doc_router, inputs=[file_doc, example_url_doc, chart_parsing_switch, doc_unwarping_switch, doc_orientation_switch], outputs=[md_preview_doc, vis_image_doc, md_raw_doc]) # ===================== Element-level Recognition ===================== with gr.Tab("Element-level Recognition"): with gr.Row(): with gr.Column(scale=5): file_vl = gr.File(label="Upload Image", file_count="single", type="filepath", file_types=["image"]) preview_vl_html = gr.HTML(value="", elem_id="image_preview_vl", visible=False) gr.Markdown("_(Best for images with a **simple, single-column layout** (e.g., pure text), or for a **pre-cropped single element** like a table, formula, or chart.)_") gr.Markdown("Choose a recognition type:") with gr.Row(elem_classes=["prompt-grid"]): btn_ocr = gr.Button("Text Recognition", variant="secondary") btn_formula = gr.Button("Formula Recognition", variant="secondary") with gr.Row(elem_classes=["prompt-grid"]): btn_table = gr.Button("Table Recognition", variant="secondary") btn_chart = gr.Button("Chart Recognition", variant="secondary") example_url_vl = gr.State(value=None) if targeted_recognition_examples: targeted_paths = [e[0] for e in targeted_recognition_examples] targeted_state = gr.State(targeted_paths) gallery_targeted = gr.Gallery( value=targeted_paths, columns=4, height=400, preview=False, label=None, allow_preview=False ) def on_gallery_select_for_vl(paths, evt: gr.SelectData): idx = evt.index if isinstance(idx, (list, tuple)): idx = idx[0] try: url = paths[int(idx)] except Exception: raise gr.Error(f"Invalid index from gallery: {evt.index}") return url, update_preview_visibility(url) gallery_targeted.select( fn=on_gallery_select_for_vl, inputs=[targeted_state], outputs=[example_url_vl, preview_vl_html], ) with gr.Column(scale=7): with gr.Tabs(): with gr.Tab("Recognition Result"): md_preview_vl = gr.Markdown("Please upload an image and click a recognition type.", latex_delimiters=LATEX_DELIMS, elem_id="md_preview_vl") with gr.Tab("Raw Output"): md_raw_vl = gr.Code(label="Raw Output", language="markdown") def on_file_vl_change(fp): return None, update_preview_visibility(fp) file_vl.change(fn=on_file_vl_change, inputs=[file_vl], outputs=[example_url_vl, preview_vl_html]) def parse_vl_router(fp, example_url, prompt_choice): src = fp if fp else example_url if not src: raise gr.Error("Please upload an image or pick an example first.") return handle_targeted_recognition(src, prompt_choice) btn_ocr.click(fn=parse_vl_router, inputs=[file_vl, example_url_vl, gr.State("Text Recognition")], outputs=[md_preview_vl, md_raw_vl]) btn_formula.click(fn=parse_vl_router, inputs=[file_vl, example_url_vl, gr.State("Formula Recognition")], outputs=[md_preview_vl, md_raw_vl]) btn_table.click(fn=parse_vl_router, inputs=[file_vl, example_url_vl, gr.State("Table Recognition")], outputs=[md_preview_vl, md_raw_vl]) btn_chart.click(fn=parse_vl_router, inputs=[file_vl, example_url_vl, gr.State("Chart Recognition")], outputs=[md_preview_vl, md_raw_vl]) if __name__ == "__main__": port = int(os.getenv("PORT", "7860")) demo.queue(max_size=64).launch(server_name="0.0.0.0", server_port=port, share=False)