|
|
import base64 |
|
|
import io |
|
|
import json |
|
|
import os |
|
|
import re |
|
|
import tempfile |
|
|
import zipfile |
|
|
from datetime import datetime |
|
|
from pathlib import Path |
|
|
|
|
|
import gradio as gr |
|
|
import requests |
|
|
from PIL import Image |
|
|
import pdf2image |
|
|
|
|
|
API_URL = "https://t707h6d9q6oftbx3.aistudio-app.com/layout-parsing" |
|
|
TOKEN = os.getenv("API_TOKEN") |
|
|
|
|
|
|
|
|
CSS = """ |
|
|
:root { |
|
|
--sand-color: #FAF9F6; |
|
|
--white: #ffffff; |
|
|
--shadow: 0 4px 6px rgba(0, 0, 0, 0.1); |
|
|
--text-color: #FAF9F6; |
|
|
--black:#000000; |
|
|
--link-hover: #2b6cb0; |
|
|
--content-width: 1200px; |
|
|
} |
|
|
|
|
|
body { |
|
|
display: flex; |
|
|
justify-content: center; |
|
|
background-color: var(--sand-color); |
|
|
color: var(--text-color); |
|
|
font-family: Arial, sans-serif; |
|
|
} |
|
|
|
|
|
.gradio-container { |
|
|
max-width: var(--content-width) !important; |
|
|
width: 100% !important; |
|
|
margin: 20px auto; |
|
|
padding: 20px; |
|
|
background-color: var(--white); |
|
|
} |
|
|
|
|
|
#component-0, |
|
|
#tabs, |
|
|
#settings { |
|
|
background-color: var(--white) !important; |
|
|
padding: 15px; |
|
|
} |
|
|
|
|
|
.upload-section { |
|
|
width: 100%; |
|
|
margin: 0 auto 30px; |
|
|
padding: 20px; |
|
|
background-color: var(--sand-color) !important; |
|
|
border-radius: 8px; |
|
|
box-shadow: var(--shadow); |
|
|
} |
|
|
|
|
|
.center-content { |
|
|
display: flex; |
|
|
flex-direction: column; |
|
|
align-items: center; |
|
|
text-align: center; |
|
|
margin-bottom: 20px; |
|
|
} |
|
|
|
|
|
.header { |
|
|
margin-bottom: 30px; |
|
|
width: 100%; |
|
|
} |
|
|
|
|
|
.logo-container { |
|
|
width: 100%; |
|
|
margin-bottom: 20px; |
|
|
} |
|
|
|
|
|
.logo-img { |
|
|
width: 100%; |
|
|
max-width: var(--content-width); |
|
|
margin: 0 auto; |
|
|
display: block; |
|
|
} |
|
|
|
|
|
.nav-bar { |
|
|
display: flex; |
|
|
justify-content: center; |
|
|
background-color: var(--white); |
|
|
padding: 15px 0; |
|
|
box-shadow: var(--shadow); |
|
|
margin-bottom: 20px; |
|
|
} |
|
|
|
|
|
.nav-links { |
|
|
display: flex; |
|
|
gap: 30px; |
|
|
width: 100%; |
|
|
justify-content: center; |
|
|
} |
|
|
|
|
|
.nav-link { |
|
|
color: var(--black); |
|
|
text-decoration: none; |
|
|
font-weight: bold; |
|
|
font-size: 24px; |
|
|
transition: color 0.2s; |
|
|
} |
|
|
|
|
|
.nav-link:hover { |
|
|
color: var(--link-hover); |
|
|
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; |
|
|
border: none !important; |
|
|
border-radius: 4px; |
|
|
padding: 8px 16px; |
|
|
} |
|
|
button:hover { |
|
|
opacity: 0.8 !important; |
|
|
} |
|
|
|
|
|
.radio-group { |
|
|
margin-bottom: 15px !important; |
|
|
} |
|
|
|
|
|
.file-download { |
|
|
margin-top: 15px !important; |
|
|
} |
|
|
.loader { |
|
|
border: 5px solid #f3f3f3; |
|
|
border-top: 5px solid #3498db; |
|
|
border-radius: 50%; |
|
|
width: 50px; |
|
|
height: 50px; |
|
|
animation: spin 1s linear infinite; |
|
|
margin: 20px auto; |
|
|
} |
|
|
|
|
|
@keyframes spin { |
|
|
0% { transform: rotate(0deg); } |
|
|
100% { transform: rotate(360deg); } |
|
|
} |
|
|
|
|
|
.loader-container { |
|
|
text-align: center; |
|
|
margin: 20px 0; |
|
|
} |
|
|
""" |
|
|
|
|
|
def clean_markdown_text(text): |
|
|
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): |
|
|
try: |
|
|
images = pdf2image.convert_from_path(pdf_path) |
|
|
return [image for image in images] |
|
|
except: |
|
|
return None |
|
|
|
|
|
|
|
|
def process_file(file_path, file_type): |
|
|
try: |
|
|
with open(file_path, "rb") as f: |
|
|
file_bytes = f.read() |
|
|
|
|
|
file_data = base64.b64encode(file_bytes).decode("ascii") |
|
|
headers = { |
|
|
"Authorization": f"token {TOKEN}", |
|
|
"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 |
|
|
) |
|
|
response.raise_for_status() |
|
|
|
|
|
result = response.json() |
|
|
layout_results = result.get("result", {}).get("layoutParsingResults", []) |
|
|
|
|
|
markdown_contents = [] |
|
|
clean_markdown_contents = [] |
|
|
for res in layout_results: |
|
|
markdown = res.get("markdown", {}) |
|
|
if isinstance(markdown, str): |
|
|
original = markdown |
|
|
elif isinstance(markdown, dict): |
|
|
original = markdown.get("text", "") |
|
|
|
|
|
markdown_contents.append(original) |
|
|
clean_markdown_contents.append(clean_markdown_text(original)) |
|
|
|
|
|
if file_type == "pdf": |
|
|
images = pdf_to_images(file_path) |
|
|
else: |
|
|
images = [Image.open(file_path)] |
|
|
|
|
|
return { |
|
|
"original_file": file_path, |
|
|
"markdown_contents": markdown_contents, |
|
|
"clean_markdown_contents": clean_markdown_contents, |
|
|
"pdf_images": images, |
|
|
"api_response": result |
|
|
} |
|
|
|
|
|
except Exception as e: |
|
|
raise gr.Error(f"Error processing file: {str(e)}") |
|
|
|
|
|
|
|
|
def create_zip_file(results): |
|
|
try: |
|
|
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") |
|
|
zip_filename = f"analysis_results_{timestamp}.zip" |
|
|
|
|
|
temp_dir = tempfile.mkdtemp() |
|
|
zip_path = os.path.join(temp_dir, zip_filename) |
|
|
|
|
|
with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf: |
|
|
original_path = results.get("original_file", "") |
|
|
if original_path and Path(original_path).exists(): |
|
|
zipf.write(original_path, f"original/{Path(original_path).name}") |
|
|
|
|
|
markdowns = results.get("markdown_contents", []) |
|
|
for i, md_content in enumerate(markdowns): |
|
|
if md_content: |
|
|
zipf.writestr(f"markdown/original/markdown_{i + 1}.md", md_content) |
|
|
|
|
|
api_response = results.get("api_response", {}) |
|
|
zipf.writestr("api_response.json", json.dumps(api_response, indent=2, ensure_ascii=False)) |
|
|
|
|
|
return zip_path |
|
|
|
|
|
except Exception as e: |
|
|
raise gr.Error(f"Error creating ZIP file: {str(e)}") |
|
|
|
|
|
|
|
|
def export_markdown(results): |
|
|
try: |
|
|
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"original_markdown_{timestamp}.md" |
|
|
content = "\n\n".join(markdowns) |
|
|
|
|
|
temp_dir = tempfile.mkdtemp() |
|
|
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() |
|
|
|
|
|
|
|
|
with gr.Column(elem_classes=["logo-container"]): |
|
|
gr.Image("pp-structurev3.png", elem_classes=["logo-img"], show_label=False) |
|
|
|
|
|
|
|
|
with gr.Row(elem_classes=["nav-bar"]): |
|
|
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">Homepage</a> |
|
|
</div> |
|
|
""") |
|
|
|
|
|
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" |
|
|
) |
|
|
process_btn = gr.Button("Analyze document", variant="primary") |
|
|
|
|
|
loading_spinner = gr.Column( |
|
|
visible=False, |
|
|
elem_classes=["loader-container"] |
|
|
) |
|
|
with loading_spinner: |
|
|
gr.HTML(""" |
|
|
<div class="loader"></div> |
|
|
<p>Wait...</p> |
|
|
""") |
|
|
|
|
|
with gr.Row(elem_classes=["result-container"]): |
|
|
with gr.Column(elem_classes=["pdf-preview"]): |
|
|
gr.Markdown("### Original document preview") |
|
|
pdf_display = gr.Gallery( |
|
|
label="PDF page", |
|
|
show_label=False, |
|
|
elem_classes=["gallery-container"] |
|
|
) |
|
|
|
|
|
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 |
|
|
) |
|
|
markdown_display = gr.HTML(label="Analysis Results") |
|
|
with gr.Column(elem_classes=["download-section"]): |
|
|
gr.Markdown("### Result Export") |
|
|
with gr.Row(): |
|
|
download_md_btn = gr.Button("Download Original Markdown", variant="secondary") |
|
|
download_all_btn = gr.Button("Download Complete Analysis Results (ZIP)", variant="primary") |
|
|
download_file = gr.File(visible=False, label="Download file", elem_classes=["file-download"]) |
|
|
|
|
|
def toggle_spinner(): |
|
|
return gr.update(visible=True) |
|
|
|
|
|
|
|
|
def hide_spinner(): |
|
|
return gr.update(visible=False) |
|
|
|
|
|
|
|
|
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] |
|
|
).success( |
|
|
lambda res: res["pdf_images"] if res and res.get("pdf_images") else [], |
|
|
inputs=[results_state], |
|
|
outputs=[pdf_display] |
|
|
).success( |
|
|
lambda res: res["markdown_contents"][0] if res and res.get("markdown_contents") else "", |
|
|
inputs=[results_state], |
|
|
outputs=[markdown_display] |
|
|
) |
|
|
|
|
|
display_mode.change( |
|
|
lambda mode, res: ( |
|
|
res["markdown_contents"][0] if mode == "原始Markdown" |
|
|
else res["clean_markdown_contents"][0] |
|
|
) if res else "", |
|
|
inputs=[display_mode, results_state], |
|
|
outputs=[markdown_display] |
|
|
) |
|
|
|
|
|
download_md_btn.click( |
|
|
export_markdown, |
|
|
inputs=[results_state], |
|
|
outputs=[download_file] |
|
|
).then( |
|
|
lambda x: gr.update(visible=True), |
|
|
inputs=[download_file], |
|
|
outputs=[download_file] |
|
|
) |
|
|
|
|
|
download_all_btn.click( |
|
|
create_zip_file, |
|
|
inputs=[results_state], |
|
|
outputs=[download_file] |
|
|
).then( |
|
|
lambda x: gr.update(visible=True), |
|
|
inputs=[download_file], |
|
|
outputs=[download_file] |
|
|
) |
|
|
|
|
|
if __name__ == "__main__": |
|
|
demo.launch(share=True) |
