File size: 5,008 Bytes
5d6a0bb
829289b
 
 
 
5d6a0bb
 
 
 
 
 
9792e33
5d6a0bb
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
9792e33
461c5b6
5d6a0bb
 
461c5b6
 
 
 
5d6a0bb
 
9792e33
461c5b6
5d6a0bb
 
 
 
 
 
461c5b6
 
 
 
 
 
 
 
 
5d6a0bb
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
# -*- encoding: utf-8 -*-
import os

os.system('pip install -r requirements.txt')

import math
import random
from pathlib import Path
import time

import cv2
import gradio as gr
from rapidocr_onnxruntime import TextSystem
import numpy as np
from PIL import Image, ImageDraw, ImageFont

text_sys = TextSystem('config.yaml')


def draw_ocr_box_txt(image, boxes, txts, font_path,
                     scores=None, text_score=0.5):
    if not Path(font_path).exists():
        raise FileNotFoundError(f'The {font_path} does not exists! \n'
                                f'Please download the file in the https://drive.google.com/file/d/1evWVX38EFNwTq_n5gTFgnlv8tdaNcyIA/view?usp=sharing')

    h, w = image.height, image.width
    img_left = image.copy()
    img_right = Image.new('RGB', (w, h), (255, 255, 255))

    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] < text_score:
            continue

        color = (random.randint(0, 255),
                 random.randint(0, 255),
                 random.randint(0, 255))
        draw_left.polygon(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 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 visualize(image_path, boxes, rec_res, font_path="resources/fonts/FZYTK.TTF"):
    image = Image.open(image_path)
    txts = [rec_res[i][0] for i in range(len(rec_res))]
    scores = [rec_res[i][1] for i in range(len(rec_res))]

    draw_img = draw_ocr_box_txt(image, boxes,
                                txts, font_path,
                                scores,
                                text_score=0.5)

    draw_img_save = Path("./inference_results/")
    if not draw_img_save.exists():
        draw_img_save.mkdir(parents=True, exist_ok=True)

    time_stamp = time.strftime('%Y-%m-%d-%H-%M-%S', time.localtime(time.time()))
    image_save = str(draw_img_save / f'{time_stamp}_{Path(image_path).name}')
    cv2.imwrite(image_save, draw_img[:, :, ::-1])
    return image_save


def inference(img, box_thresh, unclip_ratio, text_score):
    img_path = img.name
    img = cv2.imread(img_path)
    dt_boxes, rec_res = text_sys(img,
                                 box_thresh=box_thresh,
                                 unclip_ratio=unclip_ratio,
                                 text_score=text_score)
    img_save_path = visualize(img_path, dt_boxes, rec_res)
    return img_save_path, rec_res


title = 'Rapid🗲OCR Demo (捷智OCR)'
description = 'Gradio demo for RapidOCR. Github Repo: https://github.com/RapidAI/RapidOCR'
article = "<p style='text-align: center'> Completely open source, free and support offline deployment of multi-platform and multi-language OCR SDK <a href='https://github.com/RapidAI/RapidOCR'>Github Repo</a></p>"
css = ".output_image, .input_image {height: 40rem !important; width: 100% !important;}"
gr.Interface(
    inference,
    inputs=[
        gr.inputs.Image(type='file', label='Input'),
        gr.Slider(minimum=0, maximum=1.0, value=0.5,
                  label='box_thresh', step=0.1),
        gr.Slider(minimum=1.5, maximum=2.0, value=1.6,
                  label='unclip_ratio', step=0.1),
        gr.Slider(minimum=0, maximum=1.0, value=0.5,
                  label='text_score', step=0.1),
    ],
    outputs=[
        gr.outputs.Image(type='file', label='Output_image'),
        gr.outputs.Textbox(type='text', label='Output_text')
    ],
    title=title,
    description=description,
    article=article,
    css=css,
    allow_flagging='never',
    ).launch(debug=True, enable_queue=True)