93 lines
3.0 KiB
Python
93 lines
3.0 KiB
Python
"""
|
||
批量 OCR 识别 — 多进程并行加速
|
||
|
||
原理：每个进程独立加载一份 pipeline 实例，同时处理不同图片。
|
||
适用场景：一次处理多张图片（如文件夹批量 OCR）。
|
||
|
||
用法:
|
||
python batch_ocr.py <图片目录> [--workers 4] [--threads 10]
|
||
"""
|
||
import time
|
||
import os
|
||
import sys
|
||
import argparse
|
||
from multiprocessing import Pool, cpu_count
|
||
from pathlib import Path
|
||
|
||
|
||
def ocr_single(image_path: str, threads: int) -> dict:
|
||
"""单张图片 OCR（在子进程中执行）"""
|
||
from paddle import core
|
||
core.set_num_threads(threads)
|
||
|
||
from paddleocr import PaddleOCRVL
|
||
|
||
pipeline = PaddleOCRVL(pipeline_version="v1.6")
|
||
|
||
t0 = time.perf_counter()
|
||
result = pipeline.predict(image_path)
|
||
elapsed = time.perf_counter() - t0
|
||
|
||
blocks = []
|
||
for block in result[0]["parsing_res_list"]:
|
||
if block.content.strip():
|
||
blocks.append({"label": block.label, "bbox": block.bbox, "content": block.content})
|
||
|
||
return {
|
||
"path": str(image_path),
|
||
"elapsed": round(elapsed, 2),
|
||
"blocks": blocks,
|
||
}
|
||
|
||
|
||
def main():
|
||
parser = argparse.ArgumentParser(description="Batch OCR with multiprocessing")
|
||
parser.add_argument("dir", type=str, help="图片目录")
|
||
parser.add_argument("--workers", type=int, default=4, help="并行进程数 (默认 4)")
|
||
parser.add_argument("--threads", type=int, default=None, help="每进程线程数 (默认: 总核心/workers)")
|
||
args = parser.parse_args()
|
||
|
||
image_dir = Path(args.dir)
|
||
if not image_dir.is_dir():
|
||
print(f"[ERROR] 目录不存在: {args.dir}")
|
||
sys.exit(1)
|
||
|
||
images = list(image_dir.glob("*.png")) + list(image_dir.glob("*.jpg")) + list(image_dir.glob("*.jpeg"))
|
||
if not images:
|
||
print(f"[ERROR] 目录中没有图片: {args.dir}")
|
||
sys.exit(1)
|
||
|
||
workers = min(args.workers, len(images))
|
||
threads = args.threads or max(1, cpu_count() // workers)
|
||
|
||
print(f"图片数量: {len(images)}")
|
||
print(f"并行进程: {workers}")
|
||
print(f"每进程线程: {threads}")
|
||
print(f"总核心利用: {workers * threads} / {cpu_count()}")
|
||
print("=" * 60)
|
||
|
||
t0 = time.perf_counter()
|
||
|
||
# 每个子进程独立加载模型 + 推理
|
||
with Pool(workers) as pool:
|
||
tasks = [(str(img), threads) for img in images]
|
||
results = pool.starmap(ocr_single, tasks)
|
||
|
||
total_elapsed = time.perf_counter() - t0
|
||
|
||
# 输出结果
|
||
print("\n" + "=" * 60)
|
||
for r in results:
|
||
print(f"\n[文件] {r['path']} ({r['elapsed']:.1f}s)")
|
||
for block in r["blocks"]:
|
||
print(f" [{block['label']}] {block['content'][:60]}{'...' if len(block['content']) > 60 else ''}")
|
||
|
||
print("\n" + "=" * 60)
|
||
print(f"总图片: {len(images)} | 总耗时: {total_elapsed:.1f}s")
|
||
print(f"平均每图: {total_elapsed / len(images):.1f}s")
|
||
print(f"单进程串行预计: {sum(r['elapsed'] for r in results):.1f}s")
|
||
print(f"并行加速比: {sum(r['elapsed'] for r in results) / total_elapsed:.2f}x")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main() |