498 lines
16 KiB
Python
498 lines
16 KiB
Python
"""Unified suffix-based routing for files and directories."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
import statistics
|
|
import time
|
|
from dataclasses import dataclass
|
|
from datetime import datetime
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
from .output import (
|
|
atomic_write_json,
|
|
image_output_directory,
|
|
pdf_output_root,
|
|
safe_stem,
|
|
save_image_ocr_outputs,
|
|
)
|
|
from .pdf import preflight_pdf, process_pdf
|
|
from .pdf_text import TextLayerPolicy
|
|
from .runtime import PipelineProvider
|
|
|
|
IMAGE_EXTENSIONS = {".png", ".jpg", ".jpeg", ".bmp", ".tiff", ".tif", ".webp"}
|
|
PDF_EXTENSIONS = {".pdf"}
|
|
SUPPORTED_EXTENSIONS = IMAGE_EXTENSIONS | PDF_EXTENSIONS
|
|
|
|
|
|
@dataclass
|
|
class FileProcessResult:
|
|
path: Path
|
|
kind: str
|
|
status: str
|
|
seconds: float
|
|
details: dict[str, Any]
|
|
exit_code: int = 0
|
|
|
|
|
|
def detect_input_kind(path: Path) -> str:
|
|
if path.is_dir():
|
|
return "directory"
|
|
suffix = path.suffix.lower()
|
|
if suffix in IMAGE_EXTENSIONS:
|
|
return "image"
|
|
if suffix in PDF_EXTENSIONS:
|
|
return "pdf"
|
|
return "unsupported"
|
|
|
|
|
|
def discover_supported_files(directory: Path, *, recursive: bool, output_root: Path) -> list[Path]:
|
|
iterator = directory.rglob("*") if recursive else directory.glob("*")
|
|
output_root = output_root.expanduser().resolve()
|
|
files: list[Path] = []
|
|
for path in iterator:
|
|
if not path.is_file() or path.suffix.lower() not in SUPPORTED_EXTENSIONS:
|
|
continue
|
|
resolved = path.resolve()
|
|
try:
|
|
resolved.relative_to(output_root)
|
|
except ValueError:
|
|
files.append(resolved)
|
|
return sorted(files, key=lambda value: str(value).casefold())
|
|
|
|
|
|
def _result_summary(result: list[Any]) -> dict[str, Any]:
|
|
first = result[0]
|
|
blocks = first["parsing_res_list"]
|
|
return {
|
|
"width": first.get("width"),
|
|
"height": first.get("height"),
|
|
"layout_boxes": len(first.get("layout_det_res", {}).get("boxes", [])),
|
|
"parsed_blocks": len(blocks),
|
|
"non_empty_blocks": sum(bool(block.content.strip()) for block in blocks),
|
|
}
|
|
|
|
|
|
def process_image_file(
|
|
path: Path,
|
|
*,
|
|
args,
|
|
provider: PipelineProvider,
|
|
logger: logging.Logger,
|
|
project_root: Path,
|
|
run_warmup: bool,
|
|
batch_root: Path | None,
|
|
) -> FileProcessResult:
|
|
file_started = time.perf_counter()
|
|
logger.info("FILE_ROUTED path=%s kind=image", path)
|
|
pipeline = provider.get()
|
|
predict_kwargs = {
|
|
key: value
|
|
for key, value in {
|
|
"max_new_tokens": getattr(args, "max_new_tokens", None),
|
|
"min_pixels": getattr(args, "min_pixels", None),
|
|
"max_pixels": getattr(args, "max_pixels", None),
|
|
}.items()
|
|
if value is not None
|
|
}
|
|
|
|
warmup_times: list[float] = []
|
|
if run_warmup:
|
|
for index in range(args.warmup):
|
|
started = time.perf_counter()
|
|
pipeline.predict(str(path), **predict_kwargs)
|
|
provider.synchronize()
|
|
elapsed = time.perf_counter() - started
|
|
warmup_times.append(elapsed)
|
|
logger.info(
|
|
"WARMUP_COMPLETED path=%s round=%d/%d seconds=%.3f",
|
|
path,
|
|
index + 1,
|
|
args.warmup,
|
|
elapsed,
|
|
)
|
|
|
|
inference_times: list[float] = []
|
|
result = None
|
|
for index in range(args.rounds):
|
|
provider.synchronize()
|
|
started = time.perf_counter()
|
|
result = pipeline.predict(str(path), **predict_kwargs)
|
|
provider.synchronize()
|
|
elapsed = time.perf_counter() - started
|
|
inference_times.append(elapsed)
|
|
logger.info(
|
|
"INFERENCE_COMPLETED path=%s round=%d/%d seconds=%.3f",
|
|
path,
|
|
index + 1,
|
|
args.rounds,
|
|
elapsed,
|
|
)
|
|
|
|
if not result:
|
|
raise RuntimeError("OCR pipeline 未返回图片结果")
|
|
|
|
summary = _result_summary(result)
|
|
processing_seconds = time.perf_counter() - file_started
|
|
benchmark = {
|
|
"timestamp": datetime.now().astimezone().isoformat(),
|
|
**provider.metadata(),
|
|
"image_path": str(path),
|
|
"image": summary,
|
|
"warmup_seconds": [round(value, 3) for value in warmup_times],
|
|
"inference_seconds": {
|
|
"all": [round(value, 3) for value in inference_times],
|
|
"min": round(min(inference_times), 3),
|
|
"max": round(max(inference_times), 3),
|
|
"mean": round(statistics.fmean(inference_times), 3),
|
|
"median": round(statistics.median(inference_times), 3),
|
|
"stdev": round(statistics.pstdev(inference_times), 3),
|
|
},
|
|
"gpu_memory": provider.gpu_memory(),
|
|
"processing_seconds": round(processing_seconds, 3),
|
|
"export_seconds": 0.0,
|
|
"file_total_seconds": 0.0,
|
|
}
|
|
output_dir = image_output_directory(
|
|
args.output,
|
|
path,
|
|
batch_root=batch_root,
|
|
recursive=args.recursive,
|
|
)
|
|
export_started = time.perf_counter()
|
|
output_paths = save_image_ocr_outputs(
|
|
result[0],
|
|
output_dir,
|
|
input_path=path,
|
|
benchmark=benchmark,
|
|
)
|
|
export_seconds = time.perf_counter() - export_started
|
|
total_seconds = time.perf_counter() - file_started
|
|
benchmark["export_seconds"] = round(export_seconds, 3)
|
|
benchmark["file_total_seconds"] = round(total_seconds, 3)
|
|
atomic_write_json(Path(output_paths["benchmark"]), benchmark)
|
|
if batch_root is None and args.benchmark_json:
|
|
explicit_benchmark = args.benchmark_json.expanduser().resolve()
|
|
atomic_write_json(explicit_benchmark, benchmark)
|
|
output_paths["explicit_benchmark"] = str(explicit_benchmark)
|
|
|
|
logger.info(
|
|
"IMAGE_COMPLETED path=%s width=%s height=%s layout_boxes=%d parsed_blocks=%d inference_mean_seconds=%.3f export_seconds=%.3f file_total_seconds=%.3f output=%s benchmark=%s",
|
|
path,
|
|
summary["width"],
|
|
summary["height"],
|
|
summary["layout_boxes"],
|
|
summary["parsed_blocks"],
|
|
statistics.fmean(inference_times),
|
|
export_seconds,
|
|
total_seconds,
|
|
output_paths["output_dir"],
|
|
output_paths["benchmark"],
|
|
)
|
|
if not args.no_result:
|
|
for index, block in enumerate(result[0]["parsing_res_list"], 1):
|
|
logger.info(
|
|
"OCR_BLOCK path=%s index=%d label=%s bbox=%s content=%s",
|
|
path,
|
|
index,
|
|
block.label,
|
|
block.bbox,
|
|
block.content.replace("\r", "").replace("\n", "\\n"),
|
|
)
|
|
|
|
return FileProcessResult(
|
|
path=path,
|
|
kind="image",
|
|
status="completed",
|
|
seconds=total_seconds,
|
|
details={**summary, **output_paths},
|
|
)
|
|
|
|
|
|
def process_pdf_file(
|
|
path: Path,
|
|
*,
|
|
args,
|
|
provider: PipelineProvider,
|
|
logger: logging.Logger,
|
|
batch_root: Path | None,
|
|
) -> FileProcessResult:
|
|
file_started = time.perf_counter()
|
|
logger.info("FILE_ROUTED path=%s kind=pdf mode=%s", path, args.pdf_mode)
|
|
output_root = pdf_output_root(
|
|
args.output,
|
|
path,
|
|
batch_root=batch_root,
|
|
recursive=args.recursive,
|
|
)
|
|
manifest_exists = (output_root / safe_stem(path.stem) / "manifest.json").is_file()
|
|
# Directory jobs auto-resume existing PDF manifests so rerunning a batch is
|
|
# safe. Single-file jobs still require an explicit --resume.
|
|
resume = args.resume if batch_root is None else manifest_exists and not args.overwrite
|
|
preflight = preflight_pdf(
|
|
pdf_path=path,
|
|
output_root=output_root,
|
|
pages=args.pages,
|
|
dpi=args.dpi,
|
|
password=args.password,
|
|
resume=resume,
|
|
overwrite=args.overwrite,
|
|
)
|
|
logger.info(
|
|
"PDF_PREFLIGHT_COMPLETED path=%s page_count=%d selected_pages=%d output=%s mode=%s",
|
|
path,
|
|
preflight["page_count"],
|
|
len(preflight["selected_pages"]),
|
|
preflight["document_dir"],
|
|
args.pdf_mode,
|
|
)
|
|
policy = TextLayerPolicy(
|
|
min_chars=args.text_min_chars,
|
|
min_printable_ratio=args.text_min_printable_ratio,
|
|
min_content_ratio=args.text_min_content_ratio,
|
|
max_replacement_ratio=args.text_max_replacement_ratio,
|
|
min_chars_per_megapixel=args.text_min_density,
|
|
)
|
|
if args.pdf_mode == "ocr":
|
|
provider.prepare()
|
|
summary = process_pdf(
|
|
provider=provider,
|
|
pdf_path=path,
|
|
output_root=output_root,
|
|
mode=args.pdf_mode,
|
|
text_policy=policy,
|
|
pages=args.pages,
|
|
dpi=args.dpi,
|
|
password=args.password,
|
|
resume=resume,
|
|
overwrite=args.overwrite,
|
|
keep_rendered=args.keep_rendered,
|
|
fail_fast=args.fail_fast,
|
|
predict_kwargs={
|
|
key: value
|
|
for key, value in {
|
|
"max_new_tokens": args.max_new_tokens,
|
|
"min_pixels": args.min_pixels,
|
|
"max_pixels": args.max_pixels,
|
|
}.items()
|
|
if value is not None
|
|
},
|
|
logger=logger,
|
|
)
|
|
total_seconds = time.perf_counter() - file_started
|
|
logger.info(
|
|
"PDF_COMPLETED path=%s status=%s text_pages=%d ocr_pages=%d failed_pages=%s model_used=%s model_initialized_during_task=%s resume=%s file_total_seconds=%.3f output=%s",
|
|
path,
|
|
summary["status"],
|
|
summary["text_pages"],
|
|
summary["ocr_pages"],
|
|
summary["failed_pages"],
|
|
summary["model_used"],
|
|
summary["model_initialized_during_task"],
|
|
resume,
|
|
total_seconds,
|
|
summary["document_dir"],
|
|
)
|
|
return FileProcessResult(
|
|
path=path,
|
|
kind="pdf",
|
|
status=summary["status"],
|
|
seconds=total_seconds,
|
|
details=summary,
|
|
exit_code=0 if not summary["failed_pages"] else 3,
|
|
)
|
|
|
|
|
|
def process_single_file(
|
|
path: Path,
|
|
*,
|
|
args,
|
|
provider: PipelineProvider,
|
|
logger: logging.Logger,
|
|
project_root: Path,
|
|
run_image_warmup: bool,
|
|
batch_root: Path | None = None,
|
|
) -> FileProcessResult:
|
|
path = path.expanduser().resolve()
|
|
if not path.is_file():
|
|
raise FileNotFoundError(f"文件不存在: {path}")
|
|
kind = detect_input_kind(path)
|
|
if kind == "image":
|
|
return process_image_file(
|
|
path,
|
|
args=args,
|
|
provider=provider,
|
|
logger=logger,
|
|
project_root=project_root,
|
|
run_warmup=run_image_warmup,
|
|
batch_root=batch_root,
|
|
)
|
|
if kind == "pdf":
|
|
return process_pdf_file(
|
|
path,
|
|
args=args,
|
|
provider=provider,
|
|
logger=logger,
|
|
batch_root=batch_root,
|
|
)
|
|
supported = ", ".join(sorted(SUPPORTED_EXTENSIONS))
|
|
raise ValueError(f"不支持的文件类型: {path.suffix or '<无后缀>'}；支持: {supported}")
|
|
|
|
|
|
def run_input(args, provider: PipelineProvider, logger: logging.Logger, project_root: Path) -> int:
|
|
program_started = time.perf_counter()
|
|
input_path = args.input.expanduser().resolve()
|
|
kind = detect_input_kind(input_path)
|
|
|
|
if kind != "directory":
|
|
try:
|
|
result = process_single_file(
|
|
input_path,
|
|
args=args,
|
|
provider=provider,
|
|
logger=logger,
|
|
project_root=project_root,
|
|
run_image_warmup=True,
|
|
)
|
|
except KeyboardInterrupt:
|
|
logger.warning("PROGRAM_INTERRUPTED input=%s resume_hint=--resume", input_path)
|
|
return 130
|
|
except Exception as exc:
|
|
logger.exception("FILE_FAILED path=%s error=%s", input_path, exc)
|
|
return 1
|
|
logger.info(
|
|
"PROGRAM_COMPLETED input=%s kind=%s status=%s file_seconds=%.3f program_total_seconds=%.3f",
|
|
result.path,
|
|
result.kind,
|
|
result.status,
|
|
result.seconds,
|
|
time.perf_counter() - program_started,
|
|
)
|
|
return result.exit_code
|
|
|
|
files = discover_supported_files(
|
|
input_path,
|
|
recursive=args.recursive,
|
|
output_root=args.output,
|
|
)
|
|
if not files:
|
|
logger.error("NO_SUPPORTED_FILES directory=%s recursive=%s", input_path, args.recursive)
|
|
return 1
|
|
logger.info(
|
|
"DIRECTORY_PLAN directory=%s recursive=%s files=%d image_files=%d pdf_files=%d",
|
|
input_path,
|
|
args.recursive,
|
|
len(files),
|
|
sum(path.suffix.lower() in IMAGE_EXTENSIONS for path in files),
|
|
sum(path.suffix.lower() in PDF_EXTENSIONS for path in files),
|
|
)
|
|
|
|
results: list[FileProcessResult] = []
|
|
failures: list[dict[str, str]] = []
|
|
image_warmup_pending = True
|
|
for index, path in enumerate(files, 1):
|
|
logger.info("DIRECTORY_PROGRESS_START progress=%d/%d path=%s", index, len(files), path)
|
|
try:
|
|
result = process_single_file(
|
|
path,
|
|
args=args,
|
|
provider=provider,
|
|
logger=logger,
|
|
project_root=project_root,
|
|
run_image_warmup=image_warmup_pending,
|
|
batch_root=input_path,
|
|
)
|
|
results.append(result)
|
|
if result.kind == "image":
|
|
image_warmup_pending = False
|
|
if result.exit_code:
|
|
failures.append({"path": str(path), "error": result.status})
|
|
except KeyboardInterrupt:
|
|
logger.warning("PROGRAM_INTERRUPTED path=%s progress=%d/%d", path, index, len(files))
|
|
return 130
|
|
except Exception as exc:
|
|
failures.append({"path": str(path), "error": f"{type(exc).__name__}: {exc}"})
|
|
logger.exception("FILE_FAILED path=%s progress=%d/%d", path, index, len(files))
|
|
if args.fail_fast:
|
|
break
|
|
logger.info("DIRECTORY_PROGRESS_END progress=%d/%d path=%s", index, len(files), path)
|
|
|
|
image_results = [result for result in results if result.kind == "image"]
|
|
pdf_results = [result for result in results if result.kind == "pdf"]
|
|
total_file_seconds = sum(result.seconds for result in results)
|
|
program_total = time.perf_counter() - program_started
|
|
batch_manifest_path = (
|
|
args.output.expanduser().resolve()
|
|
/ "batches"
|
|
/ f"{safe_stem(input_path.name or 'batch')}-{datetime.now():%Y%m%d-%H%M%S-%f}.json"
|
|
)
|
|
batch_manifest = {
|
|
"input_directory": str(input_path),
|
|
"recursive": args.recursive,
|
|
"device": provider.resolved_device,
|
|
"discovered_files": len(files),
|
|
"completed_files": len(results),
|
|
"failed_files": len(failures),
|
|
"image_files": len(image_results),
|
|
"pdf_files": len(pdf_results),
|
|
"model_init_seconds": round(provider.model_init_seconds, 3),
|
|
"total_file_seconds": round(total_file_seconds, 3),
|
|
"program_total_seconds": round(program_total, 3),
|
|
"results": [
|
|
{
|
|
"path": str(result.path),
|
|
"kind": result.kind,
|
|
"status": result.status,
|
|
"seconds": round(result.seconds, 3),
|
|
"exit_code": result.exit_code,
|
|
"outputs": result.details,
|
|
}
|
|
for result in results
|
|
],
|
|
"failures": failures,
|
|
}
|
|
atomic_write_json(batch_manifest_path, batch_manifest)
|
|
logger.info(
|
|
"DIRECTORY_SUMMARY discovered=%d completed=%d failed=%d images_completed=%d pdfs_completed=%d total_file_seconds=%.3f program_total_seconds=%.3f model_init_seconds=%.3f manifest=%s",
|
|
len(files),
|
|
len(results),
|
|
len(failures),
|
|
len(image_results),
|
|
len(pdf_results),
|
|
total_file_seconds,
|
|
program_total,
|
|
provider.model_init_seconds,
|
|
batch_manifest_path,
|
|
)
|
|
return 0 if not failures else 3
|
|
|
|
|
|
def run_verify(args, provider: PipelineProvider, logger: logging.Logger, project_root: Path) -> int:
|
|
started = time.perf_counter()
|
|
try:
|
|
provider.prepare()
|
|
paddle = provider._paddle
|
|
if provider.config.device == "gpu":
|
|
tensor = paddle.ones([1024, 1024], dtype="float32")
|
|
result = paddle.matmul(tensor, tensor)
|
|
provider.synchronize()
|
|
logger.info("GPU_SMOKE_TEST shape=%s", list(result.shape))
|
|
else:
|
|
from paddle import core
|
|
|
|
logger.info(
|
|
"CPU_SMOKE_TEST onednn=%s mkldnn=%s",
|
|
core.is_compiled_with_onednn(),
|
|
core.is_compiled_with_mkldnn(),
|
|
)
|
|
logger.info(
|
|
"VERIFY_COMPLETED metadata=%s seconds=%.3f",
|
|
provider.metadata(),
|
|
time.perf_counter() - started,
|
|
)
|
|
return 0
|
|
except Exception:
|
|
logger.exception("VERIFY_FAILED")
|
|
return 1
|