101 lines
3.6 KiB
Python
101 lines
3.6 KiB
Python
# test_ocr_runtime.py
|
|
# 사용 예:
|
|
# python test_ocr_runtime.py --imgs img --onnx_dir PP_Models/onnx_opt --dict ppocr_keys_v1.txt --ep auto --use_cls
|
|
|
|
import os
|
|
import glob
|
|
import argparse
|
|
import time
|
|
from ocr_runtime_module import ONNXOCR, PaddleRunner
|
|
|
|
def parse_args():
|
|
ap = argparse.ArgumentParser()
|
|
ap.add_argument("--imgs", type=str, required=True, help="image folder or single image")
|
|
ap.add_argument("--onnx_dir", type=str, default=None, help="Folder containing det/rec/cls onnx")
|
|
ap.add_argument("--det", type=str, default=None, help="Path to det onnx")
|
|
ap.add_argument("--rec", type=str, default=None, help="Path to rec onnx")
|
|
ap.add_argument("--cls", type=str, default=None, help="Path to cls onnx (optional)")
|
|
ap.add_argument("--dict", type=str, required=True, help="ppocr_keys_v1.txt")
|
|
ap.add_argument("--use_cls", action="store_true", help="run angle classifier")
|
|
ap.add_argument("--ep", type=str, default="auto", choices=["auto","cpu","cuda","trt"], help="Execution Provider preference")
|
|
ap.add_argument("--trt_fp16", action="store_true", help="Enable TensorRT FP16 (if TRT is used)")
|
|
ap.add_argument("--trt_ws", type=int, default=(2<<30), help="TRT workspace size bytes")
|
|
ap.add_argument("--rec_bs", type=int, default=6, help="REC batch size")
|
|
ap.add_argument("--bench_paddle", action="store_true", help="Also run PaddleOCR for comparison")
|
|
return ap.parse_args()
|
|
|
|
def list_images(p):
|
|
if os.path.isdir(p):
|
|
imgs = []
|
|
for ext in ["*.jpg","*.jpeg","*.png","*.bmp","*.webp","*.tif","*.tiff"]:
|
|
imgs.extend(glob.glob(os.path.join(p, ext)))
|
|
return sorted(imgs)
|
|
else:
|
|
return [p]
|
|
|
|
def main():
|
|
args = parse_args()
|
|
|
|
onnx_ocr = ONNXOCR(
|
|
onnx_dir=args.onnx_dir,
|
|
det_path=args.det,
|
|
rec_path=args.rec,
|
|
cls_path=args.cls,
|
|
dict_path=args.dict,
|
|
ep=args.ep,
|
|
trt_fp16=args.trt_fp16,
|
|
trt_workspace=args.trt_ws,
|
|
rec_bs=args.rec_bs,
|
|
use_warp_crop=True
|
|
)
|
|
|
|
print(f"[EP avail] {onnx_ocr.avail}")
|
|
print(f"[EP chosen] {onnx_ocr.providers_det}")
|
|
|
|
imgs = list_images(args.imgs)
|
|
|
|
# 옵션: Paddle 비교 세팅
|
|
paddle_runner = None
|
|
if args.bench_paddle:
|
|
try:
|
|
paddle_runner = PaddleRunner(use_angle_cls=args.use_cls, lang='ch')
|
|
except Exception as e:
|
|
print("[WARN] PaddleOCR init failed:", e)
|
|
|
|
for path in imgs:
|
|
print(f"\n[TEST] image: {path}")
|
|
import cv2
|
|
img = cv2.imdecode(np.fromfile(path, dtype=np.uint8), cv2.IMREAD_COLOR)
|
|
if img is None:
|
|
print(" ! skip (cannot read)")
|
|
continue
|
|
|
|
# ONNX warmup (간단 1회)
|
|
try:
|
|
onnx_ocr.run(img, use_cls=args.use_cls, rec_bs=args.rec_bs)
|
|
except Exception:
|
|
pass
|
|
|
|
t0 = time.perf_counter()
|
|
boxes, recs, T = onnx_ocr.run(img, use_cls=args.use_cls, rec_bs=args.rec_bs)
|
|
t1 = time.perf_counter()
|
|
|
|
print("\n[ONNX] results:")
|
|
for (txt, conf) in recs[:10]:
|
|
print(f" - {txt} ({conf:.3f})")
|
|
print(f"[ONNX] time: {(t1-t0)*1000.0:.1f} ms | T={T}")
|
|
|
|
if paddle_runner:
|
|
try:
|
|
pres, p_ms = paddle_runner.run(img)
|
|
print("\n[Paddle] results:")
|
|
for (txt, conf) in pres[:10]:
|
|
print(f" - {txt} ({conf:.3f})")
|
|
print(f"[Paddle] time: {p_ms:.1f} ms")
|
|
except Exception as e:
|
|
print("[Paddle] failed:", e)
|
|
|
|
if __name__ == "__main__":
|
|
import numpy as np # needed by imdecode
|
|
main()
|