Bug fix for converting error when processing bitmap image. (#171)

RapidAI · Apr 19, 2024 · f602d30 · f602d30
1 parent 794f120
commit f602d30
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 9 deletions.
diff --git a/python/rapidocr_onnxruntime/utils.py b/python/rapidocr_onnxruntime/utils.py
@@ -141,23 +141,30 @@ def load_img(self, img: InputType) -> np.ndarray:
         if isinstance(img, (str, Path)):
             self.verify_exist(img)
             try:
-                img = np.array(Image.open(img))
+                img = self.img_to_ndarray(Image.open(img))
             except UnidentifiedImageError as e:
                 raise LoadImageError(f"cannot identify image file {img}") from e
             return img
 
         if isinstance(img, bytes):
-            img = np.array(Image.open(BytesIO(img)))
+            img = self.img_to_ndarray(Image.open(BytesIO(img)))
             return img
 
         if isinstance(img, np.ndarray):
             return img
 
         if isinstance(img, Image.Image):
-            return np.array(img)
+            return self.img_to_ndarray(img)
 
         raise LoadImageError(f"{type(img)} is not supported!")
 
+    def img_to_ndarray(self, img: Image.Image) -> np.ndarray:
+        if img.mode == "1":
+            img = img.convert("L")
+            return np.array(img)
+        else:
+            return np.array(img)
+
     def convert_img(self, img: np.ndarray, origin_img_type):
         if img.ndim == 2:
             return cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)

diff --git a/python/rapidocr_openvino/utils.py b/python/rapidocr_openvino/utils.py
@@ -68,23 +68,30 @@ def load_img(self, img: InputType) -> np.ndarray:
         if isinstance(img, (str, Path)):
             self.verify_exist(img)
             try:
-                img = np.array(Image.open(img))
+                img = self.img_to_ndarray(Image.open(img))
             except UnidentifiedImageError as e:
                 raise LoadImageError(f"cannot identify image file {img}") from e
             return img
 
         if isinstance(img, bytes):
-            img = np.array(Image.open(BytesIO(img)))
+            img = self.img_to_ndarray(Image.open(BytesIO(img)))
             return img
 
         if isinstance(img, np.ndarray):
             return img
 
         if isinstance(img, Image.Image):
-            return np.array(img)
+            return self.img_to_ndarray(img)
 
         raise LoadImageError(f"{type(img)} is not supported!")
 
+    def img_to_ndarray(self, img: Image.Image) -> np.ndarray:
+        if img.mode == "1":
+            img = img.convert("L")
+            return np.array(img)
+        else:
+            return np.array(img)
+
     def convert_img(self, img: np.ndarray, origin_img_type):
         if img.ndim == 2:
             return cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)

diff --git a/python/rapidocr_paddle/utils.py b/python/rapidocr_paddle/utils.py
@@ -150,23 +150,30 @@ def load_img(self, img: InputType) -> np.ndarray:
         if isinstance(img, (str, Path)):
             self.verify_exist(img)
             try:
-                img = np.array(Image.open(img))
+                img = self.img_to_ndarray(Image.open(img))
             except UnidentifiedImageError as e:
                 raise LoadImageError(f"cannot identify image file {img}") from e
             return img
 
         if isinstance(img, bytes):
-            img = np.array(Image.open(BytesIO(img)))
+            img = self.img_to_ndarray(Image.open(BytesIO(img)))
             return img
 
         if isinstance(img, np.ndarray):
             return img
 
         if isinstance(img, Image.Image):
-            return np.array(img)
+            return self.img_to_ndarray(img)
 
         raise LoadImageError(f"{type(img)} is not supported!")
 
+    def img_to_ndarray(self, img: Image.Image) -> np.ndarray:
+        if img.mode == "1":
+            img = img.convert("L")
+            return np.array(img)
+        else:
+            return np.array(img)
+
     def convert_img(self, img: np.ndarray, origin_img_type):
         if img.ndim == 2:
             return cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)