Spaces:
Running
Running
| 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 = "<link href='https://fonts.googleapis.com/css2?family=Noto+Sans+SC:wght@400;700&display=swap' rel='stylesheet'>" | |
| 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（不下载、不转码），用于 <img src="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：直接用 <img src="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""" | |
| <div class="uploaded-image"> | |
| <img src="{src}" alt="Preview image" style="width:100%;height:100%;object-fit:contain;" loading="lazy"/> | |
| </div> | |
| """ | |
| 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) 构造一个可视化<img>（如果有） | |
| """ | |
| layout_results = (result or {}).get("layoutParsingResults", []) | |
| if not layout_results: | |
| return "No content was recognized.", "<p>No visualization available.</p>", "" | |
| 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 = "<p style='text-align:center; color:#888;'>No visualization image available.</p>" | |
| 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'<img src="{output_image_url}" alt="Detection Visualization" loading="lazy">' | |
| 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"""<div class="app-header"><img src="{logo_data_url}" alt="App Logo" style="max-height:10%; width: auto; margin: 10px auto; display: block;"></div>""") | |
| gr.HTML("""<div class="notice"><strong>Heads up:</strong> The Hugging Face demo can be slow at times. For a faster experience, please try <a href="https://aistudio.baidu.com/application/detail/98365" target="_blank" rel="noopener noreferrer">Baidu AI Studio</a> or <a href="https://modelscope.cn/studios/PaddlePaddle/PaddleOCR-VL_Online_Demo/summary" target="_blank" rel="noopener noreferrer">ModelScope</a>.</div>""") | |
| gr.HTML("""<div class="quick-links"><a href="https://github.com/PaddlePaddle/PaddleOCR" target="_blank">GitHub</a> | <a href="https://ernie.baidu.com/blog/publication/PaddleOCR-VL_Technical_Report.pdf" target="_blank">Technical Report</a> | <a href="https://huggingface.co/PaddlePaddle/PaddleOCR-VL" target="_blank">Model</a> | <a href="https://aistudio.baidu.com/paddleocr" target="_blank">Official Website</a></div>""") | |
| 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(""" | |
| <div class="notice"> | |
| <h3>History Updates</h3> | |
| <ul> | |
| <li> | |
| <strong>Oct 30, 2025:</strong> | |
| Added two advanced control options under the "Document Parsing" tab. These features were previously enabled by default (set to true) but are now user-configurable and default to false. | |
| <ul> | |
| <li><strong>Enable document unwarping:</strong> Corrects distortions in bent or poorly photographed documents.</li> | |
| <li><strong>Enable orientation classification:</strong> Automatically corrects the orientation of rotated or upside-down images.</li> | |
| </ul> | |
| </li> | |
| <li> | |
| <strong>Oct 16, 2025:</strong> Initial release of the demo. | |
| </li> | |
| </ul> | |
| </div> | |
| """) | |
| 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) |