openocr-python 0.0.2__py3-none-any.whl

openocr/tools/export_rec.py

@@ -0,0 +1,118 @@
+import os
+import sys
+
+__dir__ = os.path.dirname(os.path.abspath(__file__))
+sys.path.append(__dir__)
+sys.path.insert(0, os.path.abspath(os.path.join(__dir__, '..')))
+
+import torch
+
+from openrec.modeling import build_model
+from openrec.postprocess import build_post_process
+from tools.engine import Config
+from tools.infer_rec import build_rec_process
+from tools.utility import ArgsParser
+from tools.utils.ckpt import load_ckpt
+from tools.utils.logging import get_logger
+
+
+def to_onnx(model, dummy_input, dynamic_axes, sava_path='model.onnx'):
+    input_axis_name = ['batch_size', 'channel', 'in_width', 'int_height']
+    output_axis_name = ['batch_size', 'channel', 'out_width', 'out_height']
+    torch.onnx.export(
+        model.to('cpu'),
+        dummy_input,
+        sava_path,
+        input_names=['input'],
+        output_names=['output'],  # the model's output names
+        dynamic_axes={
+            'input': {axis: input_axis_name[axis]
+                      for axis in dynamic_axes},
+            'output': {axis: output_axis_name[axis]
+                       for axis in dynamic_axes},
+        },
+    )
+
+
+def export_single_model(model: torch.nn.Module, _cfg, export_dir,
+                        export_config, logger, type):
+    for layer in model.modules():
+        if hasattr(layer, 'rep') and not getattr(layer, 'is_repped'):
+            layer.rep()
+    os.makedirs(export_dir, exist_ok=True)
+
+    export_cfg = {'PostProcess': _cfg['PostProcess']}
+    export_cfg['Transforms'] = build_rec_process(_cfg)
+
+    cfg.save(os.path.join(export_dir, 'config.yaml'), export_cfg)
+
+    dummy_input = torch.randn(*export_config['export_shape'], device='cpu')
+    if type == 'script':
+        save_path = os.path.join(export_dir, 'model.pt')
+        trace_model = torch.jit.trace(model, dummy_input, strict=False)
+        torch.jit.save(trace_model, save_path)
+    elif type == 'onnx':
+        save_path = os.path.join(export_dir, 'model.onnx')
+        to_onnx(model, dummy_input, export_config.get('dynamic_axes', []),
+                save_path)
+    else:
+        raise NotImplementedError
+    logger.info(f'finish export model to {save_path}')
+
+
+def main(cfg, type):
+    _cfg = cfg.cfg
+    logger = get_logger()
+    global_config = _cfg['Global']
+    export_config = _cfg['Export']
+    # build post process
+    post_process_class = build_post_process(_cfg['PostProcess'])
+    char_num = len(getattr(post_process_class, 'character'))
+    cfg['Architecture']['Decoder']['out_channels'] = char_num
+    model = build_model(_cfg['Architecture'])
+
+    load_ckpt(model, _cfg)
+    model.eval()
+
+    export_dir = export_config.get('export_dir', '')
+    if not export_dir:
+        export_dir = os.path.join(global_config.get('output_dir', 'output'),
+                                  'export')
+
+    if _cfg['Architecture']['algorithm'] in ['Distillation'
+                                             ]:  # distillation model
+        _cfg['PostProcess'][
+            'name'] = post_process_class.__class__.__base__.__name__
+        for model_name in model.model_list:
+            sub_model_save_path = os.path.join(export_dir, model_name)
+            export_single_model(
+                model.model_list[model_name],
+                _cfg,
+                sub_model_save_path,
+                export_config,
+                logger,
+                type,
+            )
+    else:
+        export_single_model(model, _cfg, export_dir, export_config, logger,
+                            type)
+
+
+def parse_args():
+    parser = ArgsParser()
+    parser.add_argument('--type',
+                        type=str,
+                        default='onnx',
+                        help='type of export')
+    args = parser.parse_args()
+    return args
+
+
+if __name__ == '__main__':
+    FLAGS = parse_args()
+    cfg = Config(FLAGS.config)
+    FLAGS = vars(FLAGS)
+    opt = FLAGS.pop('opt')
+    cfg.merge_dict(FLAGS)
+    cfg.merge_dict(opt)
+    main(cfg, FLAGS['type'])

openocr/tools/infer/onnx_engine.py

@@ -0,0 +1,65 @@
+import os
+
+import onnxruntime
+
+
+class ONNXEngine:
+
+    def __init__(self, onnx_path, use_gpu):
+        """
+        :param onnx_path:
+        """
+        if not os.path.exists(onnx_path):
+            raise Exception(f'{onnx_path} is not exists')
+
+        providers = ['CPUExecutionProvider']
+        if use_gpu:
+            providers = ([
+                'TensorrtExecutionProvider',
+                'CUDAExecutionProvider',
+                'CPUExecutionProvider',
+            ], )
+        self.onnx_session = onnxruntime.InferenceSession(onnx_path,
+                                                         providers=providers)
+        self.input_name = self.get_input_name(self.onnx_session)
+        self.output_name = self.get_output_name(self.onnx_session)
+
+    def get_output_name(self, onnx_session):
+        """
+        output_name = onnx_session.get_outputs()[0].name
+        :param onnx_session:
+        :return:
+        """
+        output_name = []
+        for node in onnx_session.get_outputs():
+            output_name.append(node.name)
+        return output_name
+
+    def get_input_name(self, onnx_session):
+        """
+        input_name = onnx_session.get_inputs()[0].name
+        :param onnx_session:
+        :return:
+        """
+        input_name = []
+        for node in onnx_session.get_inputs():
+            input_name.append(node.name)
+        return input_name
+
+    def get_input_feed(self, input_name, image_numpy):
+        """
+        input_feed={self.input_name: image_numpy}
+        :param input_name:
+        :param image_numpy:
+        :return:
+        """
+        input_feed = {}
+        for name in input_name:
+            input_feed[name] = image_numpy
+        return input_feed
+
+    def run(self, image_numpy):
+        # 输入数据的类型必须与模型一致,以下三种写法都是可以的
+        input_feed = self.get_input_feed(self.input_name, image_numpy)
+        result = self.onnx_session.run(self.output_name, input_feed=input_feed)
+        return result

openocr/tools/infer/predict_rec.py

@@ -0,0 +1,140 @@
+import os
+import sys
+
+__dir__ = os.path.dirname(os.path.abspath(__file__))
+sys.path.append(__dir__)
+sys.path.insert(0, os.path.abspath(os.path.join(__dir__, '../..')))
+
+import math
+import time
+
+import cv2
+import numpy as np
+
+from openrec.postprocess import build_post_process
+from openrec.preprocess import create_operators, transform
+from tools.engine import Config
+from tools.infer.onnx_engine import ONNXEngine
+from tools.infer.utility import check_gpu, parse_args
+from tools.utils.logging import get_logger
+from tools.utils.utility import check_and_read, get_image_file_list
+
+logger = get_logger()
+
+
+class TextRecognizer(ONNXEngine):
+
+    def __init__(self, args):
+        if args.rec_model_dir is None or not os.path.exists(
+                args.rec_model_dir):
+            raise Exception(
+                f'args.rec_model_dir is set to {args.rec_model_dir}, but it is not exists'
+            )
+
+        onnx_path = os.path.join(args.rec_model_dir, 'model.onnx')
+        config_path = os.path.join(args.rec_model_dir, 'config.yaml')
+        super(TextRecognizer, self).__init__(onnx_path, args.use_gpu)
+
+        self.rec_image_shape = [
+            int(v) for v in args.rec_image_shape.split(',')
+        ]
+        self.rec_batch_num = args.rec_batch_num
+        self.rec_algorithm = args.rec_algorithm
+
+        cfg = Config(config_path).cfg
+        self.ops = create_operators(cfg['Transforms'][1:])
+        self.postprocess_op = build_post_process(cfg['PostProcess'])
+
+    def resize_norm_img(self, img, max_wh_ratio):
+        imgC, imgH, imgW = self.rec_image_shape
+        assert imgC == img.shape[2]
+        imgW = int((imgH * max_wh_ratio))
+        h, w = img.shape[:2]
+        ratio = w / float(h)
+        if math.ceil(imgH * ratio) > imgW:
+            resized_w = imgW
+        else:
+            resized_w = int(math.ceil(imgH * ratio))
+        resized_image = cv2.resize(img, (resized_w, imgH))
+        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
+
+    def __call__(self, img_list):
+        img_num = len(img_list)
+        # Calculate the aspect ratio of all text bars
+        width_list = []
+        for img in img_list:
+            width_list.append(img.shape[1] / float(img.shape[0]))
+        # Sorting can speed up the recognition process
+        indices = np.argsort(np.array(width_list))
+        rec_res = [['', 0.0]] * img_num
+        batch_num = self.rec_batch_num
+        st = time.time()
+        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 = []
+            imgC, imgH, imgW = self.rec_image_shape[:3]
+            max_wh_ratio = imgW / imgH
+            # max_wh_ratio = 0
+            for ino in range(beg_img_no, end_img_no):
+                h, w = img_list[indices[ino]].shape[0:2]
+                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):
+                if self.rec_algorithm == 'nrtr':
+                    norm_img = transform({'image': img_list[indices[ino]]},
+                                         self.ops)[0]
+                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)
+            norm_img_batch = np.concatenate(norm_img_batch)
+            norm_img_batch = norm_img_batch.copy()
+
+            preds = self.run(norm_img_batch)
+
+            if len(preds) == 1:
+                preds = preds[0]
+
+            rec_result = self.postprocess_op({'res': preds})
+            for rno in range(len(rec_result)):
+                rec_res[indices[beg_img_no + rno]] = rec_result[rno]
+        return rec_res, time.time() - st
+
+
+def main(args):
+    args.use_gpu = check_gpu(args.use_gpu)
+
+    image_file_list = get_image_file_list(args.image_dir)
+    text_recognizer = TextRecognizer(args)
+    valid_image_file_list = []
+    img_list = []
+
+    # warmup 2 times
+    if args.warmup:
+        img = np.random.uniform(0, 255, [48, 320, 3]).astype(np.uint8)
+        for i in range(2):
+            text_recognizer([img] * int(args.rec_batch_num))
+
+    for image_file in image_file_list:
+        img, flag, _ = check_and_read(image_file)
+        if not flag:
+            img = cv2.imread(image_file)
+        if img is None:
+            logger.info(f'error in loading image:{image_file}')
+            continue
+        valid_image_file_list.append(image_file)
+        img_list.append(img)
+    rec_res, _ = text_recognizer(img_list)
+    for ino in range(len(img_list)):
+        logger.info(f'result of {valid_image_file_list[ino]}:{rec_res[ino]}')
+
+
+if __name__ == '__main__':
+    main(parse_args())

openocr/tools/infer/utility.py

@@ -0,0 +1,234 @@
+import argparse
+import math
+
+import cv2
+import numpy as np
+import torch
+import PIL
+from PIL import Image, ImageDraw, ImageFont
+import random
+
+def str2bool(v):
+    return v.lower() in ('true', 'yes', 't', 'y', '1')
+
+
+def str2int_tuple(v):
+    return tuple([int(i.strip()) for i in v.split(',')])
+
+
+def init_args():
+    parser = argparse.ArgumentParser()
+    # params for prediction engine
+    parser.add_argument('--use_gpu', type=str2bool, default=False)
+
+    # params for text detector
+    parser.add_argument('--image_dir', type=str)
+    parser.add_argument('--det_algorithm', type=str, default='DB')
+    parser.add_argument('--det_model_dir', type=str)
+    parser.add_argument('--det_limit_side_len', type=float, default=960)
+    parser.add_argument('--det_limit_type', type=str, default='max')
+    parser.add_argument('--det_box_type', type=str, default='quad')
+
+    # DB parmas
+    parser.add_argument('--det_db_thresh', type=float, default=0.3)
+    parser.add_argument('--det_db_box_thresh', type=float, default=0.6)
+    parser.add_argument('--det_db_unclip_ratio', type=float, default=1.5)
+    parser.add_argument('--max_batch_size', type=int, default=10)
+    parser.add_argument('--use_dilation', type=str2bool, default=False)
+    parser.add_argument('--det_db_score_mode', type=str, default='fast')
+
+    # params for text recognizer
+    parser.add_argument('--rec_algorithm', type=str, default='SVTR_LCNet')
+    parser.add_argument('--rec_model_dir', type=str)
+    parser.add_argument('--rec_image_inverse', type=str2bool, default=True)
+    parser.add_argument('--rec_image_shape', type=str, default='3, 48, 320')
+    parser.add_argument('--rec_batch_num', type=int, default=6)
+    parser.add_argument('--max_text_length', type=int, default=25)
+    parser.add_argument('--vis_font_path',
+                        type=str,
+                        default='./doc/fonts/simfang.ttf')
+    parser.add_argument('--drop_score', type=float, default=0.5)
+
+    # params for text classifier
+    parser.add_argument('--use_angle_cls', type=str2bool, default=False)
+    parser.add_argument('--cls_model_dir', type=str)
+    parser.add_argument('--cls_image_shape', type=str, default='3, 48, 192')
+    parser.add_argument('--label_list', type=list, default=['0', '180'])
+    parser.add_argument('--cls_batch_num', type=int, default=6)
+    parser.add_argument('--cls_thresh', type=float, default=0.9)
+
+    parser.add_argument('--warmup', type=str2bool, default=False)
+
+    #
+    parser.add_argument('--output', type=str, default='./inference_results')
+    parser.add_argument('--save_crop_res', type=str2bool, default=False)
+    parser.add_argument('--crop_res_save_dir', type=str, default='./output')
+
+    # multi-process
+    parser.add_argument('--use_mp', type=str2bool, default=False)
+    parser.add_argument('--total_process_num', type=int, default=1)
+    parser.add_argument('--process_id', type=int, default=0)
+
+    parser.add_argument('--show_log', type=str2bool, default=True)
+    return parser
+
+
+def parse_args():
+    parser = init_args()
+    return parser.parse_args()
+
+def create_font(txt, sz, font_path="./doc/fonts/simfang.ttf"):
+    font_size = int(sz[1] * 0.99)
+    font = ImageFont.truetype(font_path, font_size, encoding="utf-8")
+    if int(PIL.__version__.split(".")[0]) < 10:
+        length = font.getsize(txt)[0]
+    else:
+        length = font.getlength(txt)
+
+    if length > sz[0]:
+        font_size = int(font_size * sz[0] / length)
+        font = ImageFont.truetype(font_path, font_size, encoding="utf-8")
+    return font
+
+def draw_box_txt_fine(img_size, box, txt, font_path="./doc/fonts/simfang.ttf"):
+    box_height = int(
+        math.sqrt((box[0][0] - box[3][0]) ** 2 + (box[0][1] - box[3][1]) ** 2)
+    )
+    box_width = int(
+        math.sqrt((box[0][0] - box[1][0]) ** 2 + (box[0][1] - box[1][1]) ** 2)
+    )
+
+    if box_height > 2 * box_width and box_height > 30:
+        img_text = Image.new("RGB", (box_height, box_width), (255, 255, 255))
+        draw_text = ImageDraw.Draw(img_text)
+        if txt:
+            font = create_font(txt, (box_height, box_width), font_path)
+            draw_text.text([0, 0], txt, fill=(0, 0, 0), font=font)
+        img_text = img_text.transpose(Image.ROTATE_270)
+    else:
+        img_text = Image.new("RGB", (box_width, box_height), (255, 255, 255))
+        draw_text = ImageDraw.Draw(img_text)
+        if txt:
+            font = create_font(txt, (box_width, box_height), font_path)
+            draw_text.text([0, 0], txt, fill=(0, 0, 0), font=font)
+
+    pts1 = np.float32(
+        [[0, 0], [box_width, 0], [box_width, box_height], [0, box_height]]
+    )
+    pts2 = np.array(box, dtype=np.float32)
+    M = cv2.getPerspectiveTransform(pts1, pts2)
+
+    img_text = np.array(img_text, dtype=np.uint8)
+    img_right_text = cv2.warpPerspective(
+        img_text,
+        M,
+        img_size,
+        flags=cv2.INTER_NEAREST,
+        borderMode=cv2.BORDER_CONSTANT,
+        borderValue=(255, 255, 255),
+    )
+    return img_right_text
+
+def draw_ocr_box_txt(
+    image,
+    boxes,
+    txts=None,
+    scores=None,
+    drop_score=0.5,
+    font_path="./doc/fonts/simfang.ttf",
+):
+    h, w = image.height, image.width
+    img_left = image.copy()
+    img_right = np.ones((h, w, 3), dtype=np.uint8) * 255
+    random.seed(0)
+
+    draw_left = ImageDraw.Draw(img_left)
+    if txts is None or len(txts) != len(boxes):
+        txts = [None] * len(boxes)
+    for idx, (box, txt) in enumerate(zip(boxes, txts)):
+        if scores is not None and scores[idx] < drop_score:
+            continue
+        color = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255))
+        if isinstance(box[0], list):
+            box = list(map(tuple, box))
+        draw_left.polygon(box, fill=color)
+        img_right_text = draw_box_txt_fine((w, h), box, txt, font_path)
+        pts = np.array(box, np.int32).reshape((-1, 1, 2))
+        cv2.polylines(img_right_text, [pts], True, color, 1)
+        img_right = cv2.bitwise_and(img_right, img_right_text)
+    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(Image.fromarray(img_right), (w, 0, w * 2, h))
+    return np.array(img_show)
+
+
+def get_rotate_crop_image(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
+    """
+    assert len(points) == 4, 'shape of points must be 4*2'
+    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
+
+
+def get_minarea_rect_crop(img, points):
+    bounding_box = cv2.minAreaRect(np.array(points).astype(np.int32))
+    points = sorted(list(cv2.boxPoints(bounding_box)), key=lambda x: x[0])
+
+    index_a, index_b, index_c, index_d = 0, 1, 2, 3
+    if points[1][1] > points[0][1]:
+        index_a = 0
+        index_d = 1
+    else:
+        index_a = 1
+        index_d = 0
+    if points[3][1] > points[2][1]:
+        index_b = 2
+        index_c = 3
+    else:
+        index_b = 3
+        index_c = 2
+
+    box = [points[index_a], points[index_b], points[index_c], points[index_d]]
+    crop_img = get_rotate_crop_image(img, np.array(box))
+    return crop_img
+
+
+def check_gpu(use_gpu):
+    if use_gpu and not torch.cuda.is_available():
+        use_gpu = False
+    return use_gpu
+
+
+if __name__ == '__main__':
+    pass