topdu's picture
openocr demo
29f689c
import io
import cv2
import numpy as np
from PIL import Image
from .db_resize_for_test import DetResizeForTest
class NormalizeImage(object):
"""normalize image such as substract mean, divide std"""
def __init__(self, scale=None, mean=None, std=None, order='chw', **kwargs):
if isinstance(scale, str):
scale = eval(scale)
self.scale = np.float32(scale if scale is not None else 1.0 / 255.0)
mean = mean if mean is not None else [0.485, 0.456, 0.406]
std = std if std is not None else [0.229, 0.224, 0.225]
shape = (3, 1, 1) if order == 'chw' else (1, 1, 3)
self.mean = np.array(mean).reshape(shape).astype('float32')
self.std = np.array(std).reshape(shape).astype('float32')
def __call__(self, data):
img = data['image']
from PIL import Image
if isinstance(img, Image.Image):
img = np.array(img)
assert isinstance(img,
np.ndarray), "invalid input 'img' in NormalizeImage"
data['image'] = (img.astype('float32') * self.scale -
self.mean) / self.std
return data
class ToCHWImage(object):
"""convert hwc image to chw image"""
def __init__(self, **kwargs):
pass
def __call__(self, data):
img = data['image']
from PIL import Image
if isinstance(img, Image.Image):
img = np.array(img)
data['image'] = img.transpose((2, 0, 1))
return data
class KeepKeys(object):
def __init__(self, keep_keys, **kwargs):
self.keep_keys = keep_keys
def __call__(self, data):
data_list = []
for key in self.keep_keys:
data_list.append(data[key])
return data_list
def transform(data, ops=None):
"""transform."""
if ops is None:
ops = []
for op in ops:
data = op(data)
if data is None:
return None
return data
class DecodeImage(object):
"""decode image."""
def __init__(self,
img_mode='RGB',
channel_first=False,
ignore_orientation=False,
**kwargs):
self.img_mode = img_mode
self.channel_first = channel_first
self.ignore_orientation = ignore_orientation
def __call__(self, data):
img = data['image']
assert type(img) is bytes and len(
img) > 0, "invalid input 'img' in DecodeImage"
img = np.frombuffer(img, dtype='uint8')
if self.ignore_orientation:
img = cv2.imdecode(
img, cv2.IMREAD_IGNORE_ORIENTATION | cv2.IMREAD_COLOR)
else:
img = cv2.imdecode(img, 1)
if img is None:
return None
if self.img_mode == 'GRAY':
img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
elif self.img_mode == 'RGB':
assert img.shape[2] == 3, 'invalid shape of image[%s]' % (
img.shape)
img = img[:, :, ::-1]
if self.channel_first:
img = img.transpose((2, 0, 1))
data['image'] = img
return data
class DecodeImagePIL(object):
"""decode image."""
def __init__(self, img_mode='RGB', **kwargs):
self.img_mode = img_mode
def __call__(self, data):
img = data['image']
assert type(img) is bytes and len(
img) > 0, "invalid input 'img' in DecodeImage"
img = data['image']
buf = io.BytesIO(img)
img = Image.open(buf).convert('RGB')
if self.img_mode == 'Gray':
img = img.convert('L')
elif self.img_mode == 'BGR':
img = np.array(img)[:, :, ::-1] # 将图片转为numpy格式,并将最后一维通道倒序
img = Image.fromarray(np.uint8(img))
data['image'] = img
return data
def create_operators(op_param_list, global_config=None):
"""create operators based on the config.
Args:
params(list): a dict list, used to create some operators
"""
assert isinstance(op_param_list, list), 'operator config should be a list'
ops = []
for operator in op_param_list:
assert isinstance(operator,
dict) and len(operator) == 1, 'yaml format error'
op_name = list(operator)[0]
param = {} if operator[op_name] is None else operator[op_name]
if global_config is not None:
param.update(global_config)
op = eval(op_name)(**param)
ops.append(op)
return ops