PaddleOCR-VL-1.6_Demo/ocr.py

84 lines
2.8 KiB
Python
Raw Permalink Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

"""Unified launcher.
Examples:
python ocr.py data/images/手写01.png --device cpu
python ocr.py data/documents/sample.pdf --device cpu --pdf-mode hybrid
python ocr.py data/ --recursive --device cpu
python ocr.py benchmark --profile smoke --device cpu
python ocr.py verify --device gpu
The launcher deliberately executes the selected isolated uv project, so CPU
and GPU Paddle packages never share the same virtual environment.
"""
from __future__ import annotations
import shutil
import subprocess
import sys
from pathlib import Path
ROOT = Path(__file__).resolve().parent
def _requested_device(argv: list[str]) -> str:
for index, value in enumerate(argv):
if value == "--device":
if index + 1 >= len(argv):
raise SystemExit("--device 需要 cpu 或 gpu")
return argv[index + 1].lower()
if value.startswith("--device="):
return value.split("=", 1)[1].lower()
return "cpu"
def main() -> int:
if len(sys.argv) > 1 and sys.argv[1] == "benchmark":
command = [sys.executable, str(ROOT / "benchmark.py"), *sys.argv[2:]]
return subprocess.run(command, cwd=ROOT).returncode
try:
device = _requested_device(sys.argv[1:])
except SystemExit as exc:
print(exc, file=sys.stderr)
return 2
if device not in {"cpu", "gpu"}:
print(f"不支持的设备: {device},可选 cpu/gpu", file=sys.stderr)
return 2
project = ROOT / device
runner = project / "runner.py"
if sys.platform == "win32":
environment_python = project / ".venv" / "Scripts" / "python.exe"
else:
environment_python = project / ".venv" / "bin" / "python"
if device == "gpu":
ready_marker = project / ".gpu-ready"
if not ready_marker.is_file() or not environment_python.is_file():
print(
"GPU 环境尚未安装完成。请根据目标 CUDA 版本运行:\n"
" python gpu/setup_env.py --cuda cu118\n"
"或:\n"
" python gpu/setup_env.py --cuda cu126",
file=sys.stderr,
)
return 1
command = [str(environment_python), str(runner), *sys.argv[1:]]
return subprocess.run(command, cwd=ROOT).returncode
if environment_python.is_file():
command = [str(environment_python), str(runner), *sys.argv[1:]]
return subprocess.run(command, cwd=ROOT).returncode
uv = shutil.which("uv")
if not uv:
print("未找到 CPU 虚拟环境和 uv请先安装 uv。", file=sys.stderr)
return 1
command = [uv, "run", "--project", str(project), "python", str(runner), *sys.argv[1:]]
return subprocess.run(command, cwd=ROOT).returncode
if __name__ == "__main__":
raise SystemExit(main())