Optimize python inference part

python/rapid_ocr_api.py

@@ -0,0 +1,308 @@
+# !/usr/bin/env python
+# -*- encoding: utf-8 -*-
+import argparse
+import copy
+import imghdr
+import math
+import random
+from pathlib import Path
+
+import cv2
+import numpy as np
+from PIL import Image, ImageDraw, ImageFont
+
+from ch_ppocr_mobile_v2_cls import TextClassifier
+from ch_ppocr_mobile_v2_det import TextDetector
+from ch_ppocr_mobile_v2_rec import TextRecognizer
+
+
+def get_image_file_list(img_file):
+    if img_file is None or not Path(img_file).exists():
+        raise FileExistsError(f"not found any img file in {img_file}")
+
+    imgs_lists = []
+    img_end = ['jpg', 'bmp', 'png', 'jpeg', 'rgb', 'tif', 'tiff', 'gif']
+
+    def _get_img_path(img_file):
+        # 真实后缀
+        real_img_suffix = imghdr.what(img_file)
+
+        # 现在后缀
+        now_img_suffix = img_file.suffix[1:].lower()
+
+        if real_img_suffix in img_end:
+            # 真实后缀与现有后缀不一致
+            if now_img_suffix in img_end:
+                return img_file
+            else:
+                # 按照真实后缀重命名
+                new_img_file = img_file.with_suffix(f'.{real_img_suffix}')
+                img_file.rename(new_img_file)
+                return new_img_file
+        else:
+            return None
+
+    img_file = Path(img_file)
+    if img_file.is_file():
+        tmp_img_path = _get_img_path(img_file)
+        if tmp_img_path is not None:
+            imgs_lists.append(str(tmp_img_path))
+
+    elif img_file.is_dir():
+        for img_path in img_file.iterdir():
+            tmp_file_path = _get_img_path(img_path)
+            if tmp_file_path is not None:
+                imgs_lists.append(str(tmp_file_path))
+    else:
+        raise ValueError(f"The {img_file}'s format is not support!")
+
+    if len(imgs_lists) == 0:
+        raise FileExistsError(f"not found any img file in {img_file}")
+
+    return imgs_lists
+
+
+def check_and_read_gif(img_path):
+    if Path(img_path).name[-3:] in ['gif', 'GIF']:
+        gif = cv2.VideoCapture(img_path)
+        ret, frame = gif.read()
+        if not ret:
+            print("Cannot read {}. This gif image maybe corrupted.")
+            return None, False
+        if len(frame.shape) == 2 or frame.shape[-1] == 1:
+            frame = cv2.cvtColor(frame, cv2.COLOR_GRAY2RGB)
+        imgvalue = frame[:, :, ::-1]
+        return imgvalue, True
+    return None, False
+
+
+def draw_ocr_box_txt(image, boxes, txts,
+                     scores=None, text_score=0.5,
+                     font_path="./models/msyh.ttc"):
+    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/drive/folders/1x_a9KpCo_1blxH1xFOfgKVkw1HYRVywY')
+
+    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):
+    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, 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)
+
+    image_save = str(draw_img_save / f'infer_{Path(image_path).name}')
+    cv2.imwrite(image_save, draw_img[:, :, ::-1])
+    print(f'The infer result has saved in {image_save}')
+
+
+class TextSystem(object):
+    def __init__(self, det_model_path,
+                 rec_model_path,
+                 use_angle_cls=False,
+                 cls_model_path=None,
+                 keys_path=None):
+        super(TextSystem).__init__()
+        self.text_score = 0.5
+
+        self.text_detector = TextDetector(det_model_path)
+        self.text_recognizer = TextRecognizer(rec_model_path, keys_path)
+        self.use_angle_cls = use_angle_cls
+        if self.use_angle_cls:
+            self.text_classifier = TextClassifier(cls_model_path)
+
+    def get_rotate_crop_image(self, img, points):
+        '''
+        img_height, img_width = img.shape[0:2]
+        left = int(np.min(points[:, 0]))
+        right = int(np.max(points[:, 0]))
+        top = int(np.min(points[:, 1]))
+        bottom = int(np.max(points[:, 1]))
+        img_crop = img[top:bottom, left:right, :].copy()
+        points[:, 0] = points[:, 0] - left
+        points[:, 1] = points[:, 1] - top
+        '''
+        img_crop_width = int(
+            max(
+                np.linalg.norm(points[0] - points[1]),
+                np.linalg.norm(points[2] - points[3])))
+        img_crop_height = int(
+            max(
+                np.linalg.norm(points[0] - points[3]),
+                np.linalg.norm(points[1] - points[2])))
+        pts_std = np.float32([[0, 0], [img_crop_width, 0],
+                              [img_crop_width, img_crop_height],
+                              [0, img_crop_height]])
+        M = cv2.getPerspectiveTransform(points, pts_std)
+        dst_img = cv2.warpPerspective(
+            img,
+            M, (img_crop_width, img_crop_height),
+            borderMode=cv2.BORDER_REPLICATE,
+            flags=cv2.INTER_CUBIC)
+        dst_img_height, dst_img_width = dst_img.shape[0:2]
+        if dst_img_height * 1.0 / dst_img_width >= 1.5:
+            dst_img = np.rot90(dst_img)
+        return dst_img
+
+    @staticmethod
+    def sorted_boxes(dt_boxes):
+        """
+        Sort text boxes in order from top to bottom, left to right
+        args:
+            dt_boxes(array):detected text boxes with shape [4, 2]
+        return:
+            sorted boxes(array) with shape [4, 2]
+        """
+        num_boxes = dt_boxes.shape[0]
+        sorted_boxes = sorted(dt_boxes, key=lambda x: (x[0][1], x[0][0]))
+        _boxes = list(sorted_boxes)
+
+        for i in range(num_boxes - 1):
+            if abs(_boxes[i + 1][0][1] - _boxes[i][0][1]) < 10 and \
+                    (_boxes[i + 1][0][0] < _boxes[i][0][0]):
+                tmp = _boxes[i]
+                _boxes[i] = _boxes[i + 1]
+                _boxes[i + 1] = tmp
+        return _boxes
+
+    @staticmethod
+    def load_image(image_path):
+        img, flag = check_and_read_gif(image_path)
+        if not flag:
+            img = cv2.imread(image_path)
+        if img is None:
+            raise ValueError(f"error in loading image:{image_path}")
+        return img
+
+    def __call__(self, image_path):
+        img = self.load_image(image_path)
+        dt_boxes, elapse = self.text_detector(img)
+        print("dt_boxes num : {}, elapse : {}".format(
+            len(dt_boxes), elapse))
+        if dt_boxes is None or len(dt_boxes) < 1:
+            return None, None
+        dt_boxes = self.sorted_boxes(dt_boxes)
+
+        img_crop_list = []
+        for bno in range(len(dt_boxes)):
+            tmp_box = copy.deepcopy(dt_boxes[bno])
+            img_crop = self.get_rotate_crop_image(img, tmp_box)
+            img_crop_list.append(img_crop)
+
+        if self.use_angle_cls:
+            img_crop_list, angle_list, elapse = self.text_classifier(
+                img_crop_list)
+            print("cls num  : {}, elapse : {}".format(
+                len(img_crop_list), elapse))
+
+        rec_res, elapse = self.text_recognizer(img_crop_list)
+        print("rec_res num  : {}, elapse : {}".format(
+            len(rec_res), elapse))
+
+        filter_boxes, filter_rec_res = [], []
+        for box, rec_reuslt in zip(dt_boxes, rec_res):
+            text, score = rec_reuslt
+            if score >= self.text_score:
+                filter_boxes.append(box)
+                filter_rec_res.append(rec_reuslt)
+        return filter_boxes, filter_rec_res
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument('--det_model_path', type=str,
+                        default='models/ch_ppocr_mobile_v2.0_det_infer.onnx')
+    parser.add_argument('--cls_model_path', type=str,
+                        default='models/ch_ppocr_mobile_v2.0_cls_infer.onnx')
+    parser.add_argument('--rec_model_path', type=str,
+                        default='models/ch_ppocr_mobile_v2.0_rec_infer.onnx')
+
+    parser.add_argument('--image_path', type=str,
+                        default='test_images/det_images/')
+
+    parser.add_argument('--keys_path', type=str,
+                        default="ch_ppocr_mobile_v2_rec/ppocr_keys_v1.txt")
+    args = parser.parse_args()
+
+    if args.det_model_path.find('server') != -1:
+        from ch_ppocr_server_v2_det import TextDetector
+    else:
+        from ch_ppocr_mobile_v2_det import TextDetector
+
+    if args.rec_model_path.find('server') != -1:
+        from ch_ppocr_server_v2_rec import TextRecognizer
+    else:
+        if args.rec_model_path.find('en') != -1:
+            from en_number_ppocr_mobile_v2_rec import TextRecognizer
+        elif args.rec_model_path.find('japan') != -1:
+            from japan_ppocr_mobile_v2_rec import TextRecognizer
+        else:
+            from ch_ppocr_mobile_v2_rec import TextRecognizer
+
+    text_sys = TextSystem(args.det_model_path,
+                          args.rec_model_path,
+                          use_angle_cls=True,
+                          cls_model_path=args.cls_model_path,
+                          keys_path=args.keys_path)
+
+    image_file_list = get_image_file_list(args.image_path)
+    for image_path in image_file_list:
+        dt_boxes, rec_res = text_sys(image_path)
+        visualize(image_path, dt_boxes, rec_res)
+        print('=' * 10)