mtfyun/rapidocr
1
0 Fork 0
mirror of https://github.com/RapidAI/RapidOCR.git synced 2026-11-11 05:20:49 +08:00
Code Issues Packages Projects Releases Wiki Activity

Format by black

Browse Source
This commit is contained in:
SWHL committed 2023-06-24 10:27:12 +08:00
1 parent ad4631fdd4
commit f4fd5cdcba
40 files changed
+908 -854
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

python/tests/base_module.py
+6 -6
View File
@@ -8,25 +8,25 @@ from pathlib import Path
import yaml
class BaseModule():
def __init__(self, package_name: str = 'rapidocr_onnxruntime'):
class BaseModule:
def __init__(self, package_name: str = "rapidocr_onnxruntime"):
self.package_name = package_name
self.root_dir = Path(__file__).resolve().parent.parent
self.package_dir = self.root_dir / self.package_name
self.tests_dir = self.root_dir / 'tests'
self.tests_dir = self.root_dir / "tests"
sys.path.append(str(self.root_dir))
sys.path.append(str(self.package_dir))
def init_module(self, module_name: str, class_name: str = None):
if class_name is None:
module_part = importlib.import_module(f'{self.package_name}')
module_part = importlib.import_module(f"{self.package_name}")
return module_part
module_part = importlib.import_module(f'{self.package_name}.{module_name}')
module_part = importlib.import_module(f"{self.package_name}.{module_name}")
return getattr(module_part, class_name)
@staticmethod
def read_yaml(yaml_path: str):
with open(yaml_path, 'rb') as f:
with open(yaml_path, "rb") as f:
data = yaml.load(f, Loader=yaml.Loader)
return data
python/tests/benchmark/benchmark.py
+10 -8
View File
@@ -14,22 +14,22 @@ import cv2
from rapidocr_onnxruntime import RapidOCR
if __name__ == '__main__':
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument('--yaml_path', type=str, default='config.yaml')
parser.add_argument("--yaml_path", type=str, default="config.yaml")
args = parser.parse_args()
yaml_path = cur_dir / args.yaml_path
rapid_ocr = RapidOCR(yaml_path)
image_dir = cur_dir / 'test_images_benchmark'
image_dir = cur_dir / "test_images_benchmark"
if not image_dir.exists():
raise FileNotFoundError(f'{image_dir} does not exits!!')
raise FileNotFoundError(f"{image_dir} does not exits!!")
image_list = list(image_dir.iterdir())
cost_time_list = []
for image_path in tqdm(image_list, desc='Test'):
for image_path in tqdm(image_list, desc="Test"):
img = cv2.imread(str(image_path))
start_time = time.time()
@@ -40,6 +40,8 @@ if __name__ == '__main__':
total_time = sum(cost_time_list)
avg_time = total_time / len(cost_time_list)
print(f'Total Files: {len(image_list)}, '
f'Total Time: {total_time:.5f}, '
f'Average Time: {avg_time:.5f}')
print(
f"Total Files: {len(image_list)}, "
f"Total Time: {total_time:.5f}, "
f"Average Time: {avg_time:.5f}"
)
python/tests/test_all_ort.py
+17 -17
View File
@@ -16,14 +16,14 @@ sys.path.append(str(root_dir))
from rapidocr_onnxruntime import RapidOCR, LoadImageError
rapid_ocr = RapidOCR()
tests_dir = root_dir / 'tests' / 'test_files'
tests_dir = root_dir / "tests" / "test_files"
def test_normal():
image_path = tests_dir / 'ch_en_num.jpg'
image_path = tests_dir / "ch_en_num.jpg"
img = cv2.imread(str(image_path))
result, _ = rapid_ocr(img)
assert result[0][1] == '正品促销'
assert result[0][1] == "正品促销"
assert len(result) == 17
@@ -42,29 +42,29 @@ def test_zeros():
def test_input_str():
image_path = tests_dir / 'ch_en_num.jpg'
image_path = tests_dir / "ch_en_num.jpg"
result, _ = rapid_ocr(str(image_path))
assert result[0][1] == '正品促销'
assert result[0][1] == "正品促销"
assert len(result) == 17
def test_input_bytes():
image_path = tests_dir / 'ch_en_num.jpg'
with open(image_path, 'rb') as f:
image_path = tests_dir / "ch_en_num.jpg"
with open(image_path, "rb") as f:
result, _ = rapid_ocr(f.read())
assert result[0][1] == '正品促销'
assert result[0][1] == "正品促销"
assert len(result) == 17
def test_input_path():
image_path = tests_dir / 'ch_en_num.jpg'
image_path = tests_dir / "ch_en_num.jpg"
result, _ = rapid_ocr(image_path)
assert result[0][1] == '正品促销'
assert result[0][1] == "正品促销"
assert len(result) == 17
def test_input_parameters():
image_path = tests_dir / 'ch_en_num.jpg'
image_path = tests_dir / "ch_en_num.jpg"
rapid_ocr = RapidOCR(text_score=1)
result, _ = rapid_ocr(image_path)
@@ -72,27 +72,27 @@ def test_input_parameters():
def test_input_det_parameters():
image_path = tests_dir / 'ch_en_num.jpg'
image_path = tests_dir / "ch_en_num.jpg"
with pytest.raises(FileNotFoundError) as exc_info:
rapid_ocr = RapidOCR(det_model_path='1.onnx')
rapid_ocr = RapidOCR(det_model_path="1.onnx")
result, _ = rapid_ocr(image_path)
raise FileNotFoundError()
assert exc_info.type is FileNotFoundError
def test_input_cls_parameters():
image_path = tests_dir / 'ch_en_num.jpg'
image_path = tests_dir / "ch_en_num.jpg"
with pytest.raises(FileNotFoundError) as exc_info:
rapid_ocr = RapidOCR(cls_model_path='1.onnx')
rapid_ocr = RapidOCR(cls_model_path="1.onnx")
result, _ = rapid_ocr(image_path)
raise FileNotFoundError()
assert exc_info.type is FileNotFoundError
def test_input_rec_parameters():
image_path = tests_dir / 'ch_en_num.jpg'
image_path = tests_dir / "ch_en_num.jpg"
with pytest.raises(FileNotFoundError) as exc_info:
rapid_ocr = RapidOCR(rec_model_path='1.onnx')
rapid_ocr = RapidOCR(rec_model_path="1.onnx")
result, _ = rapid_ocr(image_path)
raise FileNotFoundError()
assert exc_info.type is FileNotFoundError
python/tests/test_all_vino.py
+17 -17
View File
@@ -16,14 +16,14 @@ from rapidocr_openvino import RapidOCR, LoadImageError
rapid_ocr = RapidOCR()
tests_dir = root_dir / 'tests' / 'test_files'
tests_dir = root_dir / "tests" / "test_files"
def test_normal():
image_path = tests_dir / 'ch_en_num.jpg'
image_path = tests_dir / "ch_en_num.jpg"
img = cv2.imread(str(image_path))
result, _ = rapid_ocr(img)
assert result[0][1] == '正品促销'
assert result[0][1] == "正品促销"
assert len(result) == 17
@@ -42,29 +42,29 @@ def test_zeros():
def test_input_str():
image_path = tests_dir / 'ch_en_num.jpg'
image_path = tests_dir / "ch_en_num.jpg"
result, _ = rapid_ocr(str(image_path))
assert result[0][1] == '正品促销'
assert result[0][1] == "正品促销"
assert len(result) == 17
def test_input_bytes():
image_path = tests_dir / 'ch_en_num.jpg'
with open(image_path, 'rb') as f:
image_path = tests_dir / "ch_en_num.jpg"
with open(image_path, "rb") as f:
result, _ = rapid_ocr(f.read())
assert result[0][1] == '正品促销'
assert result[0][1] == "正品促销"
assert len(result) == 17
def test_input_path():
image_path = tests_dir / 'ch_en_num.jpg'
image_path = tests_dir / "ch_en_num.jpg"
result, _ = rapid_ocr(image_path)
assert result[0][1] == '正品促销'
assert result[0][1] == "正品促销"
assert len(result) == 17
def test_input_parameters():
image_path = tests_dir / 'ch_en_num.jpg'
image_path = tests_dir / "ch_en_num.jpg"
rapid_ocr = RapidOCR(text_score=1)
result, _ = rapid_ocr(image_path)
@@ -72,27 +72,27 @@ def test_input_parameters():
def test_input_det_parameters():
image_path = tests_dir / 'ch_en_num.jpg'
image_path = tests_dir / "ch_en_num.jpg"
with pytest.raises(FileNotFoundError) as exc_info:
rapid_ocr = RapidOCR(det_model_path='1.onnx')
rapid_ocr = RapidOCR(det_model_path="1.onnx")
result, _ = rapid_ocr(image_path)
raise FileNotFoundError()
assert exc_info.type is FileNotFoundError
def test_input_cls_parameters():
image_path = tests_dir / 'ch_en_num.jpg'
image_path = tests_dir / "ch_en_num.jpg"
with pytest.raises(FileNotFoundError) as exc_info:
rapid_ocr = RapidOCR(cls_model_path='1.onnx')
rapid_ocr = RapidOCR(cls_model_path="1.onnx")
result, _ = rapid_ocr(image_path)
raise FileNotFoundError()
assert exc_info.type is FileNotFoundError
def test_input_rec_parameters():
image_path = tests_dir / 'ch_en_num.jpg'
image_path = tests_dir / "ch_en_num.jpg"
with pytest.raises(FileNotFoundError) as exc_info:
rapid_ocr = RapidOCR(rec_model_path='1.onnx')
rapid_ocr = RapidOCR(rec_model_path="1.onnx")
result, _ = rapid_ocr(image_path)
raise FileNotFoundError()
assert exc_info.type is FileNotFoundError
python/tests/test_cls.py
+7 -9
View File
@@ -7,24 +7,22 @@ from base_module import BaseModule
@pytest.mark.parametrize(
'package_name',
[('rapidocr_onnxruntime'),
('rapidocr_openvino')]
"package_name", [("rapidocr_onnxruntime"), ("rapidocr_openvino")]
)
def test_cls(package_name: str):
module_name = 'ch_ppocr_v2_cls'
class_name = 'TextClassifier'
module_name = "ch_ppocr_v2_cls"
class_name = "TextClassifier"
base = BaseModule(package_name=package_name)
TextClassifier = base.init_module(module_name, class_name)
yaml_path = base.package_dir / module_name / 'config.yaml'
yaml_path = base.package_dir / module_name / "config.yaml"
config = base.read_yaml(str(yaml_path))
config['model_path'] = str(base.package_dir / config['model_path'])
config["model_path"] = str(base.package_dir / config["model_path"])
text_cls = TextClassifier(config)
img_path = base.tests_dir / 'test_files' / 'text_cls.jpg'
img_path = base.tests_dir / "test_files" / "text_cls.jpg"
img = cv2.imread(str(img_path))
result = text_cls([img])
assert result[1][0][0] == '180'
assert result[1][0][0] == "180"
python/tests/test_det.py
+6 -9
View File
@@ -6,23 +6,20 @@ import pytest
from base_module import BaseModule
@pytest.mark.parametrize(
'package_name',
['rapidocr_onnxruntime', 'rapidocr_openvino']
)
@pytest.mark.parametrize("package_name", ["rapidocr_onnxruntime", "rapidocr_openvino"])
def test_det(package_name):
module_name = 'ch_ppocr_v3_det'
class_name = 'TextDetector'
module_name = "ch_ppocr_v3_det"
class_name = "TextDetector"
base = BaseModule(package_name)
TextDetector = base.init_module(module_name, class_name)
yaml_path = base.package_dir / module_name / 'config.yaml'
yaml_path = base.package_dir / module_name / "config.yaml"
config = base.read_yaml(str(yaml_path))
config['model_path'] = str(base.package_dir / config['model_path'])
config["model_path"] = str(base.package_dir / config["model_path"])
text_det = TextDetector(config)
img_path = base.tests_dir / 'test_files' / 'text_det.jpg'
img_path = base.tests_dir / "test_files" / "text_det.jpg"
img = cv2.imread(str(img_path))
dt_boxes, elapse = text_det(img)
assert dt_boxes.shape == (18, 4, 2)
python/tests/test_rec.py
+7 -10
View File
@@ -6,24 +6,21 @@ import pytest
from base_module import BaseModule
@pytest.mark.parametrize(
'package_name',
['rapidocr_onnxruntime', 'rapidocr_openvino']
)
@pytest.mark.parametrize("package_name", ["rapidocr_onnxruntime", "rapidocr_openvino"])
def test_det(package_name):
module_name = 'ch_ppocr_v3_rec'
class_name = 'TextRecognizer'
module_name = "ch_ppocr_v3_rec"
class_name = "TextRecognizer"
base = BaseModule(package_name)
TextRecognizer = base.init_module(module_name, class_name)
yaml_path = base.package_dir / module_name / 'config.yaml'
yaml_path = base.package_dir / module_name / "config.yaml"
config = base.read_yaml(str(yaml_path))
config['model_path'] = str(base.package_dir / config['model_path'])
config["model_path"] = str(base.package_dir / config["model_path"])
text_rec = TextRecognizer(config)
img_path = base.tests_dir / 'test_files' / 'text_rec.jpg'
img_path = base.tests_dir / "test_files" / "text_rec.jpg"
img = cv2.imread(str(img_path))
rec_res, elapse = text_rec([img])
assert rec_res[0][0] == '韩国小馆'
assert rec_res[0][0] == "韩国小馆"