Spaces:
Sleeping
Sleeping
| from pathlib import Path | |
| import cv2 | |
| import os | |
| import json | |
| from paddleocr.tools.infer.predict_system import TextSystem | |
| import matplotlib.image as mpimg | |
| import matplotlib.pyplot as plt | |
| from typing import Union, Any, Tuple, List, Optional, Dict | |
| from PIL import Image, ImageDraw, ImageFont | |
| import math | |
| import numpy as np | |
| from const import default_config | |
| def draw_ocr_box_txt(image, | |
| boxes, | |
| txts, | |
| scores=None, | |
| show_score=False, | |
| drop_score=0.5, | |
| font_path="./fonts/simfang.ttf"): | |
| 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, score) in enumerate(zip(boxes, txts, scores)): | |
| 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([(i[0], i[1]) for i in box], 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 show_score: | |
| txt = txt + ':' + str(score) | |
| 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_bbox = font.getbbox(c) | |
| draw_right.text( | |
| (box[0][0] + 3, cur_y), c, fill=(0, 0, 0), font=font) | |
| cur_y += char_bbox[3] - char_bbox[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 draw_ocr_results(image_fp: Union[str, Path, Image.Image], ocr_outs, font_path): | |
| # Credits: adapted from https://github.com/PaddlePaddle/PaddleOCR | |
| import cv2 | |
| if isinstance(image_fp, (str, Path)): | |
| img = Image.open(image_fp).convert('RGB') | |
| else: | |
| img = image_fp | |
| txts = [] | |
| scores = [] | |
| boxes = [] | |
| for _out in ocr_outs: | |
| txts.append(_out[1][0]) | |
| scores.append(_out[1][1]) | |
| boxes.append(_out[0]) | |
| draw_img = draw_ocr_box_txt( | |
| img, boxes, txts, scores, drop_score=0.0, font_path=font_path | |
| ) | |
| # cv2.imwrite(out_draw_fp, draw_img[:, :, ::-1]) | |
| #plt.figure(figsize=(draw_img.shape[0]/100, draw_img.shape[1]/100), dpi=100) | |
| #plt.imshow(draw_img) | |
| #plt.axis('off') # Turn off axis numbers | |
| #plt.show() | |
| return draw_img | |
| class MyObject: | |
| def __init__(self, dictionary): | |
| for key in dictionary: | |
| setattr(self, key, dictionary[key]) | |
| def json_to_class_obj(args): | |
| return MyObject(args) | |
| def start_ocr(params): | |
| args = json_to_class_obj({ | |
| **default_config, | |
| **params, | |
| }) | |
| imgs = [] | |
| if os.path.isdir(args.image_dir): | |
| for i in os.listdir(args.image_dir): | |
| imgs.append(cv2.imread(os.path.join(args.image_dir, i))) | |
| else: | |
| imgs.append(cv2.imread(args.image_dir)) | |
| ps = TextSystem(args) | |
| res = [] | |
| for img in imgs: | |
| dt_boxes, rec_res, _ = ps.__call__(img, args.use_angle_cls) | |
| ocr_res = [[box.tolist(), res] | |
| for box, res in zip(dt_boxes, rec_res)] | |
| res.append(ocr_res) | |
| return res | |