app.py

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) -> 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,
        "layoutMergeBboxesMode": "union",
    }
    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) -> 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)
    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)
    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; }
"""

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></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"):
                        chart_parsing_switch = gr.Checkbox(label="Enable chart parsing", value=False, scale=1)
                        btn_parse = gr.Button("Parse Document", variant="primary", scale=2)

                    if complex_document_examples:
                        complex_paths = [e[0] for e in complex_document_examples]  # 这里是 List[str]
                        complex_state = gr.State(complex_paths)

                        gallery_complex = gr.Gallery(
                            value=complex_paths, columns=4, height=400,
                            preview=False, label=None, allow_preview=False
                        )

                        # 2) 回调：用 evt.index 到 paths（State）里取 URL
                        def on_gallery_select_for_doc(paths, evt: gr.SelectData):
                            # 某些版本 evt.index 可能是 (row, col) 或 list，做个兜底
                            idx = evt.index
                            if isinstance(idx, (list, tuple)):
                                # 常见是一个 int；如果是 (row, col) 形式，通常线性下标 == row
                                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)

                        # 3) 绑定：把 State 作为 inputs 传给回调，outputs 写入 example_url_doc 和预览 HTML
                        gallery_complex.select(
                            fn=on_gallery_select_for_doc,
                            inputs=[complex_state],
                            outputs=[example_url_doc, preview_doc_html],
                        )


                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):
                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)

            btn_parse.click(fn=parse_doc_router, inputs=[file_doc, example_url_doc, chart_parsing_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]  # List[str]
                        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)