Spaces:
Running
Running
File size: 4,310 Bytes
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 |
# 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.
import argparse
import copy
import math
import time
from typing import List
import cv2
import numpy as np
try:
from .utils import ClsPostProcess, read_yaml, OrtInferSession
except:
from utils import ClsPostProcess, read_yaml, OrtInferSession
class TextClassifier(object):
def __init__(self, config):
self.cls_image_shape = config['cls_image_shape']
self.cls_batch_num = config['cls_batch_num']
self.cls_thresh = config['cls_thresh']
self.postprocess_op = ClsPostProcess(config['label_list'])
session_instance = OrtInferSession(config)
self.session = session_instance.session
self.input_name = session_instance.get_input_name()
def __call__(self, img_list: List[np.ndarray]):
if isinstance(img_list, np.ndarray):
img_list = [img_list]
img_list = copy.deepcopy(img_list)
# Calculate the aspect ratio of all text bars
width_list = [img.shape[1] / float(img.shape[0]) for img in img_list]
# Sorting can speed up the cls process
indices = np.argsort(np.array(width_list))
img_num = len(img_list)
cls_res = [['', 0.0]] * img_num
batch_num = self.cls_batch_num
elapse = 0
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 = []
for ino in range(beg_img_no, end_img_no):
norm_img = self.resize_norm_img(img_list[indices[ino]])
norm_img = norm_img[np.newaxis, :]
norm_img_batch.append(norm_img)
norm_img_batch = np.concatenate(norm_img_batch).astype(np.float32)
starttime = time.time()
onnx_inputs = {self.input_name: norm_img_batch}
prob_out = self.session.run(None, onnx_inputs)[0]
cls_result = self.postprocess_op(prob_out)
elapse += time.time() - starttime
for rno in range(len(cls_result)):
label, score = cls_result[rno]
cls_res[indices[beg_img_no + rno]] = [label, score]
if '180' in label and score > self.cls_thresh:
img_list[indices[beg_img_no + rno]] = cv2.rotate(
img_list[indices[beg_img_no + rno]], 1)
return img_list, cls_res, elapse
def resize_norm_img(self, img):
img_c, img_h, img_w = self.cls_image_shape
h, w = img.shape[:2]
ratio = w / float(h)
if math.ceil(img_h * ratio) > img_w:
resized_w = img_w
else:
resized_w = int(math.ceil(img_h * ratio))
resized_image = cv2.resize(img, (resized_w, img_h))
resized_image = resized_image.astype('float32')
if img_c == 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
padding_im = np.zeros((img_c, img_h, img_w), dtype=np.float32)
padding_im[:, :, :resized_w] = resized_image
return padding_im
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument('--image_path', type=str, help='image_dir|image_path')
parser.add_argument('--config_path', type=str, default='config.yaml')
args = parser.parse_args()
config = read_yaml(args.config_path)
text_classifier = TextClassifier(config)
img = cv2.imread(args.image_path)
img_list, cls_res, predict_time = text_classifier(img)
for ino in range(len(img_list)):
print(f"cls result:{cls_res[ino]}")
|