github/ddddocr
5
0
Fork 0
mirror of https://github.com/sml2h3/ddddocr.git synced 2025-05-04 16:50:25 +08:00
Code Issues Projects Releases Wiki Activity

Merge pull request #33 from lededev/mtp

Browse Source 
classification() add image type Image.Image and pathlib.Path
This commit is contained in:
Sml2h3 2022-02-08 15:03:13 +08:00 committed by GitHub
commit 97f23ada2f
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 12 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

16
ddddocr/__init__.py
View File

@ -5,6 +5,7 @@ warnings.filterwarnings('ignore')
import io import io
import os import os
import base64 import base64
import pathlib
import onnxruntime import onnxruntime
from PIL import Image, ImageChops from PIL import Image, ImageChops
import numpy as np import numpy as np
@ -1580,13 +1581,20 @@ class DdddOcr(object):
return [] return []
return result return result
def classification(self, img_bytes: bytes = None, img_base64: str = None): def classification(self, img):
if self.det: if self.det:
raise TypeError("当前识别类型为目标检测") raise TypeError("当前识别类型为目标检测")
if img_bytes: if not isinstance(img, (bytes, str, pathlib.PurePath, Image.Image)):
image = Image.open(io.BytesIO(img_bytes)) raise TypeError("未知图片类型")
if isinstance(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)
else: else:
image = base64_to_image(img_base64) 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 = image.resize((int(image.size[0] * (64 / image.size[1])), 64), Image.ANTIALIAS).convert('L')
image = np.array(image).astype(np.float32) image = np.array(image).astype(np.float32)
image = np.expand_dims(image, axis=0) / 255. image = np.expand_dims(image, axis=0) / 255.