5
0
mirror of https://github.com/sml2h3/ddddocr.git synced 2025-05-03 21:39:42 +08:00

Merge pull request #33 from lededev/mtp

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

View File

@ -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:
image = Image.open(io.BytesIO(img_bytes))
if not isinstance(img, (bytes, str, pathlib.PurePath, Image.Image)):
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:
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 = np.array(image).astype(np.float32)
image = np.expand_dims(image, axis=0) / 255.