| 23 | |
| 24 | |
| 25 | class TextRecognizer(ONNXEngine): |
| 26 | |
| 27 | def __init__(self, args): |
| 28 | if args.rec_model_dir is None or not os.path.exists( |
| 29 | args.rec_model_dir): |
| 30 | raise Exception( |
| 31 | f'args.rec_model_dir is set to {args.rec_model_dir}, but it is not exists' |
| 32 | ) |
| 33 | |
| 34 | onnx_path = os.path.join(args.rec_model_dir, 'model.onnx') |
| 35 | config_path = os.path.join(args.rec_model_dir, 'config.yaml') |
| 36 | super(TextRecognizer, self).__init__(onnx_path, args.use_gpu) |
| 37 | |
| 38 | self.rec_image_shape = [ |
| 39 | int(v) for v in args.rec_image_shape.split(',') |
| 40 | ] |
| 41 | self.rec_batch_num = args.rec_batch_num |
| 42 | self.rec_algorithm = args.rec_algorithm |
| 43 | |
| 44 | cfg = Config(config_path).cfg |
| 45 | self.ops = create_operators(cfg['Transforms'][1:]) |
| 46 | self.postprocess_op = build_post_process(cfg['PostProcess']) |
| 47 | |
| 48 | def resize_norm_img(self, img, max_wh_ratio): |
| 49 | imgC, imgH, imgW = self.rec_image_shape |
| 50 | assert imgC == img.shape[2] |
| 51 | imgW = int((imgH * max_wh_ratio)) |
| 52 | h, w = img.shape[:2] |
| 53 | ratio = w / float(h) |
| 54 | if math.ceil(imgH * ratio) > imgW: |
| 55 | resized_w = imgW |
| 56 | else: |
| 57 | resized_w = int(math.ceil(imgH * ratio)) |
| 58 | resized_image = cv2.resize(img, (resized_w, imgH)) |
| 59 | resized_image = resized_image.astype('float32') |
| 60 | resized_image = resized_image.transpose((2, 0, 1)) / 255 |
| 61 | resized_image -= 0.5 |
| 62 | resized_image /= 0.5 |
| 63 | padding_im = np.zeros((imgC, imgH, imgW), dtype=np.float32) |
| 64 | padding_im[:, :, 0:resized_w] = resized_image |
| 65 | return padding_im |
| 66 | |
| 67 | def __call__(self, img_list): |
| 68 | img_num = len(img_list) |
| 69 | # Calculate the aspect ratio of all text bars |
| 70 | width_list = [] |
| 71 | for img in img_list: |
| 72 | width_list.append(img.shape[1] / float(img.shape[0])) |
| 73 | # Sorting can speed up the recognition process |
| 74 | indices = np.argsort(np.array(width_list)) |
| 75 | rec_res = [['', 0.0]] * img_num |
| 76 | batch_num = self.rec_batch_num |
| 77 | st = time.time() |
| 78 | for beg_img_no in range(0, img_num, batch_num): |
| 79 | end_img_no = min(img_num, beg_img_no + batch_num) |
| 80 | norm_img_batch = [] |
| 81 | imgC, imgH, imgW = self.rec_image_shape[:3] |
| 82 | max_wh_ratio = imgW / imgH |