diff --git a/.github/workflows/gen_whl_to_pypi_rapidocr_ort.yml b/.github/workflows/gen_whl_to_pypi_rapidocr_ort.yml index 821a8b4..b24dcd8 100644 --- a/.github/workflows/gen_whl_to_pypi_rapidocr_ort.yml +++ b/.github/workflows/gen_whl_to_pypi_rapidocr_ort.yml @@ -11,6 +11,7 @@ on: env: RESOURCES_URL: https://github.com/RapidAI/RapidOCR/releases/download/v1.1.0/required_for_whl_v1.3.0.zip + TWO_DIM_IMAGE: https://github.com/RapidAI/RapidOCR/releases/download/v1.1.0/two_dim_image.npy jobs: UnitTesting: @@ -39,6 +40,8 @@ jobs: cd python pip install -r requirements.txt cd tests + wget $TWO_DIM_IMAGE -O two_dim_image.npy + mv two_dim_image.npy test_files/ pytest test_det.py test_cls.py test_rec.py test_all_ort.py GenerateWHL_PushPyPi: diff --git a/.github/workflows/gen_whl_to_pypi_rapidocr_vino.yml b/.github/workflows/gen_whl_to_pypi_rapidocr_vino.yml index f6b5b16..7bdedde 100644 --- a/.github/workflows/gen_whl_to_pypi_rapidocr_vino.yml +++ b/.github/workflows/gen_whl_to_pypi_rapidocr_vino.yml @@ -11,6 +11,7 @@ on: env: RESOURCES_URL: https://github.com/RapidAI/RapidOCR/releases/download/v1.1.0/required_for_whl_v1.3.0.zip + TWO_DIM_IMAGE: https://github.com/RapidAI/RapidOCR/releases/download/v1.1.0/two_dim_image.npy jobs: UnitTesting: @@ -39,6 +40,8 @@ jobs: cd python pip install -r requirements.txt cd tests + wget $TWO_DIM_IMAGE -O two_dim_image.npy + mv two_dim_image.npy test_files/ pytest test_det.py test_cls.py test_rec.py test_all_vino.py GenerateWHL_PushPyPi: diff --git a/.gitignore b/.gitignore index edaddfb..4c21ff5 100755 --- a/.gitignore +++ b/.gitignore @@ -156,3 +156,4 @@ long1.jpg *.pdmodel .DS_Store +*.npy \ No newline at end of file diff --git a/python/rapidocr_onnxruntime/utils.py b/python/rapidocr_onnxruntime/utils.py index f5f21f5..8c1a1c0 100644 --- a/python/rapidocr_onnxruntime/utils.py +++ b/python/rapidocr_onnxruntime/utils.py @@ -124,8 +124,12 @@ class LoadImage: if img.ndim == 2: return cv2.cvtColor(img, cv2.COLOR_GRAY2BGR) - if img.ndim == 3 and img.shape[2] == 4: - return self.cvt_four_to_three(img) + if img.ndim == 3: + if img.shape[2] == 4: + return self.cvt_four_to_three(img) + + if img.shape[2] == 2: + return self.cvt_two_to_three(img) return img @@ -151,7 +155,7 @@ class LoadImage: @staticmethod def cvt_four_to_three(img: np.ndarray) -> np.ndarray: - """RGBA → RGB""" + """RGBA → BGR""" r, g, b, a = cv2.split(img) new_img = cv2.merge((b, g, r)) @@ -162,6 +166,20 @@ class LoadImage: new_img = cv2.add(new_img, not_a) return new_img + @staticmethod + def cvt_two_to_three(img: np.ndarray) -> np.ndarray: + """gray + alpha → BGR""" + img_gray = img[..., 0] + img_bgr = cv2.cvtColor(img_gray, cv2.COLOR_GRAY2BGR) + + img_alpha = img[..., 1] + not_a = cv2.bitwise_not(img_alpha) + not_a = cv2.cvtColor(not_a, cv2.COLOR_GRAY2BGR) + + new_img = cv2.bitwise_and(img_bgr, img_bgr, mask=img_alpha) + new_img = cv2.add(new_img, not_a) + return new_img + @staticmethod def verify_exist(file_path: Union[str, Path]): if not Path(file_path).exists(): @@ -200,7 +218,7 @@ def init_args(): global_group.add_argument("--width_height_ratio", type=int, default=8) det_group = parser.add_argument_group(title="Det") - det_group.add_argument("--det_use_cuda", action='store_true', default=False) + det_group.add_argument("--det_use_cuda", action="store_true", default=False) det_group.add_argument("--det_model_path", type=str, default=None) det_group.add_argument("--det_limit_side_len", type=float, default=736) det_group.add_argument( @@ -215,7 +233,7 @@ def init_args(): ) cls_group = parser.add_argument_group(title="Cls") - cls_group.add_argument("--cls_use_cuda", action='store_true', default=False) + cls_group.add_argument("--cls_use_cuda", action="store_true", default=False) cls_group.add_argument("--cls_model_path", type=str, default=None) cls_group.add_argument("--cls_image_shape", type=list, default=[3, 48, 192]) cls_group.add_argument("--cls_label_list", type=list, default=["0", "180"]) @@ -223,7 +241,7 @@ def init_args(): cls_group.add_argument("--cls_thresh", type=float, default=0.9) rec_group = parser.add_argument_group(title="Rec") - rec_group.add_argument("--rec_use_cuda", action='store_true', default=False) + rec_group.add_argument("--rec_use_cuda", action="store_true", default=False) rec_group.add_argument("--rec_model_path", type=str, default=None) rec_group.add_argument("--rec_img_shape", type=list, default=[3, 48, 320]) rec_group.add_argument("--rec_batch_num", type=int, default=6) @@ -269,7 +287,7 @@ class UpdateParameters: return config det_dict = {k.split("det_")[1]: v for k, v in det_dict.items()} - model_path = det_dict.get('model_path', None) + model_path = det_dict.get("model_path", None) if not model_path: det_dict["model_path"] = str(root_dir / config["model_path"]) @@ -281,9 +299,9 @@ class UpdateParameters: return config need_remove_prefix = ["cls_label_list", "cls_model_path", "cls_use_cuda"] - new_cls_dict = self.remove_prefix(cls_dict, 'cls_', need_remove_prefix) + new_cls_dict = self.remove_prefix(cls_dict, "cls_", need_remove_prefix) - model_path = new_cls_dict.get('model_path', None) + model_path = new_cls_dict.get("model_path", None) if model_path: new_cls_dict["model_path"] = str(root_dir / config["model_path"]) @@ -295,9 +313,9 @@ class UpdateParameters: return config need_remove_prefix = ["rec_model_path", "rec_use_cuda"] - new_rec_dict = self.remove_prefix(rec_dict, 'rec_', need_remove_prefix) + new_rec_dict = self.remove_prefix(rec_dict, "rec_", need_remove_prefix) - model_path = new_rec_dict.get('model_path', None) + model_path = new_rec_dict.get("model_path", None) if not model_path: new_rec_dict["model_path"] = str(root_dir / config["model_path"]) diff --git a/python/rapidocr_openvino/utils.py b/python/rapidocr_openvino/utils.py index 3ee117f..c8bab09 100644 --- a/python/rapidocr_openvino/utils.py +++ b/python/rapidocr_openvino/utils.py @@ -56,8 +56,12 @@ class LoadImage: if img.ndim == 2: return cv2.cvtColor(img, cv2.COLOR_GRAY2BGR) - if img.ndim == 3 and img.shape[2] == 4: - return self.cvt_four_to_three(img) + if img.ndim == 3: + if img.shape[2] == 4: + return self.cvt_four_to_three(img) + + if img.shape[2] == 2: + return self.cvt_two_to_three(img) return img @@ -83,7 +87,7 @@ class LoadImage: @staticmethod def cvt_four_to_three(img: np.ndarray) -> np.ndarray: - """RGBA → RGB""" + """RGBA → BGR""" r, g, b, a = cv2.split(img) new_img = cv2.merge((b, g, r)) @@ -94,6 +98,20 @@ class LoadImage: new_img = cv2.add(new_img, not_a) return new_img + @staticmethod + def cvt_two_to_three(img: np.ndarray) -> np.ndarray: + """gray + alpha → BGR""" + img_gray = img[..., 0] + img_bgr = cv2.cvtColor(img_gray, cv2.COLOR_GRAY2BGR) + + img_alpha = img[..., 1] + not_a = cv2.bitwise_not(img_alpha) + not_a = cv2.cvtColor(not_a, cv2.COLOR_GRAY2BGR) + + new_img = cv2.bitwise_and(img_bgr, img_bgr, mask=img_alpha) + new_img = cv2.add(new_img, not_a) + return new_img + @staticmethod def verify_exist(file_path: Union[str, Path]): if not Path(file_path).exists(): @@ -198,7 +216,7 @@ class UpdateParameters: return config det_dict = {k.split("det_")[1]: v for k, v in det_dict.items()} - model_path = det_dict.get('model_path', None) + model_path = det_dict.get("model_path", None) if not model_path: det_dict["model_path"] = str(root_dir / config["model_path"]) @@ -210,9 +228,9 @@ class UpdateParameters: return config need_remove_prefix = ["cls_label_list", "cls_model_path", "cls_use_cuda"] - new_cls_dict = self.remove_prefix(cls_dict, 'cls_', need_remove_prefix) + new_cls_dict = self.remove_prefix(cls_dict, "cls_", need_remove_prefix) - model_path = new_cls_dict.get('model_path', None) + model_path = new_cls_dict.get("model_path", None) if model_path: new_cls_dict["model_path"] = str(root_dir / config["model_path"]) @@ -224,9 +242,9 @@ class UpdateParameters: return config need_remove_prefix = ["rec_model_path", "rec_use_cuda"] - new_rec_dict = self.remove_prefix(rec_dict, 'rec_', need_remove_prefix) + new_rec_dict = self.remove_prefix(rec_dict, "rec_", need_remove_prefix) - model_path = new_rec_dict.get('model_path', None) + model_path = new_rec_dict.get("model_path", None) if not model_path: new_rec_dict["model_path"] = str(root_dir / config["model_path"]) diff --git a/python/tests/test_all_ort.py b/python/tests/test_all_ort.py index 99f818d..80722af 100644 --- a/python/tests/test_all_ort.py +++ b/python/tests/test_all_ort.py @@ -95,3 +95,11 @@ def test_input_rec_parameters(): result, _ = rapid_ocr(image_path) raise FileNotFoundError() assert exc_info.type is FileNotFoundError + + +def test_input_two_ndim(): + img_npy = tests_dir / "two_dim_image.npy" + image_array = np.load(str(img_npy)) + result, _ = rapid_ocr(image_array) + + assert len(result) == 87 diff --git a/python/tests/test_all_vino.py b/python/tests/test_all_vino.py index 1ef92b8..5175bfb 100644 --- a/python/tests/test_all_vino.py +++ b/python/tests/test_all_vino.py @@ -95,3 +95,11 @@ def test_input_rec_parameters(): result, _ = rapid_ocr(image_path) raise FileNotFoundError() assert exc_info.type is FileNotFoundError + + +def test_input_two_ndim(): + img_npy = tests_dir / "two_dim_image.npy" + image_array = np.load(str(img_npy)) + result, _ = rapid_ocr(image_array) + + assert len(result) == 87 diff --git a/python/tests/test_files/two_dim_image.npy b/python/tests/test_files/two_dim_image.npy new file mode 100644 index 0000000..c373986 Binary files /dev/null and b/python/tests/test_files/two_dim_image.npy differ