Spaces:
Sleeping
Sleeping
# -*- coding: utf-8 -*- | |
# time: 2022/10/17 13:25 | |
# file: ocr_utils.py | |
import cv2 | |
import math | |
import numpy as np | |
from PIL import Image, ImageDraw, ImageFont | |
def resize_img(img, input_size=600): | |
""" | |
resize img and limit the longest side of the image to input_size | |
""" | |
img = np.array(img) | |
im_shape = img.shape | |
im_size_max = np.max(im_shape[0:2]) | |
im_scale = float(input_size) / float(im_size_max) | |
img = cv2.resize(img, None, None, fx=im_scale, fy=im_scale) | |
return img | |
def draw_ocr( | |
image, | |
boxes, | |
txts=None, | |
scores=None, | |
drop_score=0.5, | |
font_path="./fonts/font.ttf" | |
): | |
""" | |
Visualize the results of OCR detection and recognition | |
args: | |
image(Image|array): RGB image | |
boxes(list): boxes with shape(N, 4, 2) | |
txts(list): the texts | |
scores(list): txxs corresponding scores | |
drop_score(float): only scores greater than drop_threshold will be visualized | |
font_path: the path of font which is used to draw text | |
return(array): | |
the visualized img | |
""" | |
if scores is None: | |
scores = [1] * len(boxes) | |
box_num = len(boxes) | |
for i in range(box_num): | |
if scores is not None and (scores[i] < drop_score or math.isnan(scores[i])): | |
continue | |
box = np.reshape(np.array(boxes[i]), [-1, 1, 2]).astype(np.int64) | |
image = cv2.polylines(np.array(image), [box], True, (255, 0, 0), 2) | |
if txts is not None: | |
img = np.array(resize_img(image, input_size=600)) | |
txt_img = text_visual( | |
txts, | |
scores, | |
img_h=img.shape[0], | |
img_w=600, | |
threshold=drop_score, | |
font_path=font_path | |
) | |
img = np.concatenate([np.array(img), np.array(txt_img)], axis=1) | |
return img | |
return image | |
def draw_ocr_box_txt( | |
image, | |
boxes, | |
txts, | |
scores=None, | |
drop_score=0.5, | |
font_path="./fonts/font.ttf" | |
): | |
image = Image.fromarray(image) | |
h, w = image.height, image.width | |
img_left = image.copy() | |
img_right = Image.new('RGB', (w, h), (255, 255, 255)) | |
import random | |
random.seed(0) | |
draw_left = ImageDraw.Draw(img_left) | |
draw_right = ImageDraw.Draw(img_right) | |
for idx, (box, txt) in enumerate(zip(boxes, txts)): | |
if scores is not None and scores[idx] < drop_score: | |
continue | |
color = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255)) | |
draw_left.polygon( | |
[ | |
box[0][0], box[0][1], box[1][0], box[1][1], box[2][0], | |
box[2][1], box[3][0], box[3][1] | |
], | |
fill=color) | |
draw_right.polygon( | |
[ | |
box[0][0], box[0][1], box[1][0], box[1][1], box[2][0], | |
box[2][1], box[3][0], box[3][1] | |
], | |
outline=color) | |
box_height = math.sqrt((box[0][0] - box[3][0])**2 + (box[0][1] - box[3][ | |
1])**2) | |
box_width = math.sqrt((box[0][0] - box[1][0])**2 + (box[0][1] - box[1][ | |
1])**2) | |
if box_height > 2 * box_width: | |
font_size = max(int(box_width * 0.9), 10) | |
font = ImageFont.truetype(font_path, font_size, encoding="utf-8") | |
cur_y = box[0][1] | |
for c in txt: | |
char_size = font.getsize(c) | |
draw_right.text((box[0][0] + 3, cur_y), c, fill=(0, 0, 0), font=font) | |
cur_y += char_size[1] | |
else: | |
font_size = max(int(box_height * 0.8), 10) | |
font = ImageFont.truetype(font_path, font_size, encoding="utf-8") | |
draw_right.text([box[0][0], box[0][1]], txt, fill=(0, 0, 0), font=font) | |
img_left = Image.blend(image, img_left, 0.5) | |
img_show = Image.new('RGB', (w * 2, h), (255, 255, 255)) | |
img_show.paste(img_left, (0, 0, w, h)) | |
img_show.paste(img_right, (w, 0, w * 2, h)) | |
return np.array(img_show) | |
def str_count(s): | |
""" | |
Count the number of Chinese characters, | |
a single English character and a single number | |
equal to half the length of Chinese characters. | |
args: | |
s(string): the input of string | |
return(int): | |
the number of Chinese characters | |
""" | |
import string | |
count_zh = count_pu = 0 | |
s_len = len(s) | |
en_dg_count = 0 | |
for c in s: | |
if c in string.ascii_letters or c.isdigit() or c.isspace(): | |
en_dg_count += 1 | |
elif c.isalpha(): | |
count_zh += 1 | |
else: | |
count_pu += 1 | |
return s_len - math.ceil(en_dg_count / 2) | |
def text_visual( | |
texts, | |
scores, | |
img_h=400, | |
img_w=600, | |
threshold=0., | |
font_path="./fonts/font.ttf" | |
): | |
""" | |
create new blank img and draw txt on it | |
args: | |
texts(list): the text will be draw | |
scores(list|None): corresponding score of each txt | |
img_h(int): the height of blank img | |
img_w(int): the width of blank img | |
font_path: the path of font which is used to draw text | |
return(array): | |
""" | |
if scores is not None: | |
assert len(texts) == len(scores), "The number of txts and corresponding scores must match" | |
def create_blank_img(): | |
blank_img = np.ones(shape=[img_h, img_w], dtype=np.int8) * 255 | |
blank_img[:, img_w - 1:] = 0 | |
blank_img = Image.fromarray(blank_img).convert("RGB") | |
draw_txt = ImageDraw.Draw(blank_img) | |
return blank_img, draw_txt | |
blank_img, draw_txt = create_blank_img() | |
font_size = 20 | |
txt_color = (0, 0, 0) | |
font = ImageFont.truetype(font_path, font_size, encoding="utf-8") | |
gap = font_size + 5 | |
txt_img_list = [] | |
count, index = 1, 0 | |
for idx, txt in enumerate(texts): | |
index += 1 | |
if scores[idx] < threshold or math.isnan(scores[idx]): | |
index -= 1 | |
continue | |
first_line = True | |
while str_count(txt) >= img_w // font_size - 4: | |
tmp = txt | |
txt = tmp[:img_w // font_size - 4] | |
if first_line: | |
new_txt = str(index) + ': ' + txt | |
first_line = False | |
else: | |
new_txt = ' ' + txt | |
draw_txt.text((0, gap * count), new_txt, txt_color, font=font) | |
txt = tmp[img_w // font_size - 4:] | |
if count >= img_h // gap - 1: | |
txt_img_list.append(np.array(blank_img)) | |
blank_img, draw_txt = create_blank_img() | |
count = 0 | |
count += 1 | |
if first_line: | |
new_txt = str(index) + ': ' + txt + ' ' + '%.3f' % (scores[idx]) | |
else: | |
new_txt = " " + txt + " " + '%.3f' % (scores[idx]) | |
draw_txt.text((0, gap * count), new_txt, txt_color, font=font) | |
# whether add new blank img or not | |
if count >= img_h // gap - 1 and idx + 1 < len(texts): | |
txt_img_list.append(np.array(blank_img)) | |
blank_img, draw_txt = create_blank_img() | |
count = 0 | |
count += 1 | |
txt_img_list.append(np.array(blank_img)) | |
if len(txt_img_list) == 1: | |
blank_img = np.array(txt_img_list[0]) | |
else: | |
blank_img = np.concatenate(txt_img_list, axis=1) | |
return np.array(blank_img) | |
def base64_to_cv2(b64str): | |
import base64 | |
data = base64.b64decode(b64str.encode('utf8')) | |
data = np.fromstring(data, np.uint8) | |
data = cv2.imdecode(data, cv2.IMREAD_COLOR) | |
return data | |
def draw_boxes(image, boxes, scores=None, drop_score=0.5): | |
if scores is None: | |
scores = [1] * len(boxes) | |
for (box, score) in zip(boxes, scores): | |
if score < drop_score: | |
continue | |
box = np.reshape(np.array(box), [-1, 1, 2]).astype(np.int64) | |
image = cv2.polylines(np.array(image), [box], True, (255, 0, 0), 2) | |
return image | |
def get_rotate_crop_image(img, points): | |
''' | |
img_height, img_width = img.shape[0:2] | |
left = int(np.min(points[:, 0])) | |
right = int(np.max(points[:, 0])) | |
top = int(np.min(points[:, 1])) | |
bottom = int(np.max(points[:, 1])) | |
img_crop = img[top:bottom, left:right, :].copy() | |
points[:, 0] = points[:, 0] - left | |
points[:, 1] = points[:, 1] - top | |
''' | |
assert len(points) == 4, "shape of points must be 4*2" | |
img_crop_width = int( | |
max( | |
np.linalg.norm(points[0] - points[1]), | |
np.linalg.norm(points[2] - points[3]))) | |
img_crop_height = int( | |
max( | |
np.linalg.norm(points[0] - points[3]), | |
np.linalg.norm(points[1] - points[2]))) | |
pts_std = np.float32([[0, 0], [img_crop_width, 0], | |
[img_crop_width, img_crop_height], | |
[0, img_crop_height]]) | |
M = cv2.getPerspectiveTransform(points, pts_std) | |
dst_img = cv2.warpPerspective( | |
img, | |
M, (img_crop_width, img_crop_height), | |
borderMode=cv2.BORDER_REPLICATE, | |
flags=cv2.INTER_CUBIC | |
) | |
dst_img_height, dst_img_width = dst_img.shape[0:2] | |
if dst_img_height * 1.0 / dst_img_width >= 1.5: | |
dst_img = np.rot90(dst_img) | |
return dst_img | |
if __name__ == '__main__': | |
pass | |