基于streamlit與olmOCR大模型實現的pdf提取工具
import os
import json
import subprocess
import pandas as pd
from pathlib import Path
import shutil
import time
import re
import streamlit as st# 創建工作目錄
WORKSPACE_DIR = "olmocr_workspace"
os.makedirs(WORKSPACE_DIR, exist_ok=True)def modify_html_for_better_display(html_content):"""修改HTML以便在Streamlit中更好地顯示"""if not html_content:return html_content# 增加容器寬度html_content = html_content.replace('<div class="container">', '<div class="container" style="max-width: 100%; width: 100%;">')# 增加文本大小html_content = html_content.replace('<style>', '<style>\nbody {font-size: 16px;}\n.text-content {font-size: 16px; line-height: 1.5;}\n')# 調整圖像和文本部分的大小比例html_content = html_content.replace('<div class="row">', '<div class="row" style="display: flex; flex-wrap: wrap;">')html_content = html_content.replace('<div class="col-md-6">', '<div class="col-md-6" style="flex: 0 0 50%; max-width: 50%; padding: 15px;">')# 增加頁面之間的間距html_content = html_content.replace('<div class="page">', '<div class="page" style="margin-bottom: 30px; border-bottom: 1px solid #ccc; padding-bottom: 20px;">')# 增加圖像大小html_content = re.sub(r'<img([^>]*)style="([^"]*)"', r'<img\1style="max-width: 100%; height: auto; \2"', html_content)# 添加縮放控制zoom_controls = """<div style="position: fixed; bottom: 20px; right: 20px; background: #fff; padding: 10px; border-radius: 5px; box-shadow: 0 0 10px rgba(0,0,0,0.2); z-index: 1000;"><button onclick="document.body.style.zoom = parseFloat(document.body.style.zoom || 1) + 0.1;" style="margin-right: 5px;">放大</button><button onclick="document.body.style.zoom = parseFloat(document.body.style.zoom || 1) - 0.1;">縮小</button></div>"""html_content = html_content.replace('</body>', f'{zoom_controls}</body>')return html_contentdef process_pdf(pdf_file):"""處理PDF文件并返回結果"""if pdf_file is None:return "請上傳PDF文件", "", None, None# 創建一個唯一的工作目錄timestamp = int(time.time())work_dir = os.path.join(WORKSPACE_DIR, f"job_{timestamp}")os.makedirs(work_dir, exist_ok=True)# 復制PDF文件pdf_path = os.path.join(work_dir, "input.pdf")with open(pdf_path, "wb") as f:f.write(pdf_file.read())# 構建命令并執行cmd = ["python", "-m", "olmocr.pipeline", work_dir, "--pdfs", pdf_path]try:# 執行命令，等待完成process = subprocess.run(cmd,stdout=subprocess.PIPE,stderr=subprocess.PIPE,text=True,check=True)# 命令輸出log_text = process.stdout# 檢查結果目錄results_dir = os.path.join(work_dir, "results")if not os.path.exists(results_dir):return f"處理完成，但未生成結果目錄\n\n日志輸出:\n{log_text}", "", None, None# 查找輸出文件output_files = list(Path(results_dir).glob("output_*.jsonl"))if not output_files:return f"處理完成，但未找到輸出文件\n\n日志輸出:\n{log_text}", "", None, None# 讀取JSONL文件output_file = output_files[0]with open(output_file, "r") as f:content = f.read().strip()if not content:return f"輸出文件為空\n\n日志輸出:\n{log_text}", "", None, None# 解析JSONresult = json.loads(content)extracted_text = result.get("text", "未找到文本內容")# 生成HTML預覽try:preview_cmd = ["python", "-m", "olmocr.viewer.dolmaviewer", str(output_file)]subprocess.run(preview_cmd, check=True)except Exception as e:log_text += f"\n生成HTML預覽失敗: {str(e)}"# 查找HTML文件html_files = list(Path("dolma_previews").glob("*.html"))html_content = ""if html_files:try:with open(html_files[0], "r", encoding="utf-8") as hf:html_content = hf.read()# 修改HTML以更好地顯示html_content = modify_html_for_better_display(html_content)except Exception as e:log_text += f"\n讀取HTML預覽失敗: {str(e)}"# 創建元數據表格metadata = result.get("metadata", {})meta_rows = []for key, value in metadata.items():meta_rows.append([key, value])df = pd.DataFrame(meta_rows, columns=["屬性", "值"])return log_text, extracted_text, html_content, dfexcept subprocess.CalledProcessError as e:return f"命令執行失敗: {e.stderr}", "", None, Noneexcept Exception as e:return f"處理過程中發生錯誤: {str(e)}", "", None, None# 設置Streamlit應用
st.title("一個黑客創業者：基于streamlit和olmOCR PDF文本提取工具")# 文件上傳
uploaded_file = st.file_uploader("上傳PDF文件", type="pdf")if uploaded_file is not None:st.button("處理PDF", on_click=lambda: process_and_display(uploaded_file))def process_and_display(pdf_file):log_text, extracted_text, html_content, df = process_pdf(pdf_file)st.header("處理日志")st.text(log_text)st.header("提取的文本")st.text_area("", value=extracted_text, height=200)if html_content:st.header("HTML預覽")st.components.v1.html(html_content, height=800, scrolling=True)if df is not None and not df.empty:st.header("文檔元數據")st.dataframe(df)
)
