311 lines
13 KiB
Python
311 lines
13 KiB
Python
"""Unified suffix-based routing for PP-OCRv6 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 .result_adapter import adapt_ocr_result
|
|
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"
|
|
if path.suffix.lower() in IMAGE_EXTENSIONS:
|
|
return "image"
|
|
if path.suffix.lower() 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 _predict_kwargs(args) -> dict[str, Any]:
|
|
return {
|
|
key: value
|
|
for key, value in {
|
|
"text_det_limit_side_len": args.text_det_limit_side_len,
|
|
"text_det_limit_type": args.text_det_limit_type,
|
|
"text_det_thresh": args.text_det_thresh,
|
|
"text_det_box_thresh": args.text_det_box_thresh,
|
|
"text_det_unclip_ratio": args.text_det_unclip_ratio,
|
|
"text_rec_score_thresh": args.text_rec_score_thresh,
|
|
"return_word_box": args.return_word_box,
|
|
}.items()
|
|
if value is not None
|
|
}
|
|
|
|
|
|
def process_image_file(
|
|
path: Path,
|
|
*,
|
|
args,
|
|
provider: PipelineProvider,
|
|
logger: logging.Logger,
|
|
project_root: Path,
|
|
run_warmup: bool,
|
|
batch_root: Path | None,
|
|
) -> FileProcessResult:
|
|
del project_root
|
|
file_started = time.perf_counter()
|
|
logger.info("FILE_ROUTED path=%s kind=image", path)
|
|
pipeline = provider.get()
|
|
predict_kwargs = _predict_kwargs(args)
|
|
|
|
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()
|
|
results = pipeline.predict(str(path), **predict_kwargs)
|
|
provider.synchronize()
|
|
elapsed = time.perf_counter() - started
|
|
inference_times.append(elapsed)
|
|
if not results:
|
|
raise RuntimeError("PP-OCRv6 未返回图片结果")
|
|
result = results[0]
|
|
logger.info("INFERENCE_COMPLETED path=%s round=%d/%d seconds=%.3f", path, index + 1, args.rounds, elapsed)
|
|
assert result is not None
|
|
|
|
adapt_started = time.perf_counter()
|
|
normalized = adapt_ocr_result(
|
|
result,
|
|
input_path=path,
|
|
source_type="image_ocr",
|
|
language=provider.config.lang,
|
|
detection_model=provider.config.text_detection_model_name,
|
|
recognition_model=provider.config.text_recognition_model_name,
|
|
model_size=provider.config.model_size,
|
|
)
|
|
adapt_seconds = time.perf_counter() - adapt_started
|
|
summary = normalized["summary"]
|
|
benchmark = {
|
|
"timestamp": datetime.now().astimezone().isoformat(),
|
|
**provider.metadata(),
|
|
"image_path": str(path),
|
|
"image": normalized["image"],
|
|
"ocr_summary": 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),
|
|
},
|
|
"result_adapt_seconds": round(adapt_seconds, 3),
|
|
"gpu_memory": provider.gpu_memory(),
|
|
"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,
|
|
normalized,
|
|
output_dir,
|
|
input_path=path,
|
|
benchmark=benchmark,
|
|
save_raw_result=args.save_raw_result,
|
|
save_visualization=args.save_visualization,
|
|
)
|
|
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 = args.benchmark_json.expanduser().resolve()
|
|
atomic_write_json(explicit, benchmark)
|
|
output_paths["explicit_benchmark"] = str(explicit)
|
|
|
|
logger.info(
|
|
"IMAGE_COMPLETED path=%s detected_lines=%d non_empty_lines=%d mean_score=%s inference_mean_seconds=%.3f file_total_seconds=%.3f output=%s",
|
|
path,
|
|
summary["detected_lines"],
|
|
summary["non_empty_lines"],
|
|
summary["mean_score"],
|
|
statistics.fmean(inference_times),
|
|
total_seconds,
|
|
output_dir,
|
|
)
|
|
if not args.no_result:
|
|
for line in normalized["lines"]:
|
|
logger.info("OCR_LINE path=%s index=%d score=%s text=%s", path, line["index"], line["score"], line["text"].replace("\n", "\\n"))
|
|
return FileProcessResult(path, "image", "completed", total_seconds, {**summary, **output_paths})
|
|
|
|
|
|
def process_pdf_file(path: Path, *, args, provider: PipelineProvider, logger: logging.Logger, batch_root: Path | None) -> FileProcessResult:
|
|
started = time.perf_counter()
|
|
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()
|
|
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", path, preflight["page_count"], len(preflight["selected_pages"]))
|
|
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=_predict_kwargs(args),
|
|
logger=logger,
|
|
)
|
|
total = time.perf_counter() - started
|
|
logger.info("PDF_COMPLETED path=%s status=%s text_pages=%d ocr_pages=%d failed_pages=%s total_seconds=%.3f", path, summary["status"], summary["text_pages"], summary["ocr_pages"], summary["failed_pages"], total)
|
|
return FileProcessResult(path, "pdf", summary["status"], total, summary, 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)
|
|
raise ValueError(f"不支持的文件类型: {path.suffix or '<无后缀>'}")
|
|
|
|
|
|
def run_input(args, provider: PipelineProvider, logger: logging.Logger, project_root: Path) -> int:
|
|
program_started = time.perf_counter()
|
|
input_path = args.input.expanduser().resolve()
|
|
if detect_input_kind(input_path) != "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", 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 total_seconds=%.3f", result.path, result.kind, result.status, 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", input_path)
|
|
return 1
|
|
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:
|
|
return 130
|
|
except Exception as exc:
|
|
failures.append({"path": str(path), "error": f"{type(exc).__name__}: {exc}"})
|
|
logger.exception("FILE_FAILED path=%s", path)
|
|
if args.fail_fast:
|
|
break
|
|
program_total = time.perf_counter() - program_started
|
|
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"
|
|
atomic_write_json(
|
|
manifest_path,
|
|
{
|
|
"input_directory": str(input_path),
|
|
"recursive": args.recursive,
|
|
"device": provider.resolved_device,
|
|
"model_config": provider.model_config(),
|
|
"discovered_files": len(files),
|
|
"completed_files": len(results),
|
|
"failed_files": len(failures),
|
|
"program_total_seconds": round(program_total, 3),
|
|
"results": [{"path": str(item.path), "kind": item.kind, "status": item.status, "seconds": round(item.seconds, 3), "outputs": item.details} for item in results],
|
|
"failures": failures,
|
|
},
|
|
)
|
|
logger.info("DIRECTORY_SUMMARY discovered=%d completed=%d failed=%d manifest=%s", len(files), len(results), len(failures), manifest_path)
|
|
return 0 if not failures else 3
|
|
|
|
|
|
def run_verify(args, provider: PipelineProvider, logger: logging.Logger, project_root: Path) -> int:
|
|
del args, project_root
|
|
started = time.perf_counter()
|
|
try:
|
|
provider.prepare()
|
|
paddle = provider._paddle
|
|
if provider.config.device == "gpu":
|
|
result = paddle.matmul(paddle.ones([1024, 1024]), paddle.ones([1024, 1024]))
|
|
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
|