659 lines
26 KiB
Python
659 lines
26 KiB
Python
"""Shared PDF rendering, OCR orchestration, resume, and export logic."""
|
|
|
|
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, Callable, Iterable
|
|
|
|
import pypdfium2 as pdfium
|
|
from PIL import Image
|
|
|
|
MANIFEST_VERSION = 1
|
|
PAGE_SPEC_PATTERN = re.compile(r"^(\d+)(?:-(\d*)?)?$")
|
|
|
|
|
|
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]:
|
|
"""Parse one-based ranges such as ``1-5,8,10-`` into zero-based indexes."""
|
|
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)
|
|
if "-" not in part:
|
|
end = start
|
|
elif end_text:
|
|
end = int(end_text)
|
|
else:
|
|
end = 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)
|
|
temporary = path.with_name(f".{path.name}.tmp")
|
|
|
|
if isinstance(data, Image.Image):
|
|
image = data
|
|
else:
|
|
try:
|
|
import numpy as np
|
|
|
|
array = np.asarray(data)
|
|
if array.ndim == 3 and array.shape[2] == 4:
|
|
image = Image.fromarray(array.astype("uint8"), mode="RGBA")
|
|
elif array.ndim in (2, 3):
|
|
image = Image.fromarray(array.astype("uint8"))
|
|
else:
|
|
raise TypeError(f"unsupported image array shape: {array.shape}")
|
|
except Exception as exc:
|
|
raise TypeError(f"无法保存 Markdown 图片 {path.name}: {type(data).__name__}") from exc
|
|
|
|
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"}:
|
|
image_format = "PNG"
|
|
path = path.with_suffix(".png")
|
|
temporary = path.with_name(f".{path.name}.tmp")
|
|
image.save(temporary, format=image_format)
|
|
temporary.replace(path)
|
|
return path
|
|
|
|
|
|
def _result_markdown(result: Any, document_dir: Path, page_number: int) -> str:
|
|
markdown_data = result.markdown
|
|
if "res" in markdown_data and isinstance(markdown_data["res"], dict):
|
|
markdown_data = markdown_data["res"]
|
|
|
|
text = str(markdown_data.get("markdown_texts", ""))
|
|
markdown_images = markdown_data.get("markdown_images") or {}
|
|
page_asset_dir = document_dir / "assets" / f"page-{page_number:04d}"
|
|
if page_asset_dir.exists():
|
|
shutil.rmtree(page_asset_dir)
|
|
|
|
for index, (original_path, image_data) in enumerate(markdown_images.items(), start=1):
|
|
original = str(original_path).replace("\\", "/")
|
|
original_name = Path(original).name or f"image-{index:03d}.png"
|
|
asset_name = f"{index:03d}-{safe_stem(Path(original_name).stem)}{Path(original_name).suffix or '.png'}"
|
|
target = page_asset_dir / asset_name
|
|
target = _save_markdown_image(image_data, target)
|
|
page_relative = Path(os.path.relpath(target, document_dir / "pages")).as_posix()
|
|
text = text.replace(original, page_relative)
|
|
text = text.replace(str(original_path), page_relative)
|
|
|
|
return text.strip()
|
|
|
|
|
|
def _result_json(result: Any) -> dict[str, Any]:
|
|
data = result.json
|
|
if not isinstance(data, dict):
|
|
raise TypeError(f"OCR JSON 结果类型异常: {type(data).__name__}")
|
|
return data
|
|
|
|
|
|
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_json_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():
|
|
page_text = markdown_path.read_text(encoding="utf-8")
|
|
page_text = page_text.replace("../assets/", "assets/")
|
|
markdown_parts.append(f"\n\n---\n\n## Page {page_number}\n\n{page_text.strip()}")
|
|
page_json_results.append(
|
|
{
|
|
"page_number": page_number,
|
|
"metrics": record,
|
|
"ocr_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> OCR failed: {record.get('error', 'unknown 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_json_results,
|
|
},
|
|
)
|
|
|
|
|
|
def prepare_manifest(
|
|
*,
|
|
pdf_path: Path,
|
|
document_dir: Path,
|
|
page_count: int,
|
|
selected_pages: Iterable[int],
|
|
dpi: int,
|
|
resume: bool,
|
|
overwrite: bool,
|
|
run_metadata: dict[str, Any],
|
|
) -> dict[str, Any]:
|
|
manifest_path = document_dir / "manifest.json"
|
|
pdf_sha256 = sha256_file(pdf_path)
|
|
selected_one_based = [index + 1 for index in selected_pages]
|
|
|
|
if overwrite and document_dir.exists():
|
|
shutil.rmtree(document_dir)
|
|
|
|
if document_dir.exists() and any(document_dir.iterdir()) and not resume:
|
|
raise FileExistsError(
|
|
f"输出目录已存在: {document_dir}。请使用 --resume 继续或 --overwrite 重建。"
|
|
)
|
|
|
|
if resume:
|
|
if not manifest_path.is_file():
|
|
raise FileNotFoundError(f"无法断点续传,缺少 manifest: {manifest_path}")
|
|
manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
|
|
if manifest.get("input", {}).get("sha256") != pdf_sha256:
|
|
raise ValueError("PDF 内容已变化,不能使用现有断点;请使用 --overwrite")
|
|
if manifest.get("render", {}).get("dpi") != dpi:
|
|
raise ValueError("DPI 与现有任务不一致;请使用原 DPI 或 --overwrite")
|
|
manifest["selected_pages"] = sorted(
|
|
set(manifest.get("selected_pages", [])) | set(selected_one_based)
|
|
)
|
|
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": pdf_sha256,
|
|
"size_bytes": pdf_path.stat().st_size,
|
|
},
|
|
"page_count": page_count,
|
|
"selected_pages": selected_one_based,
|
|
"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 validate_pdf_request(
|
|
pdf_path: Path,
|
|
output_root: Path,
|
|
*,
|
|
resume: bool,
|
|
overwrite: bool,
|
|
) -> tuple[Path, Path]:
|
|
"""Validate cheap input/output conditions before loading the large model."""
|
|
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"无法断点续传,缺少 manifest: {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]:
|
|
"""Validate PDF access, page ranges, and output state before model loading."""
|
|
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 process_pdf(
|
|
*,
|
|
pipeline: Any,
|
|
pdf_path: Path,
|
|
output_root: Path,
|
|
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,
|
|
run_metadata: dict[str, Any] | None = None,
|
|
predict_kwargs: dict[str, Any] | None = None,
|
|
synchronize: Callable[[], None] | None = None,
|
|
logger: logging.Logger | None = None,
|
|
) -> dict[str, Any]:
|
|
"""Render and OCR a PDF one page at a time."""
|
|
task_started = time.perf_counter()
|
|
logger = logger or logging.getLogger(__name__)
|
|
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 之间")
|
|
|
|
predict_kwargs = predict_kwargs or {}
|
|
run_metadata = run_metadata or {}
|
|
document_dir = output_root / safe_stem(pdf_path.stem)
|
|
manifest_path = document_dir / "manifest.json"
|
|
temporary_render_dir = document_dir / ".render-cache"
|
|
|
|
pdf_open_started = time.perf_counter()
|
|
document = pdfium.PdfDocument(str(pdf_path), password=password)
|
|
pdf_open_seconds = time.perf_counter() - pdf_open_started
|
|
logger.info(
|
|
"PDF_OPENED path=%s seconds=%.3f dpi=%d resume=%s overwrite=%s keep_rendered=%s",
|
|
pdf_path,
|
|
pdf_open_seconds,
|
|
dpi,
|
|
resume,
|
|
overwrite,
|
|
keep_rendered,
|
|
)
|
|
try:
|
|
page_count = len(document)
|
|
selected_indexes = parse_page_spec(pages, page_count)
|
|
manifest_started = time.perf_counter()
|
|
manifest = prepare_manifest(
|
|
pdf_path=pdf_path,
|
|
document_dir=document_dir,
|
|
page_count=page_count,
|
|
selected_pages=selected_indexes,
|
|
dpi=dpi,
|
|
resume=resume,
|
|
overwrite=overwrite,
|
|
run_metadata=run_metadata,
|
|
)
|
|
manifest_prepare_seconds = time.perf_counter() - manifest_started
|
|
logger.info(
|
|
"MANIFEST_READY path=%s seconds=%.3f page_count=%d requested_pages=%d",
|
|
manifest_path,
|
|
manifest_prepare_seconds,
|
|
page_count,
|
|
len(selected_indexes),
|
|
)
|
|
# Resume uses the union stored in the manifest, so newly added ranges and
|
|
# previously selected pages remain one coherent document task.
|
|
selected_indexes = [page_number - 1 for page_number in manifest["selected_pages"]]
|
|
|
|
completed_before = sum(
|
|
_page_is_complete(document_dir, manifest, index + 1) for index in selected_indexes
|
|
)
|
|
pending_indexes = [
|
|
index
|
|
for index in selected_indexes
|
|
if not _page_is_complete(document_dir, manifest, index + 1)
|
|
]
|
|
logger.info(
|
|
"TASK_PLAN total_pages=%d selected_pages=%d completed_before=%d pending_pages=%d output=%s",
|
|
page_count,
|
|
len(selected_indexes),
|
|
completed_before,
|
|
len(pending_indexes),
|
|
document_dir,
|
|
)
|
|
|
|
run_page_times: list[float] = []
|
|
for position, page_index in enumerate(pending_indexes, start=1):
|
|
page_number = page_index + 1
|
|
page_started = time.perf_counter()
|
|
render_seconds = 0.0
|
|
ocr_seconds = 0.0
|
|
export_seconds = 0.0
|
|
state_save_seconds = 0.0
|
|
render_path = temporary_render_dir / f"page-{page_number:04d}.png"
|
|
if keep_rendered:
|
|
render_path = document_dir / "rendered" / f"page-{page_number:04d}.png"
|
|
|
|
logger.info(
|
|
"PAGE_START page=%d page_index=%d position=%d/%d",
|
|
page_number,
|
|
page_index,
|
|
position,
|
|
len(pending_indexes),
|
|
)
|
|
try:
|
|
render_started = 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() - render_started
|
|
logger.info(
|
|
"PAGE_RENDERED page=%d seconds=%.3f path=%s",
|
|
page_number,
|
|
render_seconds,
|
|
render_path,
|
|
)
|
|
|
|
if synchronize:
|
|
synchronize()
|
|
ocr_started = time.perf_counter()
|
|
result_list = pipeline.predict(str(render_path), **predict_kwargs)
|
|
if synchronize:
|
|
synchronize()
|
|
ocr_seconds = time.perf_counter() - ocr_started
|
|
logger.info("PAGE_OCR_COMPLETED page=%d seconds=%.3f", page_number, ocr_seconds)
|
|
if not result_list:
|
|
raise RuntimeError("OCR pipeline 未返回结果")
|
|
|
|
export_started = time.perf_counter()
|
|
result = result_list[0]
|
|
markdown_text = _result_markdown(result, document_dir, page_number)
|
|
result_json = _result_json(result)
|
|
json_payload = result_json.get("res", result_json)
|
|
if isinstance(json_payload, dict):
|
|
json_payload["input_path"] = str(pdf_path)
|
|
json_payload["page_index"] = page_index
|
|
json_payload["page_number"] = page_number
|
|
json_payload["page_count"] = page_count
|
|
json_payload["render_dpi"] = dpi
|
|
markdown_path, json_path = _page_paths(document_dir, page_number)
|
|
atomic_write_text(markdown_path, markdown_text.rstrip() + "\n")
|
|
atomic_write_json(json_path, result_json)
|
|
export_seconds = time.perf_counter() - export_started
|
|
|
|
total_seconds = time.perf_counter() - page_started
|
|
manifest["pages"][str(page_number)] = {
|
|
"status": "completed",
|
|
"page_number": page_number,
|
|
"render_seconds": round(render_seconds, 3),
|
|
"ocr_seconds": round(ocr_seconds, 3),
|
|
"export_seconds": round(export_seconds, 3),
|
|
"total_seconds": round(total_seconds, 3),
|
|
"width": result.get("width"),
|
|
"height": result.get("height"),
|
|
"layout_boxes": len(result.get("layout_det_res", {}).get("boxes", [])),
|
|
"parsed_blocks": len(result.get("parsing_res_list", [])),
|
|
"device": run_metadata.get("device"),
|
|
"completed_at": now_iso(),
|
|
}
|
|
run_page_times.append(total_seconds)
|
|
logger.info(
|
|
"PAGE_RESULT_SAVED page=%d seconds=%.3f markdown=%s json=%s width=%s height=%s layout_boxes=%d parsed_blocks=%d",
|
|
page_number,
|
|
export_seconds,
|
|
markdown_path,
|
|
json_path,
|
|
result.get("width"),
|
|
result.get("height"),
|
|
len(result.get("layout_det_res", {}).get("boxes", [])),
|
|
len(result.get("parsing_res_list", [])),
|
|
)
|
|
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 elapsed_seconds=%.3f",
|
|
page_number,
|
|
time.perf_counter() - task_started,
|
|
)
|
|
raise
|
|
except Exception as exc:
|
|
total_seconds = time.perf_counter() - page_started
|
|
manifest["pages"][str(page_number)] = {
|
|
"status": "failed",
|
|
"page_number": page_number,
|
|
"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 render_seconds=%.3f ocr_seconds=%.3f export_seconds=%.3f total_seconds=%.3f",
|
|
page_number,
|
|
render_seconds,
|
|
ocr_seconds,
|
|
export_seconds,
|
|
total_seconds,
|
|
)
|
|
if fail_fast:
|
|
manifest["status"] = "failed"
|
|
manifest["updated_at"] = now_iso()
|
|
atomic_write_json(manifest_path, manifest)
|
|
rebuild_combined_outputs(document_dir, manifest)
|
|
raise
|
|
finally:
|
|
if not keep_rendered and render_path.is_file():
|
|
render_path.unlink()
|
|
|
|
state_save_started = time.perf_counter()
|
|
manifest["updated_at"] = now_iso()
|
|
atomic_write_json(manifest_path, manifest)
|
|
rebuild_combined_outputs(document_dir, manifest)
|
|
state_save_seconds = time.perf_counter() - state_save_started
|
|
manifest["pages"][str(page_number)]["state_save_seconds"] = round(state_save_seconds, 3)
|
|
atomic_write_json(manifest_path, manifest)
|
|
|
|
processed_now = position
|
|
average = sum(run_page_times) / len(run_page_times) if run_page_times else None
|
|
remaining = len(pending_indexes) - processed_now
|
|
eta = average * remaining if average is not None else None
|
|
record = manifest["pages"][str(page_number)]
|
|
elapsed_task = time.perf_counter() - task_started
|
|
logger.info(
|
|
"PAGE_FINISHED page=%d status=%s render_seconds=%.3f ocr_seconds=%.3f export_seconds=%.3f state_save_seconds=%.3f page_total_seconds=%.3f task_elapsed_seconds=%.3f eta_seconds=%s progress=%d/%d",
|
|
page_number,
|
|
record["status"],
|
|
record.get("render_seconds", 0.0),
|
|
record.get("ocr_seconds", 0.0),
|
|
record.get("export_seconds", 0.0),
|
|
state_save_seconds,
|
|
record.get("total_seconds", 0.0),
|
|
elapsed_task,
|
|
f"{eta:.3f}" if eta is not None else "unknown",
|
|
processed_now,
|
|
len(pending_indexes),
|
|
)
|
|
|
|
if temporary_render_dir.exists():
|
|
shutil.rmtree(temporary_render_dir, ignore_errors=True)
|
|
|
|
selected_records = [
|
|
manifest.get("pages", {}).get(str(index + 1), {}) for index in selected_indexes
|
|
]
|
|
failed_pages = [
|
|
record.get("page_number") for record in selected_records if record.get("status") == "failed"
|
|
]
|
|
completed_pages = sum(record.get("status") == "completed" for record in selected_records)
|
|
completed_records = [record for record in selected_records if record.get("status") == "completed"]
|
|
render_total = sum(record.get("render_seconds", 0.0) for record in completed_records)
|
|
ocr_total = sum(record.get("ocr_seconds", 0.0) for record in completed_records)
|
|
export_total = sum(record.get("export_seconds", 0.0) for record in completed_records)
|
|
state_save_total = sum(record.get("state_save_seconds", 0.0) for record in selected_records)
|
|
page_total = sum(record.get("total_seconds", 0.0) for record in selected_records)
|
|
average_ocr = ocr_total / completed_pages if completed_pages else 0.0
|
|
average_page = page_total / len(selected_records) if selected_records else 0.0
|
|
manifest["status"] = "completed_with_errors" if failed_pages else "completed"
|
|
manifest["summary"] = {
|
|
"selected_pages": len(selected_indexes),
|
|
"completed_pages": completed_pages,
|
|
"completed_before_resume": completed_before,
|
|
"failed_pages": failed_pages,
|
|
"timing": {
|
|
"pdf_open_seconds": round(pdf_open_seconds, 3),
|
|
"manifest_prepare_seconds": round(manifest_prepare_seconds, 3),
|
|
"render_total_seconds": round(render_total, 3),
|
|
"ocr_total_seconds": round(ocr_total, 3),
|
|
"export_total_seconds": round(export_total, 3),
|
|
"state_save_total_seconds": round(state_save_total, 3),
|
|
"page_total_seconds": round(page_total, 3),
|
|
"average_ocr_seconds": round(average_ocr, 3),
|
|
"average_page_seconds": round(average_page, 3),
|
|
"finalize_seconds": 0.0,
|
|
"task_total_seconds": 0.0,
|
|
},
|
|
}
|
|
finalize_started = time.perf_counter()
|
|
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 selected_pages=%d completed_pages=%d failed_pages=%s pdf_open_seconds=%.3f manifest_prepare_seconds=%.3f render_total_seconds=%.3f ocr_total_seconds=%.3f export_total_seconds=%.3f state_save_total_seconds=%.3f page_total_seconds=%.3f average_ocr_seconds=%.3f average_page_seconds=%.3f finalize_seconds=%.3f task_total_seconds=%.3f",
|
|
manifest["status"],
|
|
len(selected_indexes),
|
|
completed_pages,
|
|
failed_pages,
|
|
pdf_open_seconds,
|
|
manifest_prepare_seconds,
|
|
render_total,
|
|
ocr_total,
|
|
export_total,
|
|
state_save_total,
|
|
page_total,
|
|
average_ocr,
|
|
average_page,
|
|
finalize_seconds,
|
|
task_total,
|
|
)
|
|
|
|
return {
|
|
"document_dir": str(document_dir),
|
|
"manifest_path": str(manifest_path),
|
|
"status": manifest["status"],
|
|
**manifest["summary"],
|
|
}
|
|
finally:
|
|
document.close()
|