[Feat] Big update

app.py

@@ -1,66 +1,37 @@
+import atexit
 import base64
 import io
 import json
 import os
-import re
 import tempfile
+import uuid
 import zipfile
-import shutil
-import atexit
-from datetime import datetime
 from pathlib import Path
 
 import gradio as gr
 import requests
 from PIL import Image
 
-try:
-    import pdf2image
-    PDF2IMAGE_AVAILABLE = True
-except ImportError:
-    PDF2IMAGE_AVAILABLE = False
-
-try:
-    import fitz  # PyMuPDF
-    PYGMUPDF_AVAILABLE = True
-except ImportError:
-    PYGMUPDF_AVAILABLE = False
-
 # API Configuration
-API_URL = "https://t707h6d9q6oftbx3.aistudio-app.com/layout-parsing"
+API_URL = "https://cf38vaydqdl2l4p2.aistudio-hub.baidu.com/layout-parsing"
 TOKEN = os.getenv("API_TOKEN")
 
-# Temporary directory management
-temp_dirs = []
-
-def cleanup():
-    """Clean up temporary directories"""
-    for dir_path in temp_dirs:
-        try:
-            shutil.rmtree(dir_path)
-        except:
-            pass
-
-atexit.register(cleanup)
+LOGO_PATH = Path(__file__).parent / "pp-structurev3.png"
+with open(LOGO_PATH, "rb") as image_file:
+    LOGO_BASE64 = (
+        f"data:image/png;base64,{base64.b64encode(image_file.read()).decode('utf-8')}"
+    )
 
-def image_to_base64(image_path):
-    """Convert image to base64 encoding"""
-    if not image_path or not Path(image_path).exists():
-        return ""
-    with open(image_path, "rb") as image_file:
-        return f"data:image/png;base64,{base64.b64encode(image_file.read()).decode('utf-8')}"
+TEMP_DIR = tempfile.TemporaryDirectory()
+atexit.register(TEMP_DIR.cleanup)
 
-# Get current directory
-current_dir = Path(__file__).parent
-logo_path = current_dir / "pp-structurev3.png"  
-logo_base64 = image_to_base64(logo_path)
 
 CSS = """
 :root {
-    --sand-color: #FAF9F6;    
-    --white: #ffffff;          
-    --shadow: 0 4px 6px rgba(0, 0, 0, 0.1); 
-    --text-color: #F3F4F7;     
+    --sand-color: #FAF9F6;
+    --white: #ffffff;
+    --shadow: 0 4px 6px rgba(0, 0, 0, 0.1);
+    --text-color: #F3F4F7;
     --black:#000000;
     --link-hover: #2b6cb0;
     --content-width: 1200px;
@@ -82,9 +53,9 @@ body {
     background-color: var(--white);
 }
 
-#component-0, 
-#tabs, 
-#settings { 
+#component-0,
+#tabs,
+#settings {
     background-color: var(--white) !important;
     padding: 15px;
 }
@@ -152,33 +123,6 @@ body {
     text-decoration: none;
 }
 
-.result-container {
-    display: flex;
-    gap: 20px;
-    margin-bottom: 30px;
-    width: 100%;
-}
-
-.pdf-preview {
-    flex: 1;
-    min-width: 0;
-}
-
-.markdown-result {
-    flex: 1;
-    min-width: 0;
-}
-
-.gallery-container {
-    width: 100% !important;
-}
-
-.gallery-item {
-    width: 100% !important;
-    height: auto !important;
-    aspect-ratio: auto !important;
-}
-
 button {
     background-color: var(--text-color) !important;
     color: var(--black) !important;
@@ -190,10 +134,6 @@ button:hover {
     opacity: 0.8 !important;
 }
 
-.radio-group {
-    margin-bottom: 15px !important;
-}
-
 .file-download {
     margin-top: 15px !important;
 }
@@ -216,100 +156,40 @@ button:hover {
     text-align: center;
     margin: 20px 0;
 }
+"""
 
-/* PDF Viewer specific styles */
-.pdf-viewer-container {
-    width: 100%;
-    height: 600px;
-    border: 1px solid #ddd;
-    margin-top: 15px;
-    background-color: #f9f9f9;
-    display: flex;
-    justify-content: center;
-    align-items: center;
-}
+MAX_NUM_PAGES = 10
 
-.pdf-viewer-container embed {
-    width: 100%;
-    height: 100%;
-}
 
-.no-preview-message {
-    color: #666;
-    font-size: 16px;
-    text-align: center;
-    padding: 20px;
-}
-"""
+def url_to_bytes(url, *, timeout=10):
+    resp = requests.get(url, timeout=timeout)
+    resp.raise_for_status()
+    return resp.content
 
-def clean_markdown_text(text):
-    """Clean markdown text from HTML tags and excessive newlines"""
-    if not text:
-        return ""
-    text = re.sub(r'<[^>]+>', '', text)
-    text = re.sub(r'\n{3,}', '\n\n', text)
-    return text.strip()
-
-def pdf_to_images(pdf_path, dpi=150):
-    """Convert PDF to list of images with fallback methods"""
-    images = []
-    
-    if PDF2IMAGE_AVAILABLE:
-        try:
-            images = pdf2image.convert_from_path(pdf_path, dpi=dpi)
-            return images
-        except Exception as e:
-            print(f"pdf2image conversion failed: {str(e)}")
-    
-    if PYGMUPDF_AVAILABLE:
-        try:
-            doc = fitz.open(pdf_path)
-            for page in doc:
-                pix = page.get_pixmap(dpi=dpi)
-                img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
-                images.append(img)
-            return images
-        except Exception as e:
-            print(f"PyMuPDF conversion failed: {str(e)}")
-    
-    return None
-
-def create_pdf_preview(pdf_path):
-    """Create PDF preview HTML with embedded viewer"""
-    if not pdf_path or not Path(pdf_path).exists():
-        return '<div class="no-preview-message">No PDF file available</div>'
-    
-    try:
-        # Convert PDF to base64 for embedding
-        with open(pdf_path, "rb") as f:
-            pdf_bytes = f.read()
-        pdf_base64 = base64.b64encode(pdf_bytes).decode("ascii")
-        
-        return f"""
-        <div class="pdf-viewer-container">
-            <embed 
-                src="data:application/pdf;base64,{pdf_base64}" 
-                type="application/pdf" 
-                width="100%" 
-                height="100%"
-            >
-        </div>
-        """
-    except Exception as e:
-        print(f"Failed to create PDF preview: {str(e)}")
-        return '<div class="no-preview-message">PDF preview generation failed</div>'
 
-def process_file(file_path, file_type):
+def bytes_to_image(image_bytes):
+    return Image.open(io.BytesIO(image_bytes))
+
+
+def embed_images_into_markdown_text(markdown_text, markdown_images):
+    for img_path, img_url in markdown_images.items():
+        # HACK
+        markdown_text = markdown_text.replace(
+            f'<img src="{img_path}"', f'<img src="{img_url}"'
+        )
+    return markdown_text
+
+
+def process_file(file_path):
     """Process uploaded file with API"""
     try:
         if not file_path:
             raise ValueError("Please upload a file first")
-            
-        if file_type == "pdf" and not str(file_path).lower().endswith('.pdf'):
-            raise ValueError("Please upload a valid PDF file")
-            
-        if file_type == "image" and not str(file_path).lower().endswith(('.jpg', '.jpeg', '.png')):
-            raise ValueError("Please upload a valid image file (JPG/JPEG/PNG)")
+
+        if Path(file_path).suffix == ".pdf":
+            file_type = "pdf"
+        else:
+            file_type = "image"
 
         # Read file content
         with open(file_path, "rb") as f:
@@ -319,14 +199,14 @@ def process_file(file_path, file_type):
         file_data = base64.b64encode(file_bytes).decode("ascii")
         headers = {
             "Authorization": f"token {TOKEN}",
-            "Content-Type": "application/json"
+            "Content-Type": "application/json",
         }
 
         response = requests.post(
             API_URL,
             json={"file": file_data, "fileType": 0 if file_type == "pdf" else 1},
             headers=headers,
-            timeout=60
+            timeout=1000,
         )
         response.raise_for_status()
 
@@ -334,30 +214,38 @@ def process_file(file_path, file_type):
         result = response.json()
         layout_results = result.get("result", {}).get("layoutParsingResults", [])
 
-        markdown_contents = []
-        clean_markdown_contents = []
+        layout_ordering_images = []
+        markdown_texts = []
+        markdown_images = []
+        markdown_content_list = []
+        input_images = []
         for res in layout_results:
-            markdown = res.get("markdown", {})
-            original = markdown if isinstance(markdown, str) else markdown.get("text", "")
-            markdown_contents.append(original)
-            clean_markdown_contents.append(clean_markdown_text(original))
-
-        # Generate preview content
-        if file_type == "pdf":
-            images = pdf_to_images(file_path)
-            pdf_preview = create_pdf_preview(file_path)
-        else:
-            images = [Image.open(file_path)]
-            pdf_preview = '<div class="no-preview-message">Image file preview</div>'
+            layout_ordering_images.append(
+                url_to_bytes(res["outputImages"]["layout_order_res"])
+            )
+            markdown = res["markdown"]
+            markdown_text = markdown["text"]
+            markdown_texts.append(markdown_text)
+            img_path_to_url = markdown["images"]
+            img_path_to_bytes = {}
+            for path, url in img_path_to_url.items():
+                img_path_to_bytes[path] = url_to_bytes(url)
+            markdown_images.append(img_path_to_bytes)
+            input_images.append(url_to_bytes(res["inputImage"]))
+            markdown_content = embed_images_into_markdown_text(
+                markdown_text, img_path_to_url
+            )
+            markdown_content_list.append(markdown_content)
 
         return {
             "original_file": file_path,
             "file_type": file_type,
-            "markdown_contents": markdown_contents,
-            "clean_markdown_contents": clean_markdown_contents,
-            "pdf_images": images,
-            "pdf_preview": pdf_preview,
-            "api_response": result
+            "layout_ordering_images": layout_ordering_images,
+            "markdown_texts": markdown_texts,
+            "markdown_images": markdown_images,
+            "markdown_content_list": markdown_content_list,
+            "input_images": input_images,
+            "api_response": result,
         }
 
     except requests.exceptions.RequestException as e:
@@ -365,148 +253,115 @@ def process_file(file_path, file_type):
     except Exception as e:
         raise gr.Error(f"Error processing file: {str(e)}")
 
-def create_zip_file(results):
+
+def export_full_results(results):
     """Create ZIP file with all analysis results"""
     try:
         if not results:
             raise ValueError("No results to export")
-            
-        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
-        zip_filename = f"analysis_results_{timestamp}.zip"
-
-        temp_dir = tempfile.mkdtemp()
-        temp_dirs.append(temp_dir)
-        zip_path = os.path.join(temp_dir, zip_filename)
-
-        with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
-            # Add original file
-            original_path = results.get("original_file", "")
-            if original_path and Path(original_path).exists():
-                zipf.write(original_path, f"original/{Path(original_path).name}")
-
-            # Add markdown content
-            for i, (orig_md, clean_md) in enumerate(zip(
-                results.get("markdown_contents", []),
-                results.get("clean_markdown_contents", [])
-            )):
-                if orig_md:
-                    zipf.writestr(f"markdown/original/page_{i+1}.md", orig_md)
-                if clean_md:
-                    zipf.writestr(f"markdown/clean/page_{i+1}.md", clean_md)
+
+        filename = Path(results["original_file"]).stem + f"_{uuid.uuid4().hex}.zip"
+        zip_path = Path(TEMP_DIR.name, filename)
+
+        with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zipf:
+            for i, img_bytes in enumerate(results["layout_ordering_images"]):
+                zipf.writestr(f"layout_ordering_images/page_{i+1}.jpg", img_bytes)
+
+            for i, (md_text, md_imgs) in enumerate(
+                zip(
+                    results["markdown_texts"],
+                    results["markdown_images"],
+                )
+            ):
+                zipf.writestr(f"markdown/page_{i+1}.md", md_text)
+                for img_path, img_bytes in md_imgs.items():
+                    zipf.writestr(f"markdown/{img_path}", img_bytes)
 
             # Add API response
             api_response = results.get("api_response", {})
-            zipf.writestr("api_response.json", json.dumps(api_response, indent=2, ensure_ascii=False))
+            zipf.writestr(
+                "api_response.json",
+                json.dumps(api_response, indent=2, ensure_ascii=False),
+            )
 
-            # Add PDF images if available
-            if results.get("file_type") == "pdf" and results.get("pdf_images"):
-                for i, img in enumerate(results["pdf_images"]):
-                    img_path = os.path.join(temp_dir, f"page_{i+1}.jpg")
-                    img.save(img_path, "JPEG", quality=85)
-                    zipf.write(img_path, f"images/page_{i+1}.jpg")
+            for i, img_bytes in enumerate(results["input_images"]):
+                zipf.writestr(f"input_images/page_{i+1}.jpg", img_bytes)
 
-        return zip_path
+        return str(zip_path)
 
     except Exception as e:
         raise gr.Error(f"Error creating ZIP file: {str(e)}")
 
-def export_markdown(results):
-    """Export markdown content to file"""
-    try:
-        if not results:
-            raise ValueError("No results to export")
-            
-        markdowns = results.get("markdown_contents", [])
-        if not markdowns:
-            raise gr.Error("No markdown content to export")
-
-        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
-        filename = f"markdown_export_{timestamp}.md"
-        content = "\n\n".join(markdowns)
-
-        temp_dir = tempfile.mkdtemp()
-        temp_dirs.append(temp_dir)
-        file_path = os.path.join(temp_dir, filename)
-
-        with open(file_path, 'w', encoding='utf-8') as f:
-            f.write(content)
-
-        return file_path
-
-    except Exception as e:
-        raise gr.Error(f"Error exporting markdown: {str(e)}")
 
 with gr.Blocks(css=CSS, title="Document Analysis System") as demo:
     results_state = gr.State()
 
     # Header with logo
     with gr.Column(elem_classes=["logo-container"]):
-        gr.HTML(f'<img src="{logo_base64}" class="logo-img">')
+        gr.HTML(f'<img src="{LOGO_BASE64}" class="logo-img">')
 
     # Navigation bar
     with gr.Row(elem_classes=["nav-bar"]):
-        gr.HTML("""
+        gr.HTML(
+            """
         <div class="nav-links">
             <a href="https://github.com/PaddlePaddle/PaddleOCR" class="nav-link" target="_blank">GitHub</a>
             <a href="https://paddleocr.ai" class="nav-link" target="_blank">paddleocr.ai</a>
         </div>
-        """)
+        """
+        )
 
     # Upload section
     with gr.Column(elem_classes=["upload-section"]):
-        file_type = gr.Radio(
-            ["pdf", "image"],
-            label="File Type",
-            value="pdf",
-            interactive=True
-        )
         file_input = gr.File(
             label="Upload Document",
             file_types=[".pdf", ".jpg", ".jpeg", ".png"],
-            type="filepath"
+            type="filepath",
         )
         process_btn = gr.Button("Analyze Document", variant="primary")
-
-        loading_spinner = gr.Column(
-            visible=False,
-            elem_classes=["loader-container"]
+        gr.Markdown(
+            f"*Please note that only the first {MAX_NUM_PAGES} pages will be processed.*"
         )
+
+        loading_spinner = gr.Column(visible=False, elem_classes=["loader-container"])
         with loading_spinner:
-            gr.HTML("""
+            gr.HTML(
+                """
             <div class="loader"></div>
             <p>Processing, please wait...</p>
-            """)
-
-    # Results display section
-    with gr.Row(elem_classes=["result-container"]):
-        with gr.Column(elem_classes=["pdf-preview"]):
-            gr.Markdown("### Original Document Preview")
-            pdf_preview = gr.HTML(label="PDF Preview")
-            pdf_gallery = gr.Gallery(
-                label="PDF Pages",
-                show_label=False,
-                elem_classes=["gallery-container"],
-                columns=[1],
-                object_fit="contain",
-                visible=False
+            """
             )
 
-        with gr.Column(elem_classes=["markdown-result"]):
-            with gr.Row(elem_classes=["radio-group"]):
-                display_mode = gr.Radio(
-                    ["Original Markdown", "Cleaned Text"],
-                    label="Display Mode",
-                    value="Original Markdown",
-                    interactive=True
+    # Results display section
+    with gr.Column():
+        gr.Markdown("### Results")
+        layout_ordering_images = []
+        markdown_display_list = []
+        for i in range(MAX_NUM_PAGES):
+            with gr.Row():
+                layout_ordering_images.append(
+                    gr.Image(
+                        label=f"Layout Ordering Image {i}",
+                        show_label=True,
+                        visible=False,
+                    )
+                )
+                markdown_display_list.append(
+                    gr.Markdown(
+                        visible=False,
+                        container=True,
+                        show_copy_button=True,
+                        latex_delimiters=[
+                            {"left": "$$", "right": "$$", "display": True},
+                            {"left": "$", "right": "$", "display": False},
+                        ],
+                    )
                 )
-            markdown_display = gr.Markdown(label="Analysis Results")
 
     # Download section
     with gr.Column(elem_classes=["download-section"]):
         gr.Markdown("### Result Export")
-        with gr.Row():
-            download_md_btn = gr.Button("Download Markdown", variant="secondary")
-            download_all_btn = gr.Button("Download Full Results (ZIP)", variant="primary")
+        download_all_btn = gr.Button("Download Full Results (ZIP)", variant="primary")
         download_file = gr.File(visible=False, label="Download File")
 
     # Interaction logic
@@ -516,81 +371,38 @@ with gr.Blocks(css=CSS, title="Document Analysis System") as demo:
     def hide_spinner():
         return gr.update(visible=False)
 
-    # In update_display()
     def update_display(results):
-        if not results:
-            return [
-                gr.update(value='<div class="no-preview-message">No file to display</div>'),
-                gr.update(visible=False),
-                gr.update(value="No content"),
-                gr.update(value=[])
-            ]
-    
-        images = results.get("pdf_images", [])
-        show_gallery = bool(images)
-        display_content = results["markdown_contents"][0] if results.get("markdown_contents") else "No content"
-    
-        return [
-            gr.update(value='<div class="no-preview-message">Preview rendered as images</div>'),
-            gr.update(visible=show_gallery),
-            gr.update(value=display_content),
-            gr.update(value=images if show_gallery else [])
-        ]
-
-
-    process_btn.click(
-        toggle_spinner,
-        outputs=[loading_spinner]
-    ).then(
-        process_file,
-        inputs=[file_input, file_type],
-        outputs=[results_state]
-    ).then(
-        hide_spinner,
-        outputs=[loading_spinner]
-    ).then(
+        ret_img = []
+        ret_cont = []
+        cnt = 0
+        for img, cont in zip(
+            results["layout_ordering_images"], results["markdown_content_list"]
+        ):
+            ret_img.append(gr.update(value=bytes_to_image(img), visible=True))
+            ret_cont.append(gr.update(value=cont, visible=True))
+            cnt += 1
+        for _ in range(cnt, MAX_NUM_PAGES):
+            ret_img.append(gr.update(visible=False))
+            ret_cont.append(gr.update(visible=False))
+        return ret_img + ret_cont
+
+    process_btn.click(toggle_spinner, outputs=[loading_spinner]).then(
+        process_file, inputs=[file_input], outputs=[results_state]
+    ).then(hide_spinner, outputs=[loading_spinner]).then(
         update_display,
         inputs=[results_state],
-        outputs=[pdf_preview, pdf_gallery, markdown_display, pdf_gallery]
-    )
-
-    display_mode.change(
-        lambda mode, res: (
-            res["markdown_contents"][0] if mode == "Original Markdown"
-            else res["clean_markdown_contents"][0]
-        ) if res and res.get("markdown_contents") else "No content",
-        inputs=[display_mode, results_state],
-        outputs=[markdown_display]
-    )
-
-    download_md_btn.click(
-        export_markdown,
-        inputs=[results_state],
-        outputs=[download_file]
-    ).then(
-        lambda: gr.update(visible=True),
-        outputs=[download_file]
+        outputs=layout_ordering_images + markdown_display_list,
     )
 
     download_all_btn.click(
-        create_zip_file,
-        inputs=[results_state],
-        outputs=[download_file]
-    ).then(
-        lambda: gr.update(visible=True),
-        outputs=[download_file]
-    )
+        export_full_results, inputs=[results_state], outputs=[download_file]
+    ).success(lambda: gr.update(visible=True), outputs=[download_file])
 
-if __name__ == "__main__":
-    # Check dependencies
-    if not PDF2IMAGE_AVAILABLE:
-        print("Warning: pdf2image not available, PDF to image conversion limited")
-    if not PYGMUPDF_AVAILABLE:
-        print("Warning: PyMuPDF not available, PDF fallback conversion disabled")
 
+if __name__ == "__main__":
     demo.launch(
         server_name="0.0.0.0",
         server_port=7860,
         share=True,
-        favicon_path=str(logo_path) if logo_path.exists() else None
-    )
+        favicon_path=LOGO_PATH,
+    )

requirements.txt

@@ -1,6 +1,3 @@
-gradio>=3.0
+gradio>=5
 requests
 pillow
-pdf2image
-python-dotenv
-PyMuPDF