test(rapidocr_onnxruntime): Add two unit testings about cuda and dml

This commit is contained in:
SWHL committed 2024-05-21 20:52:17 +08:00
1 parent 35e85ef370
commit 02829ef986
1 file changed
+17
+17
View File
@@ -1,6 +1,7 @@
# -*- encoding: utf-8 -*-
# @Author: SWHL
# @Contact: liekkaskono@163.com
import logging
import sys
from pathlib import Path
@@ -20,6 +21,22 @@ img_path = tests_dir / "ch_en_num.jpg"
package_name = "rapidocr_onnxruntime"
def test_ort_cuda_warning(caplog):
engine = RapidOCR(det_use_cuda=True)
caplog.set_level(logging.WARNING)
assert caplog.records[0].levelname == "WARNING"
assert "CUDAExecutionProvider" in caplog.records[0].message
def test_ort_dml_warning(caplog):
engine = RapidOCR(det_use_dml=True)
caplog.set_level(logging.WARNING)
assert caplog.records[0].levelname == "WARNING"
assert "DirectML" in caplog.records[0].message
def test_mode_one_img():
img_path = tests_dir / "issue_170.png"
result, _ = engine(img_path)