Fixed issue #115

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"])
 

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"])
 

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

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