From f3c95cc2065602ac420a509a510ef66140a109b7 Mon Sep 17 00:00:00 2001 From: lededev Date: Sun, 6 Feb 2022 19:11:47 +0800 Subject: [PATCH 1/2] classification() add image type Image.Image and pathlib.Path --- ddddocr/__init__.py | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/ddddocr/__init__.py b/ddddocr/__init__.py index 5c7a65c..e6295d2 100644 --- a/ddddocr/__init__.py +++ b/ddddocr/__init__.py @@ -5,6 +5,7 @@ warnings.filterwarnings('ignore') import io import os import base64 +import pathlib import onnxruntime from PIL import Image, ImageChops import numpy as np @@ -1580,13 +1581,20 @@ class DdddOcr(object): return [] return result - def classification(self, img_bytes: bytes = None, img_base64: str = None): + def classification(self, img): if self.det: raise TypeError("当前识别类型为目标检测") - if img_bytes: + if not isinstance(img, (bytes, str, pathlib.PurePath, Image.Image)): + raise TypeError("未知图片类型") + if isinstance(img, bytes): image = Image.open(io.BytesIO(img_bytes)) - else: + elif isinstance(img, Image.Image): + image = img.copy() + elif isinstance(img, str): image = base64_to_image(img_base64) + else: + assert isinstance(img, pathlib.PurePath) + image = Image.open(img) 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. From 90cf72b06c0b6bc2b16827118963ca84ce377f6f Mon Sep 17 00:00:00 2001 From: lededev Date: Sun, 6 Feb 2022 19:36:07 +0800 Subject: [PATCH 2/2] fix param name --- ddddocr/__init__.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/ddddocr/__init__.py b/ddddocr/__init__.py index e6295d2..41aa654 100644 --- a/ddddocr/__init__.py +++ b/ddddocr/__init__.py @@ -1587,11 +1587,11 @@ class DdddOcr(object): if not isinstance(img, (bytes, str, pathlib.PurePath, Image.Image)): raise TypeError("未知图片类型") if isinstance(img, bytes): - image = Image.open(io.BytesIO(img_bytes)) + image = Image.open(io.BytesIO(img)) elif isinstance(img, Image.Image): image = img.copy() elif isinstance(img, str): - image = base64_to_image(img_base64) + image = base64_to_image(img) else: assert isinstance(img, pathlib.PurePath) image = Image.open(img)