mirror of
https://github.com/RapidAI/RapidOCR.git
synced 2026-11-14 06:50:48 +08:00
test(rapidocr_onnxruntime): Add two unit testings about cuda and dml
This commit is contained in:
1 parent
35e85ef370
commit
02829ef986
1 file changed
+17
@@ -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)
|
||||
|
||||
Reference in new issue
Block a user