Skip to content

pix2text_ocr.py

概述

dataGen/pix2text_ocr.py 使用 Pix2Text 将 PDF 逐页转换为带 LaTeX 公式的 Markdown 文件。它支持顺序、批量和多进程三种处理模式,并且能够按页跳过已经具有输出文件的页面。

它的配置主要依赖于 config.toml[ocr] 节的内容。

模块级常量

常量来源说明
PAGE_STARTocr.page_start起始页(0-based)
PAGE_ENDocr.page_end结束页(此页不包含或包含需看具体代码,一般为最大页数上界)
SKIP_EXISTINGocr.skip_existing是否跳过已经存在的 MD 文件
DEVICEocr.device计算设备,例如 "cuda"
MFR_BATCH_SIZEocr.mfr_batch_size公式识别模型的 batch size
RENDER_DPIocr.render_dpiPDF 渲染为图像时的 DPI
RESIZED_SHAPEocr.resized_shapeOCR 识别时图像缩放的宽度
TEXT_CONTAIN_FORMULAocr.text_contain_formula是否在文本中检测行内公式
BATCH_PAGESocr.batch_pages批量处理时一次处理的页数
OCR_WORKERSocr.ocr_workers并发执行的工作进程数量

函数说明

_collect_pdfs() -> list[str]

扫描 config.RAW_DIR 目录下的所有 PDF 文件,并返回按字母排序的文件相对或绝对路径列表。

_get_pdf_page_count(pdf_path: str) -> int | None

借助 pypdf 获取指定 PDF 文件的总页数。如果获取失败(例如模块未安装或文件损坏),则返回 None

_iter_pages(total_pages, page_start=None, page_end=None) -> list[int]

基于指定的起始和结束参数,生成一个包含基于 0 的页码索引的列表,过滤掉不合理的范围。

_worker_process_pages(worker_id, pdf_path, page_indices, book_dir, pages_dir, device, mfr_batch_size, resized_shape, render_dpi, text_contain_formula=True) -> int

作为多进程模式(Sub-process)的执行单元,它会加载完整的 Pix2Text 引擎,遍历并渲染 page_indices 对应的页,调用 OCR 方法并将其按指定的页数写入文件。返回成功处理的页数。

_process_single_pdf(pdf_path, p2t, page_start_override=None)

管理单一 PDF 文件的全套处理流程。包含日志打印,建立输出目录、并根据 OCR_WORKERSBATCH_PAGES 的选择,委派串行、批处理或多进程等执行策略。

_save_page_md(page_obj, page_idx, book_dir, pages_dir)

辅助函数:将 Pix2Text 识别完毕的 page_obj 对象持久化写入磁盘特定的页面级别的 .md 文件中。

_render_page(pdf_path, page_idx)

使用 fitz (PyMuPDF) 将具体的某一页渲染并导出为 PIL.Image RGB 对象。

_process_pages_sequential(...)

顺序处理策略。单线程下一页页地进行渲染、OCR 识别与文件写入。

_process_pages_batched(...)

批处理策略。将多个页面放在一个缓冲器中传递给底层库一起处理,以缩少切换和推理调用的开销。

_process_pages_parallel(...)

多进程处理策略。使用 multiprocessing 或同类组件拉起多个 worker 以并行分担多页的 OCR 计算。

main()

CLI 入口点。允许用户解析可选的运行参数(如:“书名.pdf” 或起始页码)并在项目中循环调度 _process_single_pdf() 执行任务。

使用示例

bash
# 按顺序处理所有 PDF
python dataGen/pix2text_ocr.py

# 只处理指定的 PDF
python dataGen/pix2text_ocr.py "书名.pdf"

# 从指定书的第 136 页开始(命令行中数字一般作为起始页面)
python dataGen/pix2text_ocr.py "书名.pdf" 136

Math-RAG 毕业设计项目