寫了個腳本將pdf轉markdown

看到有人需要將掃描pdf文檔轉markdown，想起之前寫的一個小工具。
這個腳本是為了將pdf轉成markdown，只需要申請一個智譜的api key，并填到config里，使用的模型是4v flash，免費的，所以可以放心使用。
效果如下圖:
在這里插入圖片描述

腳本里的提示詞可以根據個人需要進行修改。以下是原始代碼：

#!/usr/bin/env python3
# -*- coding: utf-8 -*-"""
PDF轉Markdown自動化系統
功能：監控input/目錄下的PDF文件，轉換為Markdown格式
作者：您的專屬程序員
日期：2025-04-03
版本：2.0.0
"""import base64
import logging
import time
import json
import os
import fitz  # PyMuPDF
from pathlib import Path
from typing import Optional, Dict, Any, List, Generator
from zhipuai import ZhipuAI
from zhipuai.core._errors import ZhipuAIError# 配置日志系統
logging.basicConfig(level=logging.INFO,format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',handlers=[logging.FileHandler('pdf2md.log'),logging.StreamHandler()]
)
logger = logging.getLogger(__name__)class GLM4VTester:"""GLM-4V 模型測試工具類"""def __init__(self, api_key: str, model_name: str = "glm-4v-flash"):self.client = ZhipuAI(api_key=api_key)self.model_name = model_nameself.total_tokens = 0self.total_requests = 0self.total_time = 0.0def analyze_image(self, image_path: str, prompt: str = "你是一個OCR助手，請把圖中內容按原有格式輸出出來，如果有公式則輸出為LaTeX") -> Dict[str, Any]:"""分析圖片內容:param image_path: 圖片路徑:param prompt: 提示詞:return: API響應結果"""start_time = time.time()# 讀取圖片并轉為base64with open(image_path, "rb") as image_file:base64_image = base64.b64encode(image_file.read()).decode('utf-8')# 調用APIresponse = self.client.chat.completions.create(model=self.model_name,messages=[{"role": "user", "content": [{"type": "text", "text": prompt},{"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{base64_image}"}}]}])# 更新統計信息elapsed_time = time.time() - start_timeself.total_requests += 1self.total_time += elapsed_timeif hasattr(response, 'usage') and response.usage:self.total_tokens += response.usage.total_tokenslogger.info(f"API請求完成，耗時: {elapsed_time:.2f}秒")return {"response": response, "time": elapsed_time}def generate_markdown_report(self, image_path: str, result: Dict[str, Any], output_path: str) -> str:"""生成Markdown格式的分析報告:param image_path: 原始圖片路徑:param result: API響應結果:param output_path: 輸出文件路徑:return: 生成的Markdown內容"""response = result["response"]elapsed_time = result["time"]# 提取文本內容content = response.choices[0].message.content# 生成Markdownmarkdown = f"""# 圖像分析報告## 原始圖像
![原始圖像]({os.path.abspath(image_path)})## 分析結果
{content}## 統計信息
- 處理時間: {elapsed_time:.2f}秒
- 總請求數: {self.total_requests}
- 總Token數: {self.total_tokens}
- 平均響應時間: {self.total_time/self.total_requests:.2f}秒
"""# 保存到文件with open(output_path, 'w', encoding='utf-8') as f:f.write(markdown)return markdownclass ProcessingConfig:"""PDF處理配置類"""def __init__(self, config_dict: Dict[str, Any]):self.api_key = config_dict.get("api_key", "")self.input_dir = config_dict.get("input_dir", "input")self.output_dir = config_dict.get("output_dir", "output")self.model = config_dict.get("model", "glm-4v-flash")self.dpi = config_dict.get("dpi", 600)self.api_interval = config_dict.get("api_interval", 3.0)self.max_retries = config_dict.get("max_retries", 3)self.retry_backoff = config_dict.get("retry_backoff", 0.5)self.prompt = config_dict.get("prompt", "你是一個OCR助手，請把圖中內容按原有格式輸出出來，不要翻譯，如果有公式則輸出為LaTeX，圖片忽略不管")class PDFProcessor:"""PDF處理核心類"""def __init__(self, config: ProcessingConfig, ocr_engine: GLM4VTester):"""初始化PDF處理器:param config: 處理配置:param ocr_engine: OCR引擎實例"""self.config = configself.ocr_engine = ocr_engineself.temp_dir = "temp_images"os.makedirs(self.temp_dir, exist_ok=True)def _convert_page_to_image(self, page, page_num: int) -> str:"""將PDF頁面轉換為圖片:param page: PyMuPDF頁面對象:param page_num: 頁碼:return: 圖片文件路徑"""pix = page.get_pixmap(dpi=self.config.dpi)img_path = os.path.join(self.temp_dir, f"page_{page_num}.png")pix.save(img_path)return img_pathdef _safe_api_call(self, image_path: str) -> str:"""安全的API調用方法，包含重試機制:param image_path: 圖片路徑:return: OCR結果文本"""retries = 0while retries <= self.config.max_retries:try:time.sleep(self.config.api_interval + (retries * self.config.retry_backoff))result = self.ocr_engine.analyze_image(image_path, self.config.prompt)return result["response"].choices[0].message.contentexcept ZhipuAIError as e:logger.warning(f"API調用失敗(重試 {retries}/{self.config.max_retries}): {e}")retries += 1raise Exception(f"API調用失敗，超過最大重試次數 {self.config.max_retries}")def _format_page(self, content: str, page_num: int) -> str:"""格式化單頁內容為Markdown:param content: OCR原始內容:param page_num: 頁碼:return: 格式化后的Markdown"""return f"## 第 {page_num} 頁\n\n{content}\n\n---\n"def process_pdf(self, pdf_path: str) -> Generator[str, None, None]:"""處理單個PDF文件:param pdf_path: PDF文件路徑:return: 生成Markdown內容"""logger.info(f"開始處理PDF文件: {pdf_path}")with fitz.open(pdf_path) as doc:for page_num, page in enumerate(doc, start=1):try:# 轉換為圖片img_path = self._convert_page_to_image(page, page_num)# OCR識別content = self._safe_api_call(img_path)# 格式化輸出yield self._format_page(content, page_num)# 清理臨時圖片os.remove(img_path)except Exception as e:logger.error(f"處理第{page_num}頁時出錯: {e}")yield f"## 第 {page_num} 頁\n\n[處理錯誤: {str(e)}]\n\n"logger.info(f"完成PDF處理: {pdf_path}")def process_single_image(config: ProcessingConfig, image_path: str, output_path: str):"""處理單張圖片模式"""try:tester = GLM4VTester(api_key=config.api_key, model_name=config.model)logger.info(f"開始分析文件: {image_path}")result = tester.analyze_image(image_path, config.prompt)markdown = tester.generate_markdown_report(image_path, result, output_path)print(f"\n分析完成! 結果已保存到: {output_path}\n")return Trueexcept Exception as e:logger.error(f"文件處理失敗: {e}")return Falsedef process_pdf_file(config: ProcessingConfig, pdf_path: str, output_path: str):"""處理PDF文件模式"""try:tester = GLM4VTester(api_key=config.api_key, model_name=config.model)processor = PDFProcessor(config, tester)with open(output_path, 'w', encoding='utf-8') as f:for page_content in processor.process_pdf(pdf_path):f.write(page_content)logger.info(f"PDF轉換完成! 結果已保存到: {output_path}")return Trueexcept Exception as e:logger.error(f"PDF處理失敗: {e}")return Falsedef batch_process_pdfs(config: ProcessingConfig):"""批量處理input/目錄下的PDF文件"""tester = GLM4VTester(api_key=config.api_key, model_name=config.model)processor = PDFProcessor(config, tester)input_dir = config.input_diroutput_dir = config.output_diros.makedirs(input_dir, exist_ok=True)os.makedirs(output_dir, exist_ok=True)processed_files = set()if os.path.exists("processed.log"):with open("processed.log", "r") as f:processed_files = set(f.read().splitlines())while True:try:for filename in os.listdir(input_dir):if filename.lower().endswith('.pdf') and filename not in processed_files:pdf_path = os.path.join(input_dir, filename)output_path = os.path.join(output_dir, f"{os.path.splitext(filename)[0]}.md")logger.info(f"開始處理: {filename}")with open(output_path, 'w', encoding='utf-8') as f:for page_content in processor.process_pdf(pdf_path):f.write(page_content)# 記錄已處理文件with open("processed.log", "a") as f:f.write(f"{filename}\n")processed_files.add(filename)logger.info(f"處理完成: {filename} -> {output_path}")time.sleep(10)  # 每10秒檢查一次新文件except KeyboardInterrupt:logger.info("收到中斷信號，停止處理")breakexcept Exception as e:logger.error(f"批量處理出錯: {e}")time.sleep(30)  # 出錯后等待30秒再重試def load_config():"""加載配置文件"""config_path = "config.json"default_config = {"api_key": "","input_dir": "input","output_dir": "output","model": "glm-4v-flash","dpi": 600,"api_interval": 3.0,"max_retries": 3,"retry_backoff": 0.5,"prompt": "你是一個OCR助手，請把圖中內容按原有格式輸出出來，如果有公式則輸出為LaTeX，圖片請用《》描述"}try:with open(config_path, 'r') as f:config = json.load(f)# 合并配置，優先使用配置文件中的值return {**default_config, **config}except FileNotFoundError:logger.warning(f"配置文件 {config_path} 未找到，使用默認配置")# 創建默認配置文件with open(config_path, 'w') as f:json.dump(default_config, f, indent=2)return default_configexcept json.JSONDecodeError as e:logger.error(f"配置文件格式錯誤: {e}")return default_configdef main():"""主函數"""config_dict = load_config()config = ProcessingConfig(config_dict)# 檢查API密鑰是否設置if not config.api_key:logger.error("API密鑰未設置，請在config.json中設置api_key")exit(1)# 確保目錄存在os.makedirs(config.input_dir, exist_ok=True)os.makedirs(config.output_dir, exist_ok=True)# 直接啟動批處理模式logger.info(f"啟動批處理模式，監控目錄: {config.input_dir}")batch_process_pdfs(config)if __name__ == '__main__':main()

自己修改一下config里面的智譜api key：

{"api_key": "智譜的api_key","input_dir": "input","output_dir": "output", "model": "glm-4v-flash","dpi": 600,"api_interval": 3.0,"max_retries": 3,"retry_backoff": 0.5
}