5090关于使用paddleocr和gpu进行ocr的demo


配置

gpu: 5090
cuda: 12.9 -- 务必使用12.9,此前测试使用12.6 显卡根本无法正常返回数据,5090和126应该有适配问题。

demo代码

#!/usr/bin/env python3
import argparse
import json
import sys
from pathlib import Path
from typing import Any
import paddle
from paddleocr import PaddleOCR
def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="使用 PaddleOCR 和 NVIDIA GPU 识别图片或 PDF"
    )
    parser.add_argument("input", type=Path, help="待识别的图片或 PDF 路径")
    parser.add_argument("--gpu", type=int, default=0, help="GPU 编号,默认 0")
    parser.add_argument("--output", type=Path, help="JSON 输出文件,默认输出到终端")
    return parser.parse_args()
def check_gpu(gpu_id: int) -> str:
    if not paddle.is_compiled_with_cuda():
        raise RuntimeError(
            "当前 PaddlePaddle 不支持 CUDA,请安装 cu129 版 paddlepaddle-gpu。"
        )
    gpu_count = paddle.device.cuda.device_count()
    if gpu_count == 0:
        raise RuntimeError("PaddlePaddle 未检测到 NVIDIA GPU,请检查驱动和容器 GPU 映射。")
    if gpu_id < 0 or gpu_id >= gpu_count:
        raise ValueError(f"GPU 编号 {gpu_id} 无效,当前检测到 {gpu_count} 张 GPU。")
    device = f"gpu:{gpu_id}"
    paddle.device.set_device(device)
    capability = paddle.device.cuda.get_device_capability(gpu_id)
    name = paddle.device.cuda.get_device_name(gpu_id)
    print(
        f"使用设备: {name}, compute capability {capability[0]}.{capability[1]}",
        file=sys.stderr,
    )
    return device
def to_json_data(result: Any) -> dict[str, Any]:
    data = result.json
    data = data() if callable(data) else data
    return json.loads(data) if isinstance(data, str) else data
import time
def main() -> int:
    args = parse_args()
    if not args.input.is_file():
        raise FileNotFoundError(f"输入文件不存在: {args.input}")
    device = check_gpu(args.gpu)  
    ocr = PaddleOCR(
            device=device,
            text_detection_model_name="PP-OCRv6_medium_det",
            text_recognition_model_name="PP-OCRv6_medium_rec",
            use_doc_orientation_classify=False,
            use_doc_unwarping=False,
            use_textline_orientation=False,
        )
    n = 0
    # while n < 20:
    if True:
        n+=1
        start = time.time()
        pages = [to_json_data(result) for result in ocr.predict(str(args.input))]
        payload = {
            "input": str(args.input.resolve()),
            "device": device,
            "pages": pages,
        }
        output = json.dumps(payload, ensure_ascii=False, indent=2)
        if args.output:
            args.output.parent.mkdir(parents=True, exist_ok=True)
            args.output.write_text(output + "\n", encoding="utf-8")
            print(f"识别结果已写入: {args.output}", file=sys.stderr)
        else:
            print(output)
        end = time.time()
        print("处理时间",end - start, start, end)
    return 0
if __name__ == "__main__":
    try:
        raise SystemExit(main())
    except (FileNotFoundError, RuntimeError, ValueError) as exc:
        print(f"错误: {exc}", file=sys.stderr)
        raise SystemExit(1)

声明:一代明君的小屋|版权所有,违者必究|如未注明,均为原创|本网站采用BY-NC-SA协议进行授权

转载:转载请注明原文链接 - 5090关于使用paddleocr和gpu进行ocr的demo


欢迎来到我的小屋