Merge pull request #6 from linupychiang/base64

[feat] 增加通过图片base64编码值识别的方式

ddddocr/README.md

@@ -12,17 +12,20 @@
 
 `pip install ddddocr`
 
-```
+```python
 import ddddocr
-
 ocr = ddddocr.DdddOcr()
-
 with open('test.png', 'rb') as f:
-
     img_bytes = f.read()
-
-res = ocr.classification(img_bytes)
-
+res = ocr.classification(img_bytes=img_bytes)
+print(res)
+```
+或者传入图片 base64 编码值（不包含图片头）
+```python
+import ddddocr
+ocr = ddddocr.DdddOcr()
+img_base64 = 'img_base64' # 示例
+res = ocr.classification(img_base64=img_base64)
 print(res)
 ```
 
@@ -39,4 +42,7 @@ print(res)
 
 |  参数名   | 默认值  | 说明  |
 |  ----  | ----  | ----  |
-| img  | 0 | bytes 图片的bytes格式 |
+| img_bytes  | None | bytes 图片的bytes格式 |
+| img_base64  | None | 图片的 base64 编码值（不包含图片头） |
+
+> 说明，当 `img_bytes` 和 `img_base64` 都存在时，优先使用 `img_bytes`

ddddocr/__init__.py

@@ -4,11 +4,23 @@ import warnings
 warnings.filterwarnings('ignore')
 import io
 import os
+import base64
 import onnxruntime
 from PIL import Image
 import numpy as np
 
 
+def base64_to_image(img_base64):
+    img_data = base64.b64decode(img_base64)
+    return Image.open(io.BytesIO(img_data))
+
+
+def get_img_base64(single_image_path):
+    with open(single_image_path, 'rb') as fp:
+        img_base64 = base64.b64encode(fp.read())
+        return img_base64.decode()
+
+
 class DdddOcr(object):
     def __init__(self, use_gpu: bool = False, device_id: int = 0):
         self.__graph_path = os.path.join(os.path.dirname(__file__), 'common.onnx')
@@ -461,8 +473,11 @@ class DdddOcr(object):
                         "麋", "号", "槽", "姹", "陉", "瑯", "尉", "h", "绖", "宿", "戋", "粝", "砂", "该", "鞧", "翯", "釘", "铢", "窨",
                         "設", "⒆"]
 
-    def classification(self, img: bytes):
-        image = Image.open(io.BytesIO(img))
+    def classification(self, img_bytes: bytes = None, img_base64: str = None):
+        if img_bytes:
+            image = Image.open(io.BytesIO(img_bytes))
+        else:
+            image = base64_to_image(img_base64)
         image = image.resize((int(image.size[0] * (64 / image.size[1])), 64), Image.ANTIALIAS).convert('L')
         image = np.array(image).astype(np.float32)
         image = np.expand_dims(image, axis=0) / 255.