mirror of
https://github.com/RapidAI/RapidOCR.git
synced 2026-11-11 21:40:52 +08:00
415 lines
13 KiB
Python
Executable File
415 lines
13 KiB
Python
Executable File
"""
|
|
# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
"""
|
|
# coding:utf-8
|
|
from __future__ import (absolute_import, division, print_function,
|
|
unicode_literals)
|
|
|
|
import os
|
|
import sys
|
|
|
|
import cv2
|
|
import numpy as np
|
|
import pyclipper
|
|
import six
|
|
from shapely.geometry import Polygon
|
|
|
|
|
|
class DecodeImage(object):
|
|
""" decode image """
|
|
|
|
def __init__(self, img_mode='RGB', channel_first=False, **kwargs):
|
|
self.img_mode = img_mode
|
|
self.channel_first = channel_first
|
|
|
|
def __call__(self, data):
|
|
img = data['image']
|
|
if six.PY2:
|
|
assert type(img) is str and len(
|
|
img) > 0, "invalid input 'img' in DecodeImage"
|
|
else:
|
|
assert type(img) is bytes and len(
|
|
img) > 0, "invalid input 'img' in DecodeImage"
|
|
img = np.frombuffer(img, dtype='uint8')
|
|
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 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
|
|
|
|
|
|
class DetResizeForTest(object):
|
|
def __init__(self, **kwargs):
|
|
super(DetResizeForTest, self).__init__()
|
|
self.resize_type = 0
|
|
if 'image_shape' in kwargs:
|
|
self.image_shape = kwargs['image_shape']
|
|
self.resize_type = 1
|
|
elif 'limit_side_len' in kwargs:
|
|
self.limit_side_len = kwargs['limit_side_len']
|
|
self.limit_type = kwargs.get('limit_type', 'min')
|
|
|
|
if 'resize_long' in kwargs:
|
|
self.resize_type = 2
|
|
self.resize_long = kwargs.get('resize_long', 960)
|
|
else:
|
|
self.limit_side_len = 736
|
|
self.limit_type = 'min'
|
|
|
|
def __call__(self, data):
|
|
img = data['image']
|
|
src_h, src_w, _ = img.shape
|
|
|
|
if self.resize_type == 0:
|
|
# img, shape = self.resize_image_type0(img)
|
|
img, [ratio_h, ratio_w] = self.resize_image_type0(img)
|
|
elif self.resize_type == 2:
|
|
img, [ratio_h, ratio_w] = self.resize_image_type2(img)
|
|
else:
|
|
# img, shape = self.resize_image_type1(img)
|
|
img, [ratio_h, ratio_w] = self.resize_image_type1(img)
|
|
data['image'] = img
|
|
data['shape'] = np.array([src_h, src_w, ratio_h, ratio_w])
|
|
return data
|
|
|
|
def resize_image_type1(self, img):
|
|
resize_h, resize_w = self.image_shape
|
|
ori_h, ori_w = img.shape[:2] # (h, w, c)
|
|
ratio_h = float(resize_h) / ori_h
|
|
ratio_w = float(resize_w) / ori_w
|
|
img = cv2.resize(img, (int(resize_w), int(resize_h)))
|
|
# return img, np.array([ori_h, ori_w])
|
|
return img, [ratio_h, ratio_w]
|
|
|
|
def resize_image_type0(self, img):
|
|
"""
|
|
resize image to a size multiple of 32 which is required by the network
|
|
args:
|
|
img(array): array with shape [h, w, c]
|
|
return(tuple):
|
|
img, (ratio_h, ratio_w)
|
|
"""
|
|
limit_side_len = self.limit_side_len
|
|
h, w, _ = img.shape
|
|
|
|
# limit the max side
|
|
if self.limit_type == 'max':
|
|
if max(h, w) > limit_side_len:
|
|
if h > w:
|
|
ratio = float(limit_side_len) / h
|
|
else:
|
|
ratio = float(limit_side_len) / w
|
|
else:
|
|
ratio = 1.
|
|
else:
|
|
if min(h, w) < limit_side_len:
|
|
if h < w:
|
|
ratio = float(limit_side_len) / h
|
|
else:
|
|
ratio = float(limit_side_len) / w
|
|
else:
|
|
ratio = 1.
|
|
resize_h = int(h * ratio)
|
|
resize_w = int(w * ratio)
|
|
|
|
resize_h = int(round(resize_h / 32) * 32)
|
|
resize_w = int(round(resize_w / 32) * 32)
|
|
|
|
try:
|
|
if int(resize_w) <= 0 or int(resize_h) <= 0:
|
|
return None, (None, None)
|
|
img = cv2.resize(img, (int(resize_w), int(resize_h)))
|
|
except:
|
|
print(img.shape, resize_w, resize_h)
|
|
sys.exit(0)
|
|
ratio_h = resize_h / float(h)
|
|
ratio_w = resize_w / float(w)
|
|
return img, [ratio_h, ratio_w]
|
|
|
|
def resize_image_type2(self, img):
|
|
h, w, _ = img.shape
|
|
|
|
resize_w = w
|
|
resize_h = h
|
|
|
|
# Fix the longer side
|
|
if resize_h > resize_w:
|
|
ratio = float(self.resize_long) / resize_h
|
|
else:
|
|
ratio = float(self.resize_long) / resize_w
|
|
|
|
resize_h = int(resize_h * ratio)
|
|
resize_w = int(resize_w * ratio)
|
|
|
|
max_stride = 128
|
|
resize_h = (resize_h + max_stride - 1) // max_stride * max_stride
|
|
resize_w = (resize_w + max_stride - 1) // max_stride * max_stride
|
|
img = cv2.resize(img, (int(resize_w), int(resize_h)))
|
|
ratio_h = resize_h / float(h)
|
|
ratio_w = resize_w / float(w)
|
|
|
|
return img, [ratio_h, ratio_w]
|
|
|
|
|
|
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
|
|
|
|
|
|
def check_and_read_gif(img_path):
|
|
if os.path.basename(img_path)[-3:] in ['gif', 'GIF']:
|
|
gif = cv2.VideoCapture(img_path)
|
|
ret, frame = gif.read()
|
|
if not ret:
|
|
print("Cannot read {}. This gif image maybe corrupted.")
|
|
return None, False
|
|
if len(frame.shape) == 2 or frame.shape[-1] == 1:
|
|
frame = cv2.cvtColor(frame, cv2.COLOR_GRAY2RGB)
|
|
imgvalue = frame[:, :, ::-1]
|
|
return imgvalue, True
|
|
return None, False
|
|
|
|
|
|
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
|
|
|
|
|
|
def draw_text_det_res(dt_boxes, img_path):
|
|
src_im = cv2.imread(img_path)
|
|
for box in dt_boxes:
|
|
box = np.array(box).astype(np.int32).reshape(-1, 2)
|
|
cv2.polylines(src_im, [box], True,
|
|
color=(255, 255, 0), thickness=2)
|
|
return src_im
|
|
|
|
|
|
class DBPostProcess(object):
|
|
"""
|
|
The post process for Differentiable Binarization (DB).
|
|
"""
|
|
|
|
def __init__(self,
|
|
thresh=0.3,
|
|
box_thresh=0.7,
|
|
max_candidates=1000,
|
|
unclip_ratio=2.0,
|
|
use_dilation=False,
|
|
**kwargs):
|
|
self.thresh = thresh
|
|
self.box_thresh = box_thresh
|
|
self.max_candidates = max_candidates
|
|
self.unclip_ratio = unclip_ratio
|
|
self.min_size = 3
|
|
self.dilation_kernel = None if not use_dilation else np.array(
|
|
[[1, 1], [1, 1]])
|
|
|
|
def boxes_from_bitmap(self, pred, _bitmap, dest_width, dest_height):
|
|
'''
|
|
_bitmap: single map with shape (1, H, W),
|
|
whose values are binarized as {0, 1}
|
|
'''
|
|
|
|
bitmap = _bitmap
|
|
height, width = bitmap.shape
|
|
|
|
outs = cv2.findContours((bitmap * 255).astype(np.uint8), cv2.RETR_LIST,
|
|
cv2.CHAIN_APPROX_SIMPLE)
|
|
if len(outs) == 3:
|
|
img, contours, _ = outs[0], outs[1], outs[2]
|
|
elif len(outs) == 2:
|
|
contours, _ = outs[0], outs[1]
|
|
|
|
num_contours = min(len(contours), self.max_candidates)
|
|
|
|
boxes = []
|
|
scores = []
|
|
for index in range(num_contours):
|
|
contour = contours[index]
|
|
points, sside = self.get_mini_boxes(contour)
|
|
if sside < self.min_size:
|
|
continue
|
|
points = np.array(points)
|
|
score = self.box_score_fast(pred, points.reshape(-1, 2))
|
|
if self.box_thresh > score:
|
|
continue
|
|
|
|
box = self.unclip(points).reshape(-1, 1, 2)
|
|
box, sside = self.get_mini_boxes(box)
|
|
if sside < self.min_size + 2:
|
|
continue
|
|
box = np.array(box)
|
|
|
|
box[:, 0] = np.clip(
|
|
np.round(box[:, 0] / width * dest_width), 0, dest_width)
|
|
box[:, 1] = np.clip(
|
|
np.round(box[:, 1] / height * dest_height), 0, dest_height)
|
|
boxes.append(box.astype(np.int16))
|
|
scores.append(score)
|
|
return np.array(boxes, dtype=np.int16), scores
|
|
|
|
def unclip(self, box):
|
|
unclip_ratio = self.unclip_ratio
|
|
poly = Polygon(box)
|
|
distance = poly.area * unclip_ratio / poly.length
|
|
offset = pyclipper.PyclipperOffset()
|
|
offset.AddPath(box, pyclipper.JT_ROUND, pyclipper.ET_CLOSEDPOLYGON)
|
|
expanded = np.array(offset.Execute(distance))
|
|
return expanded
|
|
|
|
def get_mini_boxes(self, contour):
|
|
bounding_box = cv2.minAreaRect(contour)
|
|
points = sorted(list(cv2.boxPoints(bounding_box)), key=lambda x: x[0])
|
|
|
|
index_1, index_2, index_3, index_4 = 0, 1, 2, 3
|
|
if points[1][1] > points[0][1]:
|
|
index_1 = 0
|
|
index_4 = 1
|
|
else:
|
|
index_1 = 1
|
|
index_4 = 0
|
|
if points[3][1] > points[2][1]:
|
|
index_2 = 2
|
|
index_3 = 3
|
|
else:
|
|
index_2 = 3
|
|
index_3 = 2
|
|
|
|
box = [
|
|
points[index_1], points[index_2], points[index_3], points[index_4]
|
|
]
|
|
return box, min(bounding_box[1])
|
|
|
|
def box_score_fast(self, bitmap, _box):
|
|
h, w = bitmap.shape[:2]
|
|
box = _box.copy()
|
|
xmin = np.clip(np.floor(box[:, 0].min()).astype(np.int), 0, w - 1)
|
|
xmax = np.clip(np.ceil(box[:, 0].max()).astype(np.int), 0, w - 1)
|
|
ymin = np.clip(np.floor(box[:, 1].min()).astype(np.int), 0, h - 1)
|
|
ymax = np.clip(np.ceil(box[:, 1].max()).astype(np.int), 0, h - 1)
|
|
|
|
mask = np.zeros((ymax - ymin + 1, xmax - xmin + 1), dtype=np.uint8)
|
|
box[:, 0] = box[:, 0] - xmin
|
|
box[:, 1] = box[:, 1] - ymin
|
|
cv2.fillPoly(mask, box.reshape(1, -1, 2).astype(np.int32), 1)
|
|
return cv2.mean(bitmap[ymin:ymax + 1, xmin:xmax + 1], mask)[0]
|
|
|
|
def __call__(self, pred, shape_list):
|
|
pred = pred[:, 0, :, :]
|
|
segmentation = pred > self.thresh
|
|
|
|
boxes_batch = []
|
|
for batch_index in range(pred.shape[0]):
|
|
src_h, src_w, ratio_h, ratio_w = shape_list[batch_index]
|
|
if self.dilation_kernel is not None:
|
|
mask = cv2.dilate(
|
|
np.array(segmentation[batch_index]).astype(np.uint8),
|
|
self.dilation_kernel)
|
|
else:
|
|
mask = segmentation[batch_index]
|
|
boxes, scores = self.boxes_from_bitmap(pred[batch_index], mask,
|
|
src_w, src_h)
|
|
|
|
boxes_batch.append({'points': boxes})
|
|
return boxes_batch
|