330 lines
12 KiB
Python
330 lines
12 KiB
Python
"""System-friendly multiprocessing batch OCR with structured timing logs."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import logging
|
|
import os
|
|
import random
|
|
import sys
|
|
import time
|
|
from logging.handlers import QueueHandler, QueueListener
|
|
from multiprocessing import Manager, Pool, cpu_count
|
|
from pathlib import Path
|
|
|
|
from ocr_logging import default_log_path, setup_run_logger
|
|
|
|
PROJECT_ROOT = Path(__file__).resolve().parent
|
|
_WORKER_LOG_QUEUE = None
|
|
_WORKER_INIT_METRICS: dict = {}
|
|
|
|
|
|
def _worker_logger() -> logging.Logger:
|
|
logger = logging.getLogger(f"ocr.batch.worker.{os.getpid()}")
|
|
if logger.handlers:
|
|
return logger
|
|
logger.setLevel(logging.INFO)
|
|
logger.propagate = False
|
|
if _WORKER_LOG_QUEUE is not None:
|
|
logger.addHandler(QueueHandler(_WORKER_LOG_QUEUE))
|
|
return logger
|
|
|
|
|
|
def _init_worker(threads: int, stagger_max: float, log_queue) -> None:
|
|
"""Stagger startup, lower process priority, and load one model per worker."""
|
|
global _pipeline, _WORKER_LOG_QUEUE, _WORKER_INIT_METRICS
|
|
_WORKER_LOG_QUEUE = log_queue
|
|
logger = _worker_logger()
|
|
worker_started = time.perf_counter()
|
|
delay = random.uniform(0, stagger_max)
|
|
logger.info("WORKER_START threads=%d stagger_delay_seconds=%.3f", threads, delay)
|
|
time.sleep(delay)
|
|
|
|
import_started = time.perf_counter()
|
|
from paddle import core
|
|
core.set_num_threads(threads)
|
|
import_seconds = time.perf_counter() - import_started
|
|
|
|
try:
|
|
import psutil
|
|
|
|
process = psutil.Process()
|
|
if sys.platform == "win32":
|
|
process.nice(psutil.BELOW_NORMAL_PRIORITY_CLASS)
|
|
else:
|
|
process.nice(10)
|
|
priority_status = "lowered"
|
|
except Exception as exc:
|
|
priority_status = f"unchanged:{type(exc).__name__}"
|
|
|
|
model_started = time.perf_counter()
|
|
from paddleocr import PaddleOCRVL
|
|
|
|
_pipeline = PaddleOCRVL(pipeline_version="v1.6", device="cpu")
|
|
model_seconds = time.perf_counter() - model_started
|
|
startup_total = time.perf_counter() - worker_started
|
|
_WORKER_INIT_METRICS = {
|
|
"pid": os.getpid(),
|
|
"threads": threads,
|
|
"stagger_delay_seconds": round(delay, 3),
|
|
"import_seconds": round(import_seconds, 3),
|
|
"model_init_seconds": round(model_seconds, 3),
|
|
"startup_total_seconds": round(startup_total, 3),
|
|
"priority": priority_status,
|
|
}
|
|
logger.info(
|
|
"WORKER_READY threads=%d import_seconds=%.3f model_init_seconds=%.3f startup_total_seconds=%.3f priority=%s",
|
|
threads,
|
|
import_seconds,
|
|
model_seconds,
|
|
startup_total,
|
|
priority_status,
|
|
)
|
|
|
|
|
|
def _ocr_task(image_path: str) -> dict:
|
|
global _pipeline, _WORKER_INIT_METRICS
|
|
logger = _worker_logger()
|
|
started = time.perf_counter()
|
|
logger.info("IMAGE_START path=%s", image_path)
|
|
try:
|
|
result = _pipeline.predict(image_path)
|
|
elapsed = time.perf_counter() - started
|
|
first = result[0]
|
|
blocks = [
|
|
{"label": block.label, "bbox": block.bbox, "content": block.content}
|
|
for block in first["parsing_res_list"]
|
|
if block.content.strip()
|
|
]
|
|
response = {
|
|
"path": image_path,
|
|
"status": "completed",
|
|
"elapsed": round(elapsed, 3),
|
|
"width": first.get("width"),
|
|
"height": first.get("height"),
|
|
"layout_boxes": len(first["layout_det_res"]["boxes"]),
|
|
"parsed_blocks": len(first["parsing_res_list"]),
|
|
"blocks": blocks,
|
|
"worker_pid": os.getpid(),
|
|
"worker_init": _WORKER_INIT_METRICS,
|
|
}
|
|
logger.info(
|
|
"IMAGE_COMPLETED path=%s seconds=%.3f width=%s height=%s layout_boxes=%d parsed_blocks=%d non_empty_blocks=%d",
|
|
image_path,
|
|
elapsed,
|
|
response["width"],
|
|
response["height"],
|
|
response["layout_boxes"],
|
|
response["parsed_blocks"],
|
|
len(blocks),
|
|
)
|
|
return response
|
|
except Exception as exc:
|
|
elapsed = time.perf_counter() - started
|
|
logger.exception("IMAGE_FAILED path=%s seconds=%.3f error=%s", image_path, elapsed, exc)
|
|
return {
|
|
"path": image_path,
|
|
"status": "failed",
|
|
"elapsed": round(elapsed, 3),
|
|
"error": f"{type(exc).__name__}: {exc}",
|
|
"blocks": [],
|
|
"worker_pid": os.getpid(),
|
|
"worker_init": _WORKER_INIT_METRICS,
|
|
}
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser(
|
|
description="批量 OCR — 多进程并行（系统友好版）",
|
|
formatter_class=argparse.ArgumentDefaultsHelpFormatter,
|
|
)
|
|
parser.add_argument("dir", type=Path, help="图片目录")
|
|
parser.add_argument("--workers", type=int, default=2, help="并行进程数")
|
|
parser.add_argument("--threads", type=int, default=None, help="每进程线程数")
|
|
parser.add_argument("--stagger", type=float, default=15.0, help="Worker 启动错峰窗口秒数")
|
|
parser.add_argument("--log-file", type=Path, default=None, help="日志文件路径")
|
|
parser.add_argument("--verbose", action="store_true", help="输出详细日志")
|
|
parser.add_argument("--no-result", action="store_true", help="不记录 OCR 文本块")
|
|
return parser.parse_args()
|
|
|
|
|
|
def main() -> int:
|
|
program_started = time.perf_counter()
|
|
args = parse_args()
|
|
image_dir = args.dir.expanduser().resolve()
|
|
log_file = args.log_file or default_log_path(PROJECT_ROOT, "batch", image_dir.name, device="cpu")
|
|
logger = setup_run_logger("ocr.batch.main", log_file, verbose=args.verbose)
|
|
|
|
if not image_dir.is_dir():
|
|
logger.error("INPUT_DIRECTORY_NOT_FOUND path=%s", image_dir)
|
|
return 1
|
|
if args.workers < 1 or args.stagger < 0:
|
|
logger.error("INVALID_ARGUMENT workers=%d stagger=%.3f", args.workers, args.stagger)
|
|
return 2
|
|
|
|
scan_started = time.perf_counter()
|
|
extensions = ("*.png", "*.jpg", "*.jpeg", "*.bmp", "*.tiff", "*.tif", "*.webp")
|
|
images = sorted(path for extension in extensions for path in image_dir.glob(extension))
|
|
scan_seconds = time.perf_counter() - scan_started
|
|
if not images:
|
|
logger.error("NO_IMAGES_FOUND path=%s scan_seconds=%.3f", image_dir, scan_seconds)
|
|
return 1
|
|
|
|
total_cores = cpu_count()
|
|
workers = min(args.workers, len(images))
|
|
threads = args.threads or max(1, (total_cores - 1) // workers)
|
|
if threads < 1:
|
|
logger.error("INVALID_ARGUMENT threads=%d", threads)
|
|
return 2
|
|
total_cpu_used = workers * threads
|
|
estimated_mem = workers * 2.0 + 2
|
|
|
|
try:
|
|
import psutil
|
|
|
|
available_gb = psutil.virtual_memory().available / (1024**3)
|
|
except ImportError:
|
|
available_gb = None
|
|
|
|
logger.info(
|
|
"PROGRAM_STARTED directory=%s image_count=%d scan_seconds=%.3f workers=%d threads_per_worker=%d total_cores=%d planned_threads=%d reserved_cores=%d stagger_seconds=%.3f estimated_memory_gb=%.1f available_memory_gb=%s",
|
|
image_dir,
|
|
len(images),
|
|
scan_seconds,
|
|
workers,
|
|
threads,
|
|
total_cores,
|
|
total_cpu_used,
|
|
max(0, total_cores - total_cpu_used),
|
|
args.stagger,
|
|
estimated_mem,
|
|
f"{available_gb:.1f}" if available_gb is not None else "unknown",
|
|
)
|
|
|
|
if available_gb is not None and available_gb <= estimated_mem:
|
|
logger.warning(
|
|
"MEMORY_PRESSURE estimated_memory_gb=%.1f available_memory_gb=%.1f recommendation=reduce_workers",
|
|
estimated_mem,
|
|
available_gb,
|
|
)
|
|
response = input("可用内存可能不足，是否继续？[y/N] ").strip().lower()
|
|
if response != "y":
|
|
logger.warning("PROGRAM_CANCELLED_BY_USER")
|
|
return 0
|
|
|
|
pool_started = time.perf_counter()
|
|
results: list[dict] = []
|
|
try:
|
|
with Manager() as manager:
|
|
log_queue = manager.Queue()
|
|
listener_handlers = tuple(logger.handlers)
|
|
listener = QueueListener(log_queue, *listener_handlers, respect_handler_level=True)
|
|
listener.start()
|
|
try:
|
|
with Pool(
|
|
processes=workers,
|
|
initializer=_init_worker,
|
|
initargs=(threads, args.stagger, log_queue),
|
|
) as pool:
|
|
for completed, result in enumerate(
|
|
pool.imap_unordered(_ocr_task, [str(path) for path in images], chunksize=1),
|
|
start=1,
|
|
):
|
|
results.append(result)
|
|
logger.info(
|
|
"BATCH_PROGRESS completed=%d total=%d path=%s status=%s image_seconds=%.3f worker_pid=%s",
|
|
completed,
|
|
len(images),
|
|
result["path"],
|
|
result["status"],
|
|
result["elapsed"],
|
|
result.get("worker_pid"),
|
|
)
|
|
finally:
|
|
listener.stop()
|
|
except KeyboardInterrupt:
|
|
logger.warning("PROGRAM_INTERRUPTED elapsed_seconds=%.3f", time.perf_counter() - program_started)
|
|
return 130
|
|
except Exception as exc:
|
|
logger.exception("POOL_FAILED error=%s elapsed_seconds=%.3f", exc, time.perf_counter() - program_started)
|
|
return 1
|
|
|
|
pool_seconds = time.perf_counter() - pool_started
|
|
completed_results = [result for result in results if result["status"] == "completed"]
|
|
failed_results = [result for result in results if result["status"] == "failed"]
|
|
worker_metrics = {
|
|
result["worker_pid"]: result.get("worker_init", {})
|
|
for result in results
|
|
if result.get("worker_pid") is not None
|
|
}
|
|
worker_model_init_total = sum(
|
|
metrics.get("model_init_seconds", 0.0) for metrics in worker_metrics.values()
|
|
)
|
|
worker_model_init_average = (
|
|
worker_model_init_total / len(worker_metrics) if worker_metrics else 0.0
|
|
)
|
|
serial_estimate = sum(result["elapsed"] for result in results)
|
|
average = serial_estimate / len(results) if results else 0.0
|
|
speedup = serial_estimate / pool_seconds if pool_seconds else 0.0
|
|
program_total = time.perf_counter() - program_started
|
|
|
|
for worker_pid, metrics in sorted(worker_metrics.items()):
|
|
logger.info(
|
|
"WORKER_SUMMARY pid=%s threads=%s stagger_delay_seconds=%s import_seconds=%s model_init_seconds=%s startup_total_seconds=%s priority=%s",
|
|
worker_pid,
|
|
metrics.get("threads"),
|
|
metrics.get("stagger_delay_seconds"),
|
|
metrics.get("import_seconds"),
|
|
metrics.get("model_init_seconds"),
|
|
metrics.get("startup_total_seconds"),
|
|
metrics.get("priority"),
|
|
)
|
|
|
|
for result in sorted(results, key=lambda item: item["path"]):
|
|
if result["status"] == "completed":
|
|
logger.info(
|
|
"IMAGE_SUMMARY path=%s seconds=%.3f width=%s height=%s layout_boxes=%d parsed_blocks=%d",
|
|
result["path"],
|
|
result["elapsed"],
|
|
result["width"],
|
|
result["height"],
|
|
result["layout_boxes"],
|
|
result["parsed_blocks"],
|
|
)
|
|
if not args.no_result:
|
|
for index, block in enumerate(result["blocks"], start=1):
|
|
logger.info(
|
|
"OCR_BLOCK path=%s index=%d label=%s bbox=%s content=%s",
|
|
result["path"],
|
|
index,
|
|
block["label"],
|
|
block["bbox"],
|
|
block["content"].replace("\r", "").replace("\n", "\\n"),
|
|
)
|
|
else:
|
|
logger.error("IMAGE_SUMMARY path=%s status=failed error=%s", result["path"], result["error"])
|
|
|
|
logger.info(
|
|
"BATCH_SUMMARY image_count=%d completed=%d failed=%d scan_seconds=%.3f pool_seconds=%.3f worker_count=%d worker_model_init_total_seconds=%.3f worker_model_init_average_seconds=%.3f serial_estimate_seconds=%.3f average_image_seconds=%.3f speedup=%.3f program_total_seconds=%.3f workers=%d threads_per_worker=%d log=%s",
|
|
len(images),
|
|
len(completed_results),
|
|
len(failed_results),
|
|
scan_seconds,
|
|
pool_seconds,
|
|
len(worker_metrics),
|
|
worker_model_init_total,
|
|
worker_model_init_average,
|
|
serial_estimate,
|
|
average,
|
|
speedup,
|
|
program_total,
|
|
workers,
|
|
threads,
|
|
log_file.resolve(),
|
|
)
|
|
return 0 if not failed_results else 3
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|