File size: 4,125 Bytes
48c25e4
 
 
 
 
 
 
 
 
 
 
 
 
 
 
bd3c86d
48c25e4
 
 
 
 
 
 
 
 
51e90f5
48c25e4
 
 
 
 
 
 
 
 
 
 
 
 
ac810fa
48c25e4
 
 
cba2fbd
48c25e4
d1f7105
e30d296
ed9d291
 
48c25e4
 
 
 
51e90f5
 
 
 
 
 
48c25e4
 
 
 
 
51e90f5
 
48c25e4
 
51e90f5
48c25e4
 
 
 
 
 
8ff3b16
 
1309224
 
48c25e4
51e90f5
8ff3b16
 
ed9d291
8ff3b16
 
 
ed9d291
8ff3b16
48c25e4
 
d4a70e1
 
48c25e4
51e90f5
 
d4a70e1
48c25e4
 
 
 
 
 
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
# Scene Text Recognition Model Hub
# Copyright 2022 Darwin Bautista
#
# 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
#
#     https://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.

import glob

import torch
from torchvision import transforms as T

import gradio as gr


class App:

    title = 'Scene Text Recognition with<br/>Permuted Autoregressive Sequence Models'
    models = ['parseq', 'parseq_tiny', 'abinet', 'crnn', 'trba', 'vitstr']

    def __init__(self):
        self._model_cache = {}
        self._preprocess = T.Compose([
            T.Resize((32, 128), T.InterpolationMode.BICUBIC),
            T.ToTensor(),
            T.Normalize(0.5, 0.5)
        ])

    def _get_model(self, name):
        if name in self._model_cache:
            return self._model_cache[name]
        model = torch.hub.load('baudm/parseq', name, pretrained=True, trust_repo=True).eval()
        self._model_cache[name] = model
        return model

    @torch.inference_mode()
    def __call__(self, model_name, image):
        if image is None:
            return '', []
        if isinstance(image, dict):  # Extact image from ImageEditor output
            image = image['composite']
        model = self._get_model(model_name)
        image = self._preprocess(image.convert('RGB')).unsqueeze(0)
        # Greedy decoding
        pred = model(image).softmax(-1)
        label, _ = model.tokenizer.decode(pred)
        raw_label, raw_confidence = model.tokenizer.decode(pred, raw=True)
        # Format confidence values
        max_len = 25 if model_name == 'crnn' else len(label[0]) + 1
        conf = list(map('{:0.1f}'.format, raw_confidence[0][:max_len].tolist()))
        return label[0], [raw_label[0][:max_len], conf]


def main():
    app = App()

    with gr.Blocks(analytics_enabled=False, title=app.title.replace('<br/>', ' ')) as demo:
        gr.Markdown(f"""
            <div align="center">

            # {app.title}
            [![GitHub](https://img.shields.io/badge/baudm-parseq-blue?logo=github)](https://github.com/baudm/parseq)

            </div>

            To use this interactive demo for PARSeq and reproduced models:
            1. Select which model you want to use.
            2. Upload your own cropped image (or select from the given examples), or sketch on the canvas.
            3. Click **Read Text**.

            *NOTE*: None of these models were trained on handwritten text datasets.
        """)
        model_name = gr.Radio(app.models, value=app.models[0], label='The STR model to use')
        with gr.Tabs():
            with gr.TabItem('Image Upload'):
                image_upload = gr.Image(type='pil', sources=['upload'], label='Image')
                gr.Examples(glob.glob('demo_images/*.*'), inputs=image_upload)
                read_upload = gr.Button('Read Text')
            with gr.TabItem('Canvas Sketch'):
                image_canvas = gr.ImageEditor(type='pil', sources=[], label='Sketch', image_mode='RGB', layers=False, canvas_size=(768, 192))
                read_canvas = gr.Button('Read Text')

        output = gr.Textbox(max_lines=1, label='Model output')
        #adv_output = gr.Checkbox(label='Show detailed output')
        raw_output = gr.Dataframe(row_count=2, col_count=0, label='Raw output with confidence values ([0, 1] interval; [B] - BLANK token; [E] - EOS token)')

        read_upload.click(app, inputs=[model_name, image_upload], outputs=[output, raw_output])
        read_canvas.click(app, inputs=[model_name, image_canvas], outputs=[output, raw_output])
        #adv_output.change(lambda x: gr.update(visible=x), inputs=adv_output, outputs=raw_output)

    demo.launch()


if __name__ == '__main__':
    main()