predict_rec.py 28.0 KB
Newer Older
L
LDOUBLEV 已提交
1 2 3 4 5 6 7 8 9 10 11 12 13
# 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.
L
LDOUBLEV 已提交
14 15
import os
import sys
T
Topdu 已提交
16
from PIL import Image
17
__dir__ = os.path.dirname(os.path.abspath(__file__))
L
LDOUBLEV 已提交
18
sys.path.append(__dir__)
littletomatodonkey's avatar
littletomatodonkey 已提交
19
sys.path.insert(0, os.path.abspath(os.path.join(__dir__, '../..')))
L
LDOUBLEV 已提交
20

L
LDOUBLEV 已提交
21 22
os.environ["FLAGS_allocator_strategy"] = 'auto_growth'

L
LDOUBLEV 已提交
23 24 25 26
import cv2
import numpy as np
import math
import time
W
WenmuZhou 已提交
27
import traceback
T
tink2123 已提交
28
import paddle
29 30

import tools.infer.utility as utility
W
WenmuZhou 已提交
31 32
from ppocr.postprocess import build_post_process
from ppocr.utils.logging import get_logger
33
from ppocr.utils.utility import get_image_file_list, check_and_read
L
LDOUBLEV 已提交
34

W
WenmuZhou 已提交
35 36
logger = get_logger()

L
LDOUBLEV 已提交
37 38 39

class TextRecognizer(object):
    def __init__(self, args):
40
        self.rec_image_shape = [int(v) for v in args.rec_image_shape.split(",")]
41
        self.rec_batch_num = args.rec_batch_num
T
tink2123 已提交
42
        self.rec_algorithm = args.rec_algorithm
W
WenmuZhou 已提交
43 44
        postprocess_params = {
            'name': 'CTCLabelDecode',
45
            "character_dict_path": args.rec_char_dict_path,
W
WenmuZhou 已提交
46
            "use_space_char": args.use_space_char
T
tink2123 已提交
47
        }
T
tink2123 已提交
48 49 50
        if self.rec_algorithm == "SRN":
            postprocess_params = {
                'name': 'SRNLabelDecode',
W
WenmuZhou 已提交
51 52 53 54 55 56
                "character_dict_path": args.rec_char_dict_path,
                "use_space_char": args.use_space_char
            }
        elif self.rec_algorithm == "RARE":
            postprocess_params = {
                'name': 'AttnLabelDecode',
T
tink2123 已提交
57 58 59
                "character_dict_path": args.rec_char_dict_path,
                "use_space_char": args.use_space_char
            }
T
Topdu 已提交
60 61 62 63 64 65
        elif self.rec_algorithm == 'NRTR':
            postprocess_params = {
                'name': 'NRTRLabelDecode',
                "character_dict_path": args.rec_char_dict_path,
                "use_space_char": args.use_space_char
            }
T
Topdu 已提交
66 67 68 69 70
        elif self.rec_algorithm == "SAR":
            postprocess_params = {
                'name': 'SARLabelDecode',
                "character_dict_path": args.rec_char_dict_path,
                "use_space_char": args.use_space_char
71
            }
A
andyjpaddle 已提交
72 73 74 75 76 77
        elif self.rec_algorithm == "VisionLAN":
            postprocess_params = {
                'name': 'VLLabelDecode',
                "character_dict_path": args.rec_char_dict_path,
                "use_space_char": args.use_space_char
            }
78 79 80 81 82 83 84 85 86 87 88
        elif self.rec_algorithm == 'ViTSTR':
            postprocess_params = {
                'name': 'ViTSTRLabelDecode',
                "character_dict_path": args.rec_char_dict_path,
                "use_space_char": args.use_space_char
            }
        elif self.rec_algorithm == 'ABINet':
            postprocess_params = {
                'name': 'ABINetLabelDecode',
                "character_dict_path": args.rec_char_dict_path,
                "use_space_char": args.use_space_char
xuyang2233's avatar
xuyang2233 已提交
89
            }
xuyang2233's avatar
add pr  
xuyang2233 已提交
90 91
        elif self.rec_algorithm == "SPIN":
            postprocess_params = {
xuyang2233's avatar
xuyang2233 已提交
92
                'name': 'SPINLabelDecode',
xuyang2233's avatar
add pr  
xuyang2233 已提交
93 94 95
                "character_dict_path": args.rec_char_dict_path,
                "use_space_char": args.use_space_char
            }
xuyang2233's avatar
xuyang2233 已提交
96 97 98 99 100 101 102
        elif self.rec_algorithm == "RobustScanner":
            postprocess_params = {
                'name': 'SARLabelDecode',
                "character_dict_path": args.rec_char_dict_path,
                "use_space_char": args.use_space_char,
                "rm_symbol": True
            }
103 104 105 106 107 108
        elif self.rec_algorithm == 'RFL':
            postprocess_params = {
                'name': 'RFLLabelDecode',
                "character_dict_path": None,
                "use_space_char": args.use_space_char
            }
109 110
        elif self.rec_algorithm == "PREN":
            postprocess_params = {'name': 'PRENLabelDecode'}
D
dorren 已提交
111 112 113
        elif self.rec_algorithm == "CAN":
            self.inverse = args.rec_image_inverse
            postprocess_params = {
114
                'name': 'CANLabelDecode',
D
dorren 已提交
115 116 117
                "character_dict_path": args.rec_char_dict_path,
                "use_space_char": args.use_space_char
            }
W
WenmuZhou 已提交
118
        self.postprocess_op = build_post_process(postprocess_params)
L
LDOUBLEV 已提交
119
        self.predictor, self.input_tensor, self.output_tensors, self.config = \
W
WenmuZhou 已提交
120
            utility.create_predictor(args, 'rec', logger)
T
tink2123 已提交
121
        self.benchmark = args.benchmark
T
tink2123 已提交
122
        self.use_onnx = args.use_onnx
T
tink2123 已提交
123 124 125
        if args.benchmark:
            import auto_log
            pid = os.getpid()
L
LDOUBLEV 已提交
126
            gpu_id = utility.get_infer_gpuid()
T
tink2123 已提交
127 128 129
            self.autolog = auto_log.AutoLogger(
                model_name="rec",
                model_precision=args.precision,
T
tink2123 已提交
130
                batch_size=args.rec_batch_num,
T
tink2123 已提交
131
                data_shape="dynamic",
132
                save_path=None,  #args.save_log_path,
T
tink2123 已提交
133 134 135
                inference_config=self.config,
                pids=pid,
                process_name=None,
L
LDOUBLEV 已提交
136
                gpu_ids=gpu_id if args.use_gpu else None,
T
tink2123 已提交
137 138 139
                time_keys=[
                    'preprocess_time', 'inference_time', 'postprocess_time'
                ],
T
tink2123 已提交
140
                warmup=0,
141
                logger=logger)
L
LDOUBLEV 已提交
142

143
    def resize_norm_img(self, img, max_wh_ratio):
L
LDOUBLEV 已提交
144
        imgC, imgH, imgW = self.rec_image_shape
145
        if self.rec_algorithm == 'NRTR' or self.rec_algorithm == 'ViTSTR':
T
Topdu 已提交
146 147 148
            img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
            # return padding_im
            image_pil = Image.fromarray(np.uint8(img))
149 150 151 152
            if self.rec_algorithm == 'ViTSTR':
                img = image_pil.resize([imgW, imgH], Image.BICUBIC)
            else:
                img = image_pil.resize([imgW, imgH], Image.ANTIALIAS)
T
Topdu 已提交
153 154 155
            img = np.array(img)
            norm_img = np.expand_dims(img, -1)
            norm_img = norm_img.transpose((2, 0, 1))
156 157 158 159 160
            if self.rec_algorithm == 'ViTSTR':
                norm_img = norm_img.astype(np.float32) / 255.
            else:
                norm_img = norm_img.astype(np.float32) / 128. - 1.
            return norm_img
161 162 163 164 165 166 167 168 169 170
        elif self.rec_algorithm == 'RFL':
            img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
            resized_image = cv2.resize(
                img, (imgW, imgH), interpolation=cv2.INTER_CUBIC)
            resized_image = resized_image.astype('float32')
            resized_image = resized_image / 255
            resized_image = resized_image[np.newaxis, :]
            resized_image -= 0.5
            resized_image /= 0.5
            return resized_image
T
Topdu 已提交
171

172
        assert imgC == img.shape[2]
A
andyjpaddle 已提交
173
        imgW = int((imgH * max_wh_ratio))
T
tink2123 已提交
174
        if self.use_onnx:
175 176 177 178
            w = self.input_tensor.shape[3:][0]
            if w is not None and w > 0:
                imgW = w

179
        h, w = img.shape[:2]
180 181 182 183 184
        ratio = w / float(h)
        if math.ceil(imgH * ratio) > imgW:
            resized_w = imgW
        else:
            resized_w = int(math.ceil(imgH * ratio))
A
andyjpaddle 已提交
185 186 187 188
        if self.rec_algorithm == 'RARE':
            if resized_w > self.rec_image_shape[2]:
                resized_w = self.rec_image_shape[2]
            imgW = self.rec_image_shape[2]
T
tink2123 已提交
189
        resized_image = cv2.resize(img, (resized_w, imgH))
L
LDOUBLEV 已提交
190 191 192 193 194 195 196
        resized_image = resized_image.astype('float32')
        resized_image = resized_image.transpose((2, 0, 1)) / 255
        resized_image -= 0.5
        resized_image /= 0.5
        padding_im = np.zeros((imgC, imgH, imgW), dtype=np.float32)
        padding_im[:, :, 0:resized_w] = resized_image
        return padding_im
T
tink2123 已提交
197

A
andyjpaddle 已提交
198 199 200 201 202 203 204 205 206 207
    def resize_norm_img_vl(self, img, image_shape):

        imgC, imgH, imgW = image_shape
        img = img[:, :, ::-1]  # bgr2rgb
        resized_image = cv2.resize(
            img, (imgW, imgH), interpolation=cv2.INTER_LINEAR)
        resized_image = resized_image.astype('float32')
        resized_image = resized_image.transpose((2, 0, 1)) / 255
        return resized_image

T
tink2123 已提交
208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279
    def resize_norm_img_srn(self, img, image_shape):
        imgC, imgH, imgW = image_shape

        img_black = np.zeros((imgH, imgW))
        im_hei = img.shape[0]
        im_wid = img.shape[1]

        if im_wid <= im_hei * 1:
            img_new = cv2.resize(img, (imgH * 1, imgH))
        elif im_wid <= im_hei * 2:
            img_new = cv2.resize(img, (imgH * 2, imgH))
        elif im_wid <= im_hei * 3:
            img_new = cv2.resize(img, (imgH * 3, imgH))
        else:
            img_new = cv2.resize(img, (imgW, imgH))

        img_np = np.asarray(img_new)
        img_np = cv2.cvtColor(img_np, cv2.COLOR_BGR2GRAY)
        img_black[:, 0:img_np.shape[1]] = img_np
        img_black = img_black[:, :, np.newaxis]

        row, col, c = img_black.shape
        c = 1

        return np.reshape(img_black, (c, row, col)).astype(np.float32)

    def srn_other_inputs(self, image_shape, num_heads, max_text_length):

        imgC, imgH, imgW = image_shape
        feature_dim = int((imgH / 8) * (imgW / 8))

        encoder_word_pos = np.array(range(0, feature_dim)).reshape(
            (feature_dim, 1)).astype('int64')
        gsrm_word_pos = np.array(range(0, max_text_length)).reshape(
            (max_text_length, 1)).astype('int64')

        gsrm_attn_bias_data = np.ones((1, max_text_length, max_text_length))
        gsrm_slf_attn_bias1 = np.triu(gsrm_attn_bias_data, 1).reshape(
            [-1, 1, max_text_length, max_text_length])
        gsrm_slf_attn_bias1 = np.tile(
            gsrm_slf_attn_bias1,
            [1, num_heads, 1, 1]).astype('float32') * [-1e9]

        gsrm_slf_attn_bias2 = np.tril(gsrm_attn_bias_data, -1).reshape(
            [-1, 1, max_text_length, max_text_length])
        gsrm_slf_attn_bias2 = np.tile(
            gsrm_slf_attn_bias2,
            [1, num_heads, 1, 1]).astype('float32') * [-1e9]

        encoder_word_pos = encoder_word_pos[np.newaxis, :]
        gsrm_word_pos = gsrm_word_pos[np.newaxis, :]

        return [
            encoder_word_pos, gsrm_word_pos, gsrm_slf_attn_bias1,
            gsrm_slf_attn_bias2
        ]

    def process_image_srn(self, img, image_shape, num_heads, max_text_length):
        norm_img = self.resize_norm_img_srn(img, image_shape)
        norm_img = norm_img[np.newaxis, :]

        [encoder_word_pos, gsrm_word_pos, gsrm_slf_attn_bias1, gsrm_slf_attn_bias2] = \
            self.srn_other_inputs(image_shape, num_heads, max_text_length)

        gsrm_slf_attn_bias1 = gsrm_slf_attn_bias1.astype(np.float32)
        gsrm_slf_attn_bias2 = gsrm_slf_attn_bias2.astype(np.float32)
        encoder_word_pos = encoder_word_pos.astype(np.int64)
        gsrm_word_pos = gsrm_word_pos.astype(np.int64)

        return (norm_img, encoder_word_pos, gsrm_word_pos, gsrm_slf_attn_bias1,
                gsrm_slf_attn_bias2)

T
Topdu 已提交
280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314
    def resize_norm_img_sar(self, img, image_shape,
                            width_downsample_ratio=0.25):
        imgC, imgH, imgW_min, imgW_max = image_shape
        h = img.shape[0]
        w = img.shape[1]
        valid_ratio = 1.0
        # make sure new_width is an integral multiple of width_divisor.
        width_divisor = int(1 / width_downsample_ratio)
        # resize
        ratio = w / float(h)
        resize_w = math.ceil(imgH * ratio)
        if resize_w % width_divisor != 0:
            resize_w = round(resize_w / width_divisor) * width_divisor
        if imgW_min is not None:
            resize_w = max(imgW_min, resize_w)
        if imgW_max is not None:
            valid_ratio = min(1.0, 1.0 * resize_w / imgW_max)
            resize_w = min(imgW_max, resize_w)
        resized_image = cv2.resize(img, (resize_w, imgH))
        resized_image = resized_image.astype('float32')
        # norm 
        if image_shape[0] == 1:
            resized_image = resized_image / 255
            resized_image = resized_image[np.newaxis, :]
        else:
            resized_image = resized_image.transpose((2, 0, 1)) / 255
        resized_image -= 0.5
        resized_image /= 0.5
        resize_shape = resized_image.shape
        padding_im = -1.0 * np.ones((imgC, imgH, imgW_max), dtype=np.float32)
        padding_im[:, :, 0:resize_w] = resized_image
        pad_shape = padding_im.shape

        return padding_im, resize_shape, pad_shape, valid_ratio

xuyang2233's avatar
add pr  
xuyang2233 已提交
315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330
    def resize_norm_img_spin(self, img):
        img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        # return padding_im
        img = cv2.resize(img, tuple([100, 32]), cv2.INTER_CUBIC)
        img = np.array(img, np.float32)
        img = np.expand_dims(img, -1)
        img = img.transpose((2, 0, 1))
        mean = [127.5]
        std = [127.5]
        mean = np.array(mean, dtype=np.float32)
        std = np.array(std, dtype=np.float32)
        mean = np.float32(mean.reshape(1, -1))
        stdinv = 1 / np.float32(std.reshape(1, -1))
        img -= mean
        img *= stdinv
        return img
A
andyjpaddle 已提交
331

332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360
    def resize_norm_img_svtr(self, img, image_shape):

        imgC, imgH, imgW = image_shape
        resized_image = cv2.resize(
            img, (imgW, imgH), interpolation=cv2.INTER_LINEAR)
        resized_image = resized_image.astype('float32')
        resized_image = resized_image.transpose((2, 0, 1)) / 255
        resized_image -= 0.5
        resized_image /= 0.5
        return resized_image

    def resize_norm_img_abinet(self, img, image_shape):

        imgC, imgH, imgW = image_shape

        resized_image = cv2.resize(
            img, (imgW, imgH), interpolation=cv2.INTER_LINEAR)
        resized_image = resized_image.astype('float32')
        resized_image = resized_image / 255.

        mean = np.array([0.485, 0.456, 0.406])
        std = np.array([0.229, 0.224, 0.225])
        resized_image = (
            resized_image - mean[None, None, ...]) / std[None, None, ...]
        resized_image = resized_image.transpose((2, 0, 1))
        resized_image = resized_image.astype('float32')

        return resized_image

D
dorren 已提交
361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384
    def norm_img_can(self, img, image_shape):

        img = cv2.cvtColor(
            img, cv2.COLOR_BGR2GRAY)  # CAN only predict gray scale image

        if self.inverse:
            img = 255 - img

        if self.rec_image_shape[0] == 1:
            h, w = img.shape
            _, imgH, imgW = self.rec_image_shape
            if h < imgH or w < imgW:
                padding_h = max(imgH - h, 0)
                padding_w = max(imgW - w, 0)
                img_padded = np.pad(img, ((0, padding_h), (0, padding_w)),
                                    'constant',
                                    constant_values=(255))
                img = img_padded

        img = np.expand_dims(img, 0) / 255.0  # h,w,c -> c,h,w
        img = img.astype('float32')

        return img

L
LDOUBLEV 已提交
385 386
    def __call__(self, img_list):
        img_num = len(img_list)
387
        # Calculate the aspect ratio of all text bars
388 389 390
        width_list = []
        for img in img_list:
            width_list.append(img.shape[1] / float(img.shape[0]))
张欣-男's avatar
张欣-男 已提交
391
        # Sorting can speed up the recognition process
392 393
        indices = np.argsort(np.array(width_list))
        rec_res = [['', 0.0]] * img_num
394
        batch_num = self.rec_batch_num
L
LDOUBLEV 已提交
395
        st = time.time()
T
tink2123 已提交
396 397
        if self.benchmark:
            self.autolog.times.start()
L
LDOUBLEV 已提交
398 399 400
        for beg_img_no in range(0, img_num, batch_num):
            end_img_no = min(img_num, beg_img_no + batch_num)
            norm_img_batch = []
X
xiaoting 已提交
401 402 403 404 405 406 407
            if self.rec_algorithm == "SRN":
                encoder_word_pos_list = []
                gsrm_word_pos_list = []
                gsrm_slf_attn_bias1_list = []
                gsrm_slf_attn_bias2_list = []
            if self.rec_algorithm == "SAR":
                valid_ratios = []
A
andyjpaddle 已提交
408
            imgC, imgH, imgW = self.rec_image_shape[:3]
A
andyjpaddle 已提交
409 410
            max_wh_ratio = imgW / imgH
            # max_wh_ratio = 0
L
LDOUBLEV 已提交
411
            for ino in range(beg_img_no, end_img_no):
412
                h, w = img_list[indices[ino]].shape[0:2]
413 414 415
                wh_ratio = w * 1.0 / h
                max_wh_ratio = max(max_wh_ratio, wh_ratio)
            for ino in range(beg_img_no, end_img_no):
T
Topdu 已提交
416
                if self.rec_algorithm == "SAR":
T
Topdu 已提交
417 418 419 420 421 422
                    norm_img, _, _, valid_ratio = self.resize_norm_img_sar(
                        img_list[indices[ino]], self.rec_image_shape)
                    norm_img = norm_img[np.newaxis, :]
                    valid_ratio = np.expand_dims(valid_ratio, axis=0)
                    valid_ratios.append(valid_ratio)
                    norm_img_batch.append(norm_img)
T
Topdu 已提交
423
                elif self.rec_algorithm == "SRN":
L
LDOUBLEV 已提交
424 425
                    norm_img = self.process_image_srn(
                        img_list[indices[ino]], self.rec_image_shape, 8, 25)
T
tink2123 已提交
426 427 428 429 430
                    encoder_word_pos_list.append(norm_img[1])
                    gsrm_word_pos_list.append(norm_img[2])
                    gsrm_slf_attn_bias1_list.append(norm_img[3])
                    gsrm_slf_attn_bias2_list.append(norm_img[4])
                    norm_img_batch.append(norm_img[0])
T
Topdu 已提交
431
                elif self.rec_algorithm == "SVTR":
T
tink2123 已提交
432 433
                    norm_img = self.resize_norm_img_svtr(img_list[indices[ino]],
                                                         self.rec_image_shape)
T
Topdu 已提交
434 435
                    norm_img = norm_img[np.newaxis, :]
                    norm_img_batch.append(norm_img)
436
                elif self.rec_algorithm in ["VisionLAN", "PREN"]:
A
andyjpaddle 已提交
437 438 439 440
                    norm_img = self.resize_norm_img_vl(img_list[indices[ino]],
                                                       self.rec_image_shape)
                    norm_img = norm_img[np.newaxis, :]
                    norm_img_batch.append(norm_img)
xuyang2233's avatar
add pr  
xuyang2233 已提交
441 442
                elif self.rec_algorithm == 'SPIN':
                    norm_img = self.resize_norm_img_spin(img_list[indices[ino]])
xuyang2233's avatar
xuyang2233 已提交
443 444
                    norm_img = norm_img[np.newaxis, :]
                    norm_img_batch.append(norm_img)
445 446 447 448
                elif self.rec_algorithm == "ABINet":
                    norm_img = self.resize_norm_img_abinet(
                        img_list[indices[ino]], self.rec_image_shape)
                    norm_img = norm_img[np.newaxis, :]
T
Topdu 已提交
449
                    norm_img_batch.append(norm_img)
xuyang2233's avatar
xuyang2233 已提交
450 451
                elif self.rec_algorithm == "RobustScanner":
                    norm_img, _, _, valid_ratio = self.resize_norm_img_sar(
452 453 454
                        img_list[indices[ino]],
                        self.rec_image_shape,
                        width_downsample_ratio=0.25)
xuyang2233's avatar
xuyang2233 已提交
455 456 457 458 459 460 461 462 463
                    norm_img = norm_img[np.newaxis, :]
                    valid_ratio = np.expand_dims(valid_ratio, axis=0)
                    valid_ratios = []
                    valid_ratios.append(valid_ratio)
                    norm_img_batch.append(norm_img)
                    word_positions_list = []
                    word_positions = np.array(range(0, 40)).astype('int64')
                    word_positions = np.expand_dims(word_positions, axis=0)
                    word_positions_list.append(word_positions)
D
dorren 已提交
464 465 466 467 468 469 470 471 472 473 474
                elif self.rec_algorithm == "CAN":
                    norm_img = self.norm_img_can(img_list[indices[ino]],
                                                 max_wh_ratio)
                    norm_img = norm_img[np.newaxis, :]
                    norm_img_batch.append(norm_img)
                    norm_image_mask = np.ones(norm_img.shape, dtype='float32')
                    word_label = np.ones([1, 36], dtype='int64')
                    norm_img_mask_batch = []
                    word_label_list = []
                    norm_img_mask_batch.append(norm_image_mask)
                    word_label_list.append(word_label)
T
Topdu 已提交
475 476 477 478 479
                else:
                    norm_img = self.resize_norm_img(img_list[indices[ino]],
                                                    max_wh_ratio)
                    norm_img = norm_img[np.newaxis, :]
                    norm_img_batch.append(norm_img)
L
LDOUBLEV 已提交
480 481
            norm_img_batch = np.concatenate(norm_img_batch)
            norm_img_batch = norm_img_batch.copy()
T
tink2123 已提交
482 483
            if self.benchmark:
                self.autolog.times.stamp()
T
tink2123 已提交
484 485 486 487 488 489 490 491 492 493 494 495 496 497 498 499

            if self.rec_algorithm == "SRN":
                encoder_word_pos_list = np.concatenate(encoder_word_pos_list)
                gsrm_word_pos_list = np.concatenate(gsrm_word_pos_list)
                gsrm_slf_attn_bias1_list = np.concatenate(
                    gsrm_slf_attn_bias1_list)
                gsrm_slf_attn_bias2_list = np.concatenate(
                    gsrm_slf_attn_bias2_list)

                inputs = [
                    norm_img_batch,
                    encoder_word_pos_list,
                    gsrm_word_pos_list,
                    gsrm_slf_attn_bias1_list,
                    gsrm_slf_attn_bias2_list,
                ]
T
tink2123 已提交
500 501 502 503 504 505 506 507 508 509 510 511 512 513 514 515 516 517 518 519
                if self.use_onnx:
                    input_dict = {}
                    input_dict[self.input_tensor.name] = norm_img_batch
                    outputs = self.predictor.run(self.output_tensors,
                                                 input_dict)
                    preds = {"predict": outputs[2]}
                else:
                    input_names = self.predictor.get_input_names()
                    for i in range(len(input_names)):
                        input_tensor = self.predictor.get_input_handle(
                            input_names[i])
                        input_tensor.copy_from_cpu(inputs[i])
                    self.predictor.run()
                    outputs = []
                    for output_tensor in self.output_tensors:
                        output = output_tensor.copy_to_cpu()
                        outputs.append(output)
                    if self.benchmark:
                        self.autolog.times.stamp()
                    preds = {"predict": outputs[2]}
T
Topdu 已提交
520 521 522 523
            elif self.rec_algorithm == "SAR":
                valid_ratios = np.concatenate(valid_ratios)
                inputs = [
                    norm_img_batch,
A
andyjpaddle 已提交
524 525
                    np.array(
                        [valid_ratios], dtype=np.float32),
T
Topdu 已提交
526
                ]
T
tink2123 已提交
527 528 529 530 531 532
                if self.use_onnx:
                    input_dict = {}
                    input_dict[self.input_tensor.name] = norm_img_batch
                    outputs = self.predictor.run(self.output_tensors,
                                                 input_dict)
                    preds = outputs[0]
T
Topdu 已提交
533
                else:
T
tink2123 已提交
534 535 536 537
                    input_names = self.predictor.get_input_names()
                    for i in range(len(input_names)):
                        input_tensor = self.predictor.get_input_handle(
                            input_names[i])
xuyang2233's avatar
xuyang2233 已提交
538 539 540 541 542 543 544 545 546 547 548 549
                        input_tensor.copy_from_cpu(inputs[i])
                    self.predictor.run()
                    outputs = []
                    for output_tensor in self.output_tensors:
                        output = output_tensor.copy_to_cpu()
                        outputs.append(output)
                    if self.benchmark:
                        self.autolog.times.stamp()
                    preds = outputs[0]
            elif self.rec_algorithm == "RobustScanner":
                valid_ratios = np.concatenate(valid_ratios)
                word_positions_list = np.concatenate(word_positions_list)
550 551
                inputs = [norm_img_batch, valid_ratios, word_positions_list]

xuyang2233's avatar
xuyang2233 已提交
552 553 554 555 556 557 558 559 560 561 562
                if self.use_onnx:
                    input_dict = {}
                    input_dict[self.input_tensor.name] = norm_img_batch
                    outputs = self.predictor.run(self.output_tensors,
                                                 input_dict)
                    preds = outputs[0]
                else:
                    input_names = self.predictor.get_input_names()
                    for i in range(len(input_names)):
                        input_tensor = self.predictor.get_input_handle(
                            input_names[i])
T
tink2123 已提交
563 564 565 566 567 568 569 570
                        input_tensor.copy_from_cpu(inputs[i])
                    self.predictor.run()
                    outputs = []
                    for output_tensor in self.output_tensors:
                        output = output_tensor.copy_to_cpu()
                        outputs.append(output)
                    if self.benchmark:
                        self.autolog.times.stamp()
T
Topdu 已提交
571
                    preds = outputs[0]
D
dorren 已提交
572 573 574 575 576 577 578 579 580 581 582 583 584 585 586 587 588 589 590 591 592 593 594 595 596 597 598
            elif self.rec_algorithm == "CAN":
                norm_img_mask_batch = np.concatenate(norm_img_mask_batch)
                word_label_list = np.concatenate(word_label_list)
                inputs = [norm_img_batch, norm_img_mask_batch, word_label_list]
                if self.use_onnx:
                    input_dict = {}
                    input_dict[self.input_tensor.name] = norm_img_batch
                    outputs = self.predictor.run(self.output_tensors,
                                                 input_dict)
                    preds = outputs
                else:
                    input_names = self.predictor.get_input_names()
                    input_tensor = []
                    for i in range(len(input_names)):
                        input_tensor_i = self.predictor.get_input_handle(
                            input_names[i])
                        input_tensor_i.copy_from_cpu(inputs[i])
                        input_tensor.append(input_tensor_i)
                    self.input_tensor = input_tensor
                    self.predictor.run()
                    outputs = []
                    for output_tensor in self.output_tensors:
                        output = output_tensor.copy_to_cpu()
                        outputs.append(output)
                    if self.benchmark:
                        self.autolog.times.stamp()
                    preds = outputs
T
tink2123 已提交
599 600 601 602 603 604 605 606 607 608 609 610 611 612 613 614 615 616 617 618
            else:
                if self.use_onnx:
                    input_dict = {}
                    input_dict[self.input_tensor.name] = norm_img_batch
                    outputs = self.predictor.run(self.output_tensors,
                                                 input_dict)
                    preds = outputs[0]
                else:
                    self.input_tensor.copy_from_cpu(norm_img_batch)
                    self.predictor.run()
                    outputs = []
                    for output_tensor in self.output_tensors:
                        output = output_tensor.copy_to_cpu()
                        outputs.append(output)
                    if self.benchmark:
                        self.autolog.times.stamp()
                    if len(outputs) != 1:
                        preds = outputs
                    else:
                        preds = outputs[0]
W
WenmuZhou 已提交
619 620 621
            rec_result = self.postprocess_op(preds)
            for rno in range(len(rec_result)):
                rec_res[indices[beg_img_no + rno]] = rec_result[rno]
T
tink2123 已提交
622 623
            if self.benchmark:
                self.autolog.times.end(stamp=True)
L
LDOUBLEV 已提交
624
        return rec_res, time.time() - st
L
LDOUBLEV 已提交
625 626


627
def main(args):
D
dyning 已提交
628
    image_file_list = get_image_file_list(args.image_dir)
L
LDOUBLEV 已提交
629 630 631
    text_recognizer = TextRecognizer(args)
    valid_image_file_list = []
    img_list = []
L
LDOUBLEV 已提交
632

T
tink2123 已提交
633 634
    logger.info(
        "In PP-OCRv3, rec_image_shape parameter defaults to '3, 48, 320', "
T
tink2123 已提交
635
        "if you are using recognition model with PP-OCRv2 or an older version, please set --rec_image_shape='3,32,320"
T
tink2123 已提交
636
    )
637
    # warmup 2 times
L
LDOUBLEV 已提交
638
    if args.warmup:
T
tink2123 已提交
639
        img = np.random.uniform(0, 255, [48, 320, 3]).astype(np.uint8)
640
        for i in range(2):
L
LDOUBLEV 已提交
641
            res = text_recognizer([img] * int(args.rec_batch_num))
L
LDOUBLEV 已提交
642

L
LDOUBLEV 已提交
643
    for image_file in image_file_list:
644
        img, flag, _ = check_and_read(image_file)
L
LDOUBLEV 已提交
645 646
        if not flag:
            img = cv2.imread(image_file)
L
LDOUBLEV 已提交
647 648 649 650 651
        if img is None:
            logger.info("error in loading image:{}".format(image_file))
            continue
        valid_image_file_list.append(image_file)
        img_list.append(img)
L
LDOUBLEV 已提交
652 653 654 655 656 657 658 659 660 661
    try:
        rec_res, _ = text_recognizer(img_list)

    except Exception as E:
        logger.info(traceback.format_exc())
        logger.info(E)
        exit()
    for ino in range(len(img_list)):
        logger.info("Predicts of {}:{}".format(valid_image_file_list[ino],
                                               rec_res[ino]))
T
tink2123 已提交
662 663
    if args.benchmark:
        text_recognizer.autolog.report()
664 665 666 667


if __name__ == "__main__":
    main(utility.parse_args())