141 lines
7.4 KiB
Python
141 lines
7.4 KiB
Python
"""Path-first unified CLI for PP-OCRv6."""
|
||
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import sys
|
||
from pathlib import Path
|
||
|
||
from .commands import run_input, run_verify
|
||
from .logging_utils import default_log_path, setup_run_logger
|
||
from .runtime import PipelineProvider, RuntimeConfig
|
||
|
||
PROJECT_ROOT = Path(__file__).resolve().parent.parent
|
||
LEGACY_COMMANDS = {"image", "pdf", "batch"}
|
||
|
||
|
||
def _add_bool_option(parser: argparse.ArgumentParser, name: str, default: bool, help_text: str) -> None:
|
||
parser.add_argument(
|
||
f"--{name}",
|
||
action=argparse.BooleanOptionalAction,
|
||
default=default,
|
||
help=help_text,
|
||
)
|
||
|
||
|
||
def _add_device_options(parser: argparse.ArgumentParser, default: str | None) -> None:
|
||
parser.add_argument("--device", choices=("cpu", "gpu"), default=default or "cpu", help="运行设备")
|
||
parser.add_argument("--device-id", type=int, default=0, help="GPU 编号")
|
||
parser.add_argument("--threads", type=int, default=None, help="CPU 线程数")
|
||
parser.add_argument("--log-file", type=Path, default=None, help="日志文件路径")
|
||
parser.add_argument("--verbose", action="store_true", help="输出详细日志")
|
||
|
||
|
||
def _add_model_options(parser: argparse.ArgumentParser) -> None:
|
||
parser.add_argument(
|
||
"--model-size",
|
||
"--model",
|
||
dest="model_size",
|
||
choices=("tiny", "small", "medium"),
|
||
default="medium",
|
||
help="PP-OCRv6 检测与识别模型规格",
|
||
)
|
||
parser.add_argument("--lang", default="ch", help="结果元数据中的语言标识，默认中文/中英混合")
|
||
parser.add_argument("--text-rec-score-thresh", type=float, default=0.0, help="识别置信度阈值")
|
||
parser.add_argument("--text-det-limit-side-len", type=int, default=None, help="文本检测边长限制")
|
||
parser.add_argument("--text-det-limit-type", choices=("min", "max"), default=None, help="检测边长限制方式")
|
||
parser.add_argument("--text-det-thresh", type=float, default=None, help="文本检测像素阈值")
|
||
parser.add_argument("--text-det-box-thresh", type=float, default=None, help="文本框阈值")
|
||
parser.add_argument("--text-det-unclip-ratio", type=float, default=None, help="文本框扩张比例")
|
||
parser.add_argument("--text-recognition-batch-size", type=int, default=6, help="文本识别批大小")
|
||
parser.add_argument("--return-word-box", action="store_true", help="返回单词级坐标")
|
||
_add_bool_option(parser, "doc-orientation-classify", True, "启用文档方向分类")
|
||
_add_bool_option(parser, "doc-unwarping", False, "启用文档去畸变")
|
||
_add_bool_option(parser, "textline-orientation", True, "启用文本行方向分类")
|
||
|
||
|
||
def build_input_parser(device_override: str | None = None) -> argparse.ArgumentParser:
|
||
parser = argparse.ArgumentParser(
|
||
prog="ocr.py",
|
||
description="使用 PP-OCRv6 自动处理图片、PDF 或目录",
|
||
formatter_class=argparse.ArgumentDefaultsHelpFormatter,
|
||
)
|
||
parser.add_argument("input", type=Path, help="图片、PDF 或目录")
|
||
parser.add_argument("--recursive", action="store_true", help="目录模式递归扫描子目录")
|
||
parser.add_argument("--output", type=Path, default=PROJECT_ROOT / "outputs", help="输出根目录")
|
||
parser.add_argument("--fail-fast", action="store_true", help="单文件失败后立即停止目录任务")
|
||
parser.add_argument("--warmup", type=int, default=0, help="首张图片预热轮数")
|
||
parser.add_argument("--rounds", type=int, default=1, help="每张图片推理轮数")
|
||
parser.add_argument("--benchmark-json", type=Path, default=None, help="额外复制单图 Benchmark JSON")
|
||
parser.add_argument("--no-result", action="store_true", help="不在日志中逐行记录识别文字")
|
||
parser.add_argument("--save-visualization", action="store_true", help="保存 OCR 可视化图片")
|
||
parser.add_argument("--save-raw-result", action="store_true", help="保存 PaddleOCR 原始 JSON")
|
||
parser.add_argument("--pdf-mode", "--mode", dest="pdf_mode", choices=("hybrid", "text", "ocr"), default="hybrid", help="PDF 处理模式")
|
||
parser.add_argument("--pages", help="PDF 页码范围，例如 1-5,8,10-")
|
||
parser.add_argument("--dpi", type=int, default=144, help="PDF OCR 页面渲染 DPI")
|
||
parser.add_argument("--password", help="PDF 密码")
|
||
parser.add_argument("--resume", action="store_true", help="PDF 断点续传")
|
||
parser.add_argument("--overwrite", action="store_true", help="覆盖已有 PDF 输出")
|
||
parser.add_argument("--keep-rendered", action="store_true", help="保留 OCR 页面 PNG")
|
||
parser.add_argument("--text-min-chars", type=int, default=50, help="有效文本最小字符数")
|
||
parser.add_argument("--text-min-printable-ratio", type=float, default=0.85, help="可打印字符比例阈值")
|
||
parser.add_argument("--text-min-content-ratio", type=float, default=0.60, help="字母/数字/CJK 比例阈值")
|
||
parser.add_argument("--text-max-replacement-ratio", type=float, default=0.02, help="替换字符最大比例")
|
||
parser.add_argument("--text-min-density", type=float, default=25.0, help="文本密度阈值")
|
||
_add_model_options(parser)
|
||
_add_device_options(parser, device_override)
|
||
return parser
|
||
|
||
|
||
def build_verify_parser(device_override: str | None = None) -> argparse.ArgumentParser:
|
||
parser = argparse.ArgumentParser(prog="ocr.py verify", description="验证 CPU/GPU Paddle 环境")
|
||
_add_model_options(parser)
|
||
_add_device_options(parser, device_override)
|
||
return parser
|
||
|
||
|
||
def normalize_argv(argv: list[str]) -> tuple[str, list[str]]:
|
||
if argv and argv[0] == "verify":
|
||
return "verify", argv[1:]
|
||
if argv and argv[0] in LEGACY_COMMANDS:
|
||
return "input", argv[1:]
|
||
return "input", argv
|
||
|
||
|
||
def main(device: str | None = None, argv: list[str] | None = None) -> int:
|
||
raw_argv = list(sys.argv[1:] if argv is None else argv)
|
||
command, normalized = normalize_argv(raw_argv)
|
||
parser = build_verify_parser(device) if command == "verify" else build_input_parser(device)
|
||
args = parser.parse_args(normalized)
|
||
if device is not None:
|
||
args.device = device
|
||
stem = "verify" if command == "verify" else (args.input.stem or args.input.name or "input")
|
||
category = "verify" if command == "verify" else "input"
|
||
log_file = args.log_file or default_log_path(PROJECT_ROOT, category, stem, device=args.device)
|
||
logger = setup_run_logger(f"ppocrv6.{category}.{args.device}", log_file, verbose=args.verbose)
|
||
provider = PipelineProvider(
|
||
RuntimeConfig(
|
||
device=args.device,
|
||
threads=args.threads,
|
||
device_id=args.device_id,
|
||
lang=args.lang,
|
||
model_size=args.model_size,
|
||
use_doc_orientation_classify=args.doc_orientation_classify,
|
||
use_doc_unwarping=args.doc_unwarping,
|
||
use_textline_orientation=args.textline_orientation,
|
||
text_recognition_batch_size=args.text_recognition_batch_size,
|
||
),
|
||
logger,
|
||
)
|
||
try:
|
||
if command == "verify":
|
||
return run_verify(args, provider, logger, PROJECT_ROOT)
|
||
if args.warmup < 0 or args.rounds < 1:
|
||
raise ValueError("--warmup 必须 >= 0，--rounds 必须 >= 1")
|
||
if args.text_recognition_batch_size < 1:
|
||
raise ValueError("--text-recognition-batch-size 必须 >= 1")
|
||
return run_input(args, provider, logger, PROJECT_ROOT)
|
||
except Exception as exc:
|
||
logger.exception("COMMAND_FAILED command=%s error=%s", command, exc)
|
||
return 1
|