"""Hybrid PDF processing: extract usable text layers and OCR only when needed.""" from __future__ import annotations import hashlib import json import logging import os import re import shutil import time from datetime import datetime from pathlib import Path from typing import Any, Iterable import pypdfium2 as pdfium from PIL import Image from .pdf_text import TextLayerPolicy, extract_page_text MANIFEST_VERSION = 2 PAGE_SPEC_PATTERN = re.compile(r"^(\d+)(?:-(\d*)?)?$") PDF_MODES = {"hybrid", "text", "ocr"} def now_iso() -> str: return datetime.now().astimezone().isoformat() def atomic_write_text(path: Path, content: str) -> None: path.parent.mkdir(parents=True, exist_ok=True) temporary = path.with_name(f".{path.name}.tmp") temporary.write_text(content, encoding="utf-8") temporary.replace(path) def atomic_write_json(path: Path, data: Any) -> None: atomic_write_text(path, json.dumps(data, ensure_ascii=False, indent=2)) def sha256_file(path: Path, chunk_size: int = 1024 * 1024) -> str: digest = hashlib.sha256() with path.open("rb") as file: while chunk := file.read(chunk_size): digest.update(chunk) return digest.hexdigest() def safe_stem(value: str) -> str: cleaned = re.sub(r"[^\w.-]+", "_", value, flags=re.UNICODE).strip("._") return cleaned or "document" def parse_page_spec(spec: str | None, page_count: int) -> list[int]: if page_count < 1: return [] if spec is None or not spec.strip(): return list(range(page_count)) selected: set[int] = set() for raw_part in spec.split(","): part = raw_part.strip() match = PAGE_SPEC_PATTERN.fullmatch(part) if not match: raise ValueError(f"无效页码范围: {part!r}，示例: 1-5,8,10-") start = int(match.group(1)) end_text = match.group(2) end = start if "-" not in part else int(end_text) if end_text else page_count if start < 1 or end < 1: raise ValueError("PDF 页码从 1 开始") if start > end: raise ValueError(f"页码起始值不能大于结束值: {part}") if start > page_count or end > page_count: raise ValueError(f"页码范围 {part} 超出 PDF 总页数 {page_count}") selected.update(range(start - 1, end)) return sorted(selected) def render_page(document: Any, page_index: int, dpi: int) -> Image.Image: page = document.get_page(page_index) bitmap = None try: bitmap = page.render(scale=dpi / 72.0) return bitmap.to_pil().convert("RGB").copy() finally: if bitmap is not None: bitmap.close() page.close() def save_png_atomic(image: Image.Image, path: Path) -> None: path.parent.mkdir(parents=True, exist_ok=True) temporary = path.with_name(f".{path.name}.tmp") image.save(temporary, format="PNG") temporary.replace(path) def _save_markdown_image(data: Any, path: Path) -> Path: path.parent.mkdir(parents=True, exist_ok=True) if isinstance(data, Image.Image): image = data else: import numpy as np array = np.asarray(data) if array.ndim not in (2, 3): raise TypeError(f"无法保存 Markdown 图片: {array.shape}") image = Image.fromarray(array.astype("uint8")) image_format = (path.suffix.lstrip(".") or "png").upper() if image_format == "JPG": image_format = "JPEG" if image_format not in {"PNG", "JPEG", "WEBP", "BMP", "TIFF"}: path = path.with_suffix(".png") image_format = "PNG" temporary = path.with_name(f".{path.name}.tmp") image.save(temporary, format=image_format) temporary.replace(path) return path def _ocr_markdown(result: Any, document_dir: Path, page_number: int) -> str: data = result.markdown if "res" in data and isinstance(data["res"], dict): data = data["res"] text = str(data.get("markdown_texts", "")) asset_dir = document_dir / "assets" / f"page-{page_number:04d}" if asset_dir.exists(): shutil.rmtree(asset_dir) for index, (original_path, image_data) in enumerate((data.get("markdown_images") or {}).items(), 1): original = str(original_path).replace("\\", "/") source_name = Path(original).name or f"image-{index:03d}.png" target = asset_dir / f"{index:03d}-{safe_stem(Path(source_name).stem)}{Path(source_name).suffix or '.png'}" target = _save_markdown_image(image_data, target) relative = Path(os.path.relpath(target, document_dir / "pages")).as_posix() text = text.replace(original, relative).replace(str(original_path), relative) return text.strip() def _page_paths(document_dir: Path, page_number: int) -> tuple[Path, Path]: stem = f"page-{page_number:04d}" return document_dir / "pages" / f"{stem}.md", document_dir / "pages" / f"{stem}.json" def _page_is_complete(document_dir: Path, manifest: dict[str, Any], page_number: int) -> bool: record = manifest.get("pages", {}).get(str(page_number), {}) markdown_path, json_path = _page_paths(document_dir, page_number) return record.get("status") == "completed" and markdown_path.is_file() and json_path.is_file() def rebuild_combined_outputs(document_dir: Path, manifest: dict[str, Any]) -> None: markdown_parts = [f"# {manifest['document_name']}"] page_results = [] for page_number in manifest.get("selected_pages", []): record = manifest.get("pages", {}).get(str(page_number), {}) markdown_path, json_path = _page_paths(document_dir, page_number) if record.get("status") == "completed" and markdown_path.is_file() and json_path.is_file(): text = markdown_path.read_text(encoding="utf-8").replace("../assets/", "assets/") source = record.get("source_type", "unknown") markdown_parts.append(f"\n\n---\n\n## Page {page_number} ({source})\n\n{text.strip()}") page_results.append( { "page_number": page_number, "source_type": source, "metrics": record, "result": json.loads(json_path.read_text(encoding="utf-8")), } ) elif record.get("status") == "failed": markdown_parts.append(f"\n\n---\n\n## Page {page_number}\n\n> Failed: {record.get('error')}") atomic_write_text(document_dir / "document.md", "".join(markdown_parts).rstrip() + "\n") atomic_write_json(document_dir / "document.json", {"manifest": manifest, "page_results": page_results}) def validate_pdf_request(pdf_path: Path, output_root: Path, *, resume: bool, overwrite: bool) -> tuple[Path, Path]: pdf_path = pdf_path.expanduser().resolve() output_root = output_root.expanduser().resolve() if not pdf_path.is_file(): raise FileNotFoundError(f"PDF 不存在: {pdf_path}") if pdf_path.suffix.lower() != ".pdf": raise ValueError(f"输入文件不是 PDF: {pdf_path}") if resume and overwrite: raise ValueError("--resume 和 --overwrite 不能同时使用") document_dir = output_root / safe_stem(pdf_path.stem) if resume and not (document_dir / "manifest.json").is_file(): raise FileNotFoundError(f"无法续传，缺少 {document_dir / 'manifest.json'}") if document_dir.exists() and any(document_dir.iterdir()) and not (resume or overwrite): raise FileExistsError(f"输出目录已存在: {document_dir}；请使用 --resume 或 --overwrite") return pdf_path, output_root def preflight_pdf( *, pdf_path: Path, output_root: Path, pages: str | None, dpi: int, password: str | None, resume: bool, overwrite: bool, ) -> dict[str, Any]: pdf_path, output_root = validate_pdf_request(pdf_path, output_root, resume=resume, overwrite=overwrite) if dpi < 72 or dpi > 600: raise ValueError("--dpi 必须在 72 到 600 之间") document = pdfium.PdfDocument(str(pdf_path), password=password) try: page_count = len(document) selected = parse_page_spec(pages, page_count) finally: document.close() return { "pdf_path": pdf_path, "output_root": output_root, "document_dir": output_root / safe_stem(pdf_path.stem), "page_count": page_count, "selected_pages": [index + 1 for index in selected], } def _prepare_manifest( *, pdf_path: Path, document_dir: Path, page_count: int, selected_pages: Iterable[int], dpi: int, mode: str, policy: TextLayerPolicy, resume: bool, overwrite: bool, run_metadata: dict[str, Any], ) -> dict[str, Any]: manifest_path = document_dir / "manifest.json" digest = sha256_file(pdf_path) selected = [index + 1 for index in selected_pages] if overwrite and document_dir.exists(): shutil.rmtree(document_dir) if resume: manifest = json.loads(manifest_path.read_text(encoding="utf-8")) if manifest.get("manifest_version") != MANIFEST_VERSION: raise ValueError("旧版 manifest 不兼容混合模式，请使用 --overwrite") if manifest.get("input", {}).get("sha256") != digest: raise ValueError("PDF 内容已变化，请使用 --overwrite") if manifest.get("render", {}).get("dpi") != dpi or manifest.get("mode") != mode: raise ValueError("DPI 或模式与原任务不一致，请使用原参数或 --overwrite") if manifest.get("text_layer_policy") != policy.__dict__: raise ValueError("文本层阈值与原任务不一致，请使用原参数或 --overwrite") manifest["selected_pages"] = sorted(set(manifest.get("selected_pages", [])) | set(selected)) manifest["run_metadata"] = run_metadata manifest["status"] = "running" manifest["updated_at"] = now_iso() else: document_dir.mkdir(parents=True, exist_ok=True) manifest = { "manifest_version": MANIFEST_VERSION, "document_name": pdf_path.stem, "input": {"path": str(pdf_path), "sha256": digest, "size_bytes": pdf_path.stat().st_size}, "page_count": page_count, "selected_pages": selected, "mode": mode, "text_layer_policy": policy.__dict__, "render": {"dpi": dpi, "format": "png"}, "run_metadata": run_metadata, "status": "running", "created_at": now_iso(), "updated_at": now_iso(), "pages": {}, } atomic_write_json(manifest_path, manifest) return manifest def process_pdf( *, provider: Any, pdf_path: Path, output_root: Path, mode: str = "hybrid", text_policy: TextLayerPolicy | None = None, pages: str | None = None, dpi: int = 144, password: str | None = None, resume: bool = False, overwrite: bool = False, keep_rendered: bool = False, fail_fast: bool = False, predict_kwargs: dict[str, Any] | None = None, logger: logging.Logger | None = None, ) -> dict[str, Any]: if mode not in PDF_MODES: raise ValueError(f"不支持的 PDF 模式: {mode}") task_started = time.perf_counter() model_init_before = provider.model_init_seconds logger = logger or logging.getLogger(__name__) text_policy = text_policy or TextLayerPolicy() predict_kwargs = predict_kwargs or {} pdf_path, output_root = validate_pdf_request(pdf_path, output_root, resume=resume, overwrite=overwrite) if dpi < 72 or dpi > 600: raise ValueError("--dpi 必须在 72 到 600 之间") document_dir = output_root / safe_stem(pdf_path.stem) manifest_path = document_dir / "manifest.json" cache_dir = document_dir / ".render-cache" opened = time.perf_counter() document = pdfium.PdfDocument(str(pdf_path), password=password) pdf_open_seconds = time.perf_counter() - opened logger.info("PDF_OPENED path=%s mode=%s seconds=%.3f dpi=%d", pdf_path, mode, pdf_open_seconds, dpi) try: page_count = len(document) selected_indexes = parse_page_spec(pages, page_count) prepared = time.perf_counter() manifest = _prepare_manifest( pdf_path=pdf_path, document_dir=document_dir, page_count=page_count, selected_pages=selected_indexes, dpi=dpi, mode=mode, policy=text_policy, resume=resume, overwrite=overwrite, run_metadata={"device": provider.resolved_device}, ) manifest_prepare_seconds = time.perf_counter() - prepared selected_indexes = [number - 1 for number in manifest["selected_pages"]] completed_before = sum(_page_is_complete(document_dir, manifest, index + 1) for index in selected_indexes) pending = [index for index in selected_indexes if not _page_is_complete(document_dir, manifest, index + 1)] logger.info( "TASK_PLAN mode=%s total_pages=%d selected_pages=%d completed_before=%d pending_pages=%d", mode, page_count, len(selected_indexes), completed_before, len(pending), ) current_run_times: list[float] = [] for position, page_index in enumerate(pending, 1): page_number = page_index + 1 started = time.perf_counter() text_extract_seconds = render_seconds = ocr_seconds = export_seconds = state_save_seconds = 0.0 source_type = "unknown" assessment_dict: dict[str, Any] = {} render_path = (document_dir / "rendered" if keep_rendered else cache_dir) / f"page-{page_number:04d}.png" logger.info("PAGE_START page=%d position=%d/%d mode=%s", page_number, position, len(pending), mode) try: text_started = time.perf_counter() page = document.get_page(page_index) try: extracted_text, assessment = extract_page_text(page, text_policy) width_points, height_points = page.get_size() finally: page.close() text_extract_seconds = time.perf_counter() - text_started assessment_dict = assessment.to_dict() use_text = mode == "text" or (mode == "hybrid" and assessment.usable) source_type = "text" if use_text else "ocr" logger.info( "PAGE_ROUTED page=%d source=%s reason=%s text_chars=%d printable_ratio=%.4f content_ratio=%.4f density=%.3f text_extract_seconds=%.3f", page_number, source_type, assessment.reason, assessment.non_whitespace_chars, assessment.printable_ratio, assessment.content_ratio, assessment.chars_per_megapixel, text_extract_seconds, ) markdown_path, json_path = _page_paths(document_dir, page_number) if source_type == "text": markdown_text = extracted_text payload = { "res": { "input_path": str(pdf_path), "page_index": page_index, "page_number": page_number, "page_count": page_count, "source_type": "text", "text": extracted_text, "text_layer": assessment_dict, "width_points": width_points, "height_points": height_points, } } layout_boxes = 0 parsed_blocks = 1 if extracted_text else 0 else: rendered = time.perf_counter() image = render_page(document, page_index, dpi) try: save_png_atomic(image, render_path) finally: image.close() render_seconds = time.perf_counter() - rendered pipeline = provider.get() provider.synchronize() ocr_started = time.perf_counter() results = pipeline.predict(str(render_path), **predict_kwargs) provider.synchronize() ocr_seconds = time.perf_counter() - ocr_started if not results: raise RuntimeError("OCR pipeline 未返回结果") result = results[0] markdown_text = _ocr_markdown(result, document_dir, page_number) payload = result.json result_payload = payload.get("res", payload) result_payload.update( { "input_path": str(pdf_path), "page_index": page_index, "page_number": page_number, "page_count": page_count, "source_type": "ocr", "ocr_reason": assessment.reason if mode == "hybrid" else "forced_ocr_mode", "text_layer": assessment_dict, "render_dpi": dpi, } ) layout_boxes = len(result.get("layout_det_res", {}).get("boxes", [])) parsed_blocks = len(result.get("parsing_res_list", [])) export_started = time.perf_counter() atomic_write_text(markdown_path, markdown_text.rstrip() + "\n") atomic_write_json(json_path, payload) export_seconds = time.perf_counter() - export_started total_seconds = time.perf_counter() - started manifest["pages"][str(page_number)] = { "status": "completed", "page_number": page_number, "source_type": source_type, "routing_reason": assessment.reason if mode == "hybrid" else f"forced_{source_type}_mode", "text_layer": assessment_dict, "text_extract_seconds": round(text_extract_seconds, 3), "render_seconds": round(render_seconds, 3), "ocr_seconds": round(ocr_seconds, 3), "export_seconds": round(export_seconds, 3), "total_seconds": round(total_seconds, 3), "layout_boxes": layout_boxes, "parsed_blocks": parsed_blocks, "completed_at": now_iso(), } current_run_times.append(total_seconds) except KeyboardInterrupt: manifest["status"] = "interrupted" manifest["updated_at"] = now_iso() atomic_write_json(manifest_path, manifest) rebuild_combined_outputs(document_dir, manifest) logger.warning("TASK_INTERRUPTED page=%d", page_number) raise except Exception as exc: total_seconds = time.perf_counter() - started manifest["pages"][str(page_number)] = { "status": "failed", "page_number": page_number, "source_type": source_type, "text_layer": assessment_dict, "text_extract_seconds": round(text_extract_seconds, 3), "render_seconds": round(render_seconds, 3), "ocr_seconds": round(ocr_seconds, 3), "export_seconds": round(export_seconds, 3), "total_seconds": round(total_seconds, 3), "error": f"{type(exc).__name__}: {exc}", "failed_at": now_iso(), } logger.exception("PAGE_FAILED page=%d source=%s", page_number, source_type) if fail_fast: raise finally: if not keep_rendered and render_path.is_file(): render_path.unlink() saved = time.perf_counter() manifest["run_metadata"] = provider.metadata() manifest["updated_at"] = now_iso() atomic_write_json(manifest_path, manifest) rebuild_combined_outputs(document_dir, manifest) state_save_seconds = time.perf_counter() - saved manifest["pages"][str(page_number)]["state_save_seconds"] = round(state_save_seconds, 3) atomic_write_json(manifest_path, manifest) average = sum(current_run_times) / len(current_run_times) if current_run_times else None eta = average * (len(pending) - position) if average is not None else None record = manifest["pages"][str(page_number)] logger.info( "PAGE_FINISHED page=%d status=%s source=%s text_extract_seconds=%.3f render_seconds=%.3f ocr_seconds=%.3f export_seconds=%.3f state_save_seconds=%.3f total_seconds=%.3f eta_seconds=%s progress=%d/%d", page_number, record["status"], record.get("source_type"), text_extract_seconds, render_seconds, ocr_seconds, export_seconds, state_save_seconds, record.get("total_seconds", 0.0), f"{eta:.3f}" if eta is not None else "unknown", position, len(pending), ) if cache_dir.exists(): shutil.rmtree(cache_dir, ignore_errors=True) records = [manifest.get("pages", {}).get(str(index + 1), {}) for index in selected_indexes] failed_pages = [record.get("page_number") for record in records if record.get("status") == "failed"] completed = [record for record in records if record.get("status") == "completed"] text_pages = sum(record.get("source_type") == "text" for record in completed) ocr_pages = sum(record.get("source_type") == "ocr" for record in completed) timing_keys = ("text_extract_seconds", "render_seconds", "ocr_seconds", "export_seconds", "state_save_seconds", "total_seconds") totals = {key: sum(record.get(key, 0.0) for record in records) for key in timing_keys} finalize_started = time.perf_counter() manifest["status"] = "completed_with_errors" if failed_pages else "completed" manifest["run_metadata"] = provider.metadata() manifest["summary"] = { "selected_pages": len(selected_indexes), "completed_pages": len(completed), "completed_before_resume": completed_before, "text_pages": text_pages, "ocr_pages": ocr_pages, "failed_pages": failed_pages, "model_used": ocr_pages > 0, "model_initialized_during_task": ( model_init_before == 0 and provider.model_init_seconds > 0 ), "model_available": provider.model_init_seconds > 0, "timing": { "pdf_open_seconds": round(pdf_open_seconds, 3), "manifest_prepare_seconds": round(manifest_prepare_seconds, 3), "text_extract_total_seconds": round(totals["text_extract_seconds"], 3), "render_total_seconds": round(totals["render_seconds"], 3), "ocr_total_seconds": round(totals["ocr_seconds"], 3), "export_total_seconds": round(totals["export_seconds"], 3), "state_save_total_seconds": round(totals["state_save_seconds"], 3), "page_total_seconds": round(totals["total_seconds"], 3), "model_init_seconds": round(provider.model_init_seconds, 3), "finalize_seconds": 0.0, "task_total_seconds": 0.0, }, } manifest["updated_at"] = now_iso() atomic_write_json(manifest_path, manifest) rebuild_combined_outputs(document_dir, manifest) finalize_seconds = time.perf_counter() - finalize_started task_total = time.perf_counter() - task_started manifest["summary"]["timing"]["finalize_seconds"] = round(finalize_seconds, 3) manifest["summary"]["timing"]["task_total_seconds"] = round(task_total, 3) atomic_write_json(manifest_path, manifest) rebuild_combined_outputs(document_dir, manifest) logger.info( "TASK_COMPLETED status=%s mode=%s selected_pages=%d text_pages=%d ocr_pages=%d failed_pages=%s model_used=%s model_initialized_during_task=%s model_available=%s model_init_seconds=%.3f text_extract_total_seconds=%.3f render_total_seconds=%.3f ocr_total_seconds=%.3f task_total_seconds=%.3f", manifest["status"], mode, len(selected_indexes), text_pages, ocr_pages, failed_pages, manifest["summary"]["model_used"], manifest["summary"]["model_initialized_during_task"], manifest["summary"]["model_available"], provider.model_init_seconds, totals["text_extract_seconds"], totals["render_seconds"], totals["ocr_seconds"], task_total, ) return {"document_dir": str(document_dir), "manifest_path": str(manifest_path), "status": manifest["status"], **manifest["summary"]} finally: document.close()