openocr-python 0.0.2__py3-none-any.whl

openocr/tools/infer_rec.py

@@ -0,0 +1,371 @@
+import os
+from pathlib import Path
+import sys
+import time
+
+__dir__ = os.path.dirname(os.path.abspath(__file__))
+
+sys.path.append(__dir__)
+sys.path.insert(0, os.path.abspath(os.path.join(__dir__, '..')))
+
+import numpy as np
+import torch
+from torchvision import transforms as T
+from torchvision.transforms import functional as F
+from tools.engine import Config
+from tools.utility import ArgsParser
+from tools.utils.ckpt import load_ckpt
+from tools.utils.logging import get_logger
+from tools.utils.utility import get_image_file_list
+from tools.infer_det import replace_batchnorm
+
+logger = get_logger()
+
+root_dir = Path(__file__).resolve().parent
+DEFAULT_CFG_PATH_REC_SERVER = str(root_dir / '../configs/det/svtrv2/svtrv2_ch.yml')
+DEFAULT_CFG_PATH_REC = str(root_dir / '../configs/rec/svtrv2/repsvtr_ch.yml')
+DEFAULT_DICT_PATH_REC = str(root_dir / './utils/ppocr_keys_v1.txt')
+
+MODEL_NAME_REC = './openocr_repsvtr_ch.pth'  # 模型文件名称
+DOWNLOAD_URL_REC = 'https://github.com/Topdu/OpenOCR/releases/download/develop0.0.1/openocr_repsvtr_ch.pth'  # 模型文件 URL
+MODEL_NAME_REC_SERVER = './openocr_svtrv2_ch.pth'  # 模型文件名称
+DOWNLOAD_URL_REC_SERVER = 'https://github.com/Topdu/OpenOCR/releases/download/develop0.0.1/openocr_svtrv2_ch.pth'  # 模型文件 URL
+
+
+def check_and_download_model(model_name: str, url: str):
+    """
+    检查预训练模型是否存在，若不存在则从指定 URL 下载到固定缓存目录。
+
+    Args:
+        model_name (str): 模型文件的名称，例如 "model.pt"
+        url (str): 模型文件的下载地址
+
+    Returns:
+        str: 模型文件的完整路径
+    """
+    if os.path.exists(model_name):
+        return model_name
+
+    # 固定缓存路径为用户主目录下的 ".cache/openocr"
+    cache_dir = Path.home() / '.cache' / 'openocr'
+    model_path = cache_dir / model_name
+
+    # 如果模型文件已存在，直接返回路径
+    if model_path.exists():
+        logger.info(f'Model already exists at: {model_path}')
+        return str(model_path)
+
+    # 如果文件不存在，下载模型
+    logger.info(f'Model not found. Downloading from {url}...')
+
+    # 创建缓存目录（如果不存在）
+    cache_dir.mkdir(parents=True, exist_ok=True)
+
+    try:
+        # 下载文件
+        import urllib.request
+        with urllib.request.urlopen(url) as response, open(model_path,
+                                                           'wb') as out_file:
+            out_file.write(response.read())
+        logger.info(f'Model downloaded and saved at: {model_path}')
+        return str(model_path)
+
+    except Exception as e:
+        logger.info(f'Error downloading the model: {e}')
+        raise
+
+
+class RatioRecTVReisze(object):
+
+    def __init__(self, cfg):
+        self.max_ratio = cfg['Eval']['loader'].get('max_ratio', 12)
+        self.base_shape = cfg['Eval']['dataset'].get(
+            'base_shape', [[64, 64], [96, 48], [112, 40], [128, 32]])
+        self.base_h = cfg['Eval']['dataset'].get('base_h', 32)
+        self.interpolation = T.InterpolationMode.BICUBIC
+        transforms = []
+        transforms.extend([
+            T.ToTensor(),
+            T.Normalize(0.5, 0.5),
+        ])
+        self.transforms = T.Compose(transforms)
+        self.ceil = cfg['Eval']['dataset'].get('ceil', False),
+
+    def __call__(self, data):
+        img = data['image']
+        imgH = self.base_h
+        w, h = img.size
+        if self.ceil:
+            gen_ratio = int(float(w) / float(h)) + 1
+        else:
+            gen_ratio = max(1, round(float(w) / float(h)))
+        ratio_resize = min(gen_ratio, self.max_ratio)
+        imgW, imgH = self.base_shape[ratio_resize -
+                                     1] if ratio_resize <= 4 else [
+                                         self.base_h *
+                                         ratio_resize, self.base_h
+                                     ]
+        resized_w = imgW
+        resized_image = F.resize(img, (imgH, resized_w),
+                                 interpolation=self.interpolation)
+        img = self.transforms(resized_image)
+        data['image'] = img
+        return data
+
+
+def build_rec_process(cfg):
+    transforms = []
+    ratio_resize_flag = True
+    for op in cfg['Eval']['dataset']['transforms']:
+        op_name = list(op)[0]
+        if 'Resize' in op_name:
+            ratio_resize_flag = False
+        if 'Label' in op_name:
+            continue
+        elif op_name in ['RecResizeImg']:
+            op[op_name]['infer_mode'] = True
+        elif op_name == 'KeepKeys':
+            if cfg['Architecture']['algorithm'] in ['SAR', 'RobustScanner']:
+                if 'valid_ratio' in op[op_name]['keep_keys']:
+                    op[op_name]['keep_keys'] = ['image', 'valid_ratio']
+                else:
+                    op[op_name]['keep_keys'] = ['image']
+            else:
+                op[op_name]['keep_keys'] = ['image']
+        transforms.append(op)
+    return transforms, ratio_resize_flag
+
+
+def set_device(device, numId=0):
+    if device == 'gpu' and torch.cuda.is_available():
+        device = torch.device(f'cuda:{numId}')
+    else:
+        device = torch.device('cpu')
+    return device
+
+
+class OpenRecognizer(object):
+
+    def __init__(self, config=None, mode='mobile', numId=0):
+        """
+        初始化方法。
+
+        Args:
+            config (dict, optional): 配置信息。默认为None。
+            mode (str, optional): 模式，'server' 或 'mobile'。默认为'mobile'。
+            numId (int, optional): 设备编号。默认为0。
+
+        Returns:
+            None
+
+        Raises:
+            无
+
+        """
+        if config is None:
+            if mode == 'server':
+                config = Config(
+                    DEFAULT_CFG_PATH_REC_SERVER).cfg  # server model
+                model_dir = check_and_download_model(MODEL_NAME_REC_SERVER,
+                                                     DOWNLOAD_URL_REC_SERVER)
+            else:
+                config = Config(DEFAULT_CFG_PATH_REC).cfg  # mobile model
+                model_dir = check_and_download_model(MODEL_NAME_REC,
+                                                     DOWNLOAD_URL_REC)
+            config['Global']['pretrained_model'] = model_dir
+        config['Global']['character_dict_path'] = DEFAULT_DICT_PATH_REC
+        global_config = config['Global']
+        self.cfg = config
+        if global_config['pretrained_model'] is None:
+            global_config[
+                'pretrained_model'] = global_config['output_dir'] + '/best.pth'
+        # build post process
+        from openrec.modeling import build_model as build_rec_model
+        from openrec.postprocess import build_post_process
+        from openrec.preprocess import create_operators, transform
+        self.transform = transform
+        self.post_process_class = build_post_process(config['PostProcess'],
+                                                     global_config)
+
+        char_num = self.post_process_class.get_character_num()
+        config['Architecture']['Decoder']['out_channels'] = char_num
+        # print(char_num)
+        self.model = build_rec_model(config['Architecture'])
+        load_ckpt(self.model, config)
+
+        # exit(0)
+        self.device = set_device(global_config['device'], numId=numId)
+        self.model.eval()
+        replace_batchnorm(self.model.encoder)
+        self.model.to(device=self.device)
+
+        transforms, ratio_resize_flag = build_rec_process(self.cfg)
+        global_config['infer_mode'] = True
+        self.ops = create_operators(transforms, global_config)
+        if ratio_resize_flag:
+            ratio_resize = RatioRecTVReisze(cfg=self.cfg)
+            self.ops.insert(-1, ratio_resize)
+
+    def __call__(self,
+                 img_path=None,
+                 img_numpy_list=None,
+                 img_numpy=None,
+                 batch_num=1):
+        """
+        调用函数，处理输入图像，并返回识别结果。
+
+        Args:
+            img_path (str, optional): 图像文件的路径。默认为 None。
+            img_numpy_list (list, optional): 包含多个图像 numpy 数组的列表。默认为 None。
+            img_numpy (numpy.ndarray, optional): 单个图像的 numpy 数组。默认为 None。
+            batch_num (int, optional): 每次处理的图像数量。默认为 1。
+
+        Returns:
+            list: 包含识别结果的列表，每个元素为一个字典，包含文件路径（如果有的话）、文本、分数和延迟时间。
+
+        Raises:
+            Exception: 如果没有提供图像路径或 numpy 数组，则引发异常。
+        """
+
+        if img_numpy is not None:
+            img_numpy_list = [img_numpy]
+            num_img = 1
+        elif img_path is not None:
+            img_path = get_image_file_list(img_path)
+            num_img = len(img_path)
+        elif img_numpy_list is not None:
+            num_img = len(img_numpy_list)
+        else:
+            raise Exception('No input image path or numpy array.')
+        results = []
+        for start_idx in range(0, num_img, batch_num):
+            batch_data = []
+            batch_others = []
+            batch_file_names = []
+
+            max_width, max_height = 0, 0
+            # Prepare batch data
+            for img_idx in range(start_idx, min(start_idx + batch_num,
+                                                num_img)):
+                if img_numpy_list is not None:
+                    img = img_numpy_list[img_idx]
+                    data = {'image': img}
+                elif img_path is not None:
+                    file_name = img_path[img_idx]
+                    with open(file_name, 'rb') as f:
+                        img = f.read()
+                        data = {'image': img}
+                    data = self.transform(data, self.ops[:1])
+                    batch_file_names.append(file_name)
+                batch = self.transform(data, self.ops[1:])
+                others = None
+                if self.cfg['Architecture']['algorithm'] in [
+                        'SAR', 'RobustScanner'
+                ]:
+                    valid_ratio = np.expand_dims(batch[-1], axis=0)
+                    batch_others.append(valid_ratio)
+                    # others = [torch.from_numpy(valid_ratio).to(device=self.device)]
+                resized_image = batch[0]
+                h, w = resized_image.shape[-2:]
+                max_width = max(max_width, w)
+                max_height = max(max_height, h)
+                batch_data.append(batch[0])
+
+            padded_batch_data = []
+            for resized_image in batch_data:
+                padded_image = np.zeros([1, 3, max_height, max_width],
+                                        dtype=np.float32)
+                h, w = resized_image.shape[-2:]
+
+                # Apply padding (bottom-right padding)
+                padded_image[:, :, :h, :
+                             w] = resized_image  # 0 is typically used for padding
+                padded_batch_data.append(padded_image)
+
+            if batch_others:
+                others = np.concatenate(batch_others, axis=0)
+            else:
+                others = None
+            images = np.concatenate(padded_batch_data, axis=0)
+            images = torch.from_numpy(images).to(device=self.device)
+
+            with torch.no_grad():
+                t_start = time.time()
+                preds = self.model(images, others)
+                torch.cuda.synchronize()
+                t_cost = time.time() - t_start
+            post_results = self.post_process_class(preds)
+
+            for i, post_result in enumerate(post_results):
+                if img_path is not None:
+                    info = {
+                        'file': batch_file_names[i],
+                        'text': post_result[0],
+                        'score': post_result[1],
+                        'elapse': t_cost
+                    }
+                else:
+                    info = {
+                        'text': post_result[0],
+                        'score': post_result[1],
+                        'elapse': t_cost
+                    }
+                results.append(info)
+
+        return results
+
+
+def main(cfg):
+    model = OpenRecognizer(cfg)
+
+    save_res_path = cfg['Global']['output_dir']
+    if not os.path.exists(save_res_path):
+        os.makedirs(save_res_path)
+
+    t_sum = 0
+    sample_num = 0
+    max_len = cfg['Global']['max_text_length']
+    text_len_time = [0 for _ in range(max_len)]
+    text_len_num = [0 for _ in range(max_len)]
+
+    sample_num = 0
+    with open(save_res_path + '/rec_results.txt', 'wb') as fout:
+        for file in get_image_file_list(cfg['Global']['infer_img']):
+
+            preds_result = model(img_path=file, batch_num=1)[0]
+
+            rec_text = preds_result['text']
+            score = preds_result['score']
+            t_cost = preds_result['elapse']
+            info = rec_text + '\t' + str(score)
+            text_len_num[min(max_len - 1, len(rec_text))] += 1
+            text_len_time[min(max_len - 1, len(rec_text))] += t_cost
+            logger.info(
+                f'{sample_num} {file}\t result: {info}, time cost: {t_cost}')
+            otstr = file + '\t' + info + '\n'
+            t_sum += t_cost
+            fout.write(otstr.encode())
+            sample_num += 1
+
+    print(text_len_num)
+    w_avg_t_cost = []
+    for l_t_cost, l_num in zip(text_len_time, text_len_num):
+        if l_num != 0:
+            w_avg_t_cost.append(l_t_cost / l_num)
+    print(w_avg_t_cost)
+    w_avg_t_cost = sum(w_avg_t_cost) / len(w_avg_t_cost)
+
+    logger.info(
+        f'Sample num: {sample_num}, Weighted Avg time cost: {t_sum/sample_num}, Avg time cost: {w_avg_t_cost}'
+    )
+    logger.info('success!')
+
+
+if __name__ == '__main__':
+    FLAGS = ArgsParser().parse_args()
+    cfg = Config(FLAGS.config)
+    FLAGS = vars(FLAGS)
+    opt = FLAGS.pop('opt')
+    cfg.merge_dict(FLAGS)
+    cfg.merge_dict(opt)
+    main(cfg.cfg)

openocr/tools/train_rec.py

@@ -0,0 +1,37 @@
+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__, '..')))
+
+from tools.engine import Config, Trainer
+from tools.utility import ArgsParser
+
+
+def parse_args():
+    parser = ArgsParser()
+    parser.add_argument(
+        '--eval',
+        action='store_true',
+        default=True,
+        help='Whether to perform evaluation in train',
+    )
+    args = parser.parse_args()
+    return args
+
+
+def main():
+    FLAGS = parse_args()
+    cfg = Config(FLAGS.config)
+    FLAGS = vars(FLAGS)
+    opt = FLAGS.pop('opt')
+    cfg.merge_dict(FLAGS)
+    cfg.merge_dict(opt)
+    trainer = Trainer(cfg, mode='train_eval' if FLAGS['eval'] else 'train')
+    trainer.train()
+
+
+if __name__ == '__main__':
+    main()

openocr/tools/utility.py

@@ -0,0 +1,45 @@
+from argparse import ArgumentParser, RawDescriptionHelpFormatter
+
+import yaml
+
+
+class ArgsParser(ArgumentParser):
+
+    def __init__(self):
+        super(ArgsParser,
+              self).__init__(formatter_class=RawDescriptionHelpFormatter)
+        self.add_argument('-c', '--config', help='configuration file to use')
+        self.add_argument('-o',
+                          '--opt',
+                          nargs='*',
+                          help='set configuration options')
+        self.add_argument('--local_rank')
+        self.add_argument('--local-rank')
+
+    def parse_args(self, argv=None):
+        args = super(ArgsParser, self).parse_args(argv)
+        assert args.config is not None, 'Please specify --config=configure_file_path.'
+        args.opt = self._parse_opt(args.opt)
+        return args
+
+    def _parse_opt(self, opts):
+        config = {}
+        if not opts:
+            return config
+        for s in opts:
+            s = s.strip()
+            k, v = s.split('=', 1)
+            if '.' not in k:
+                config[k] = yaml.load(v, Loader=yaml.Loader)
+            else:
+                keys = k.split('.')
+                if keys[0] not in config:
+                    config[keys[0]] = {}
+                cur = config[keys[0]]
+                for idx, key in enumerate(keys[1:]):
+                    if idx == len(keys) - 2:
+                        cur[key] = yaml.load(v, Loader=yaml.Loader)
+                    else:
+                        cur[key] = {}
+                        cur = cur[key]
+        return config

openocr/tools/utils/EN_symbol_dict.txt

@@ -0,0 +1,94 @@
+0
+1
+2
+3
+4
+5
+6
+7
+8
+9
+a
+b
+c
+d
+e
+f
+g
+h
+i
+j
+k
+l
+m
+n
+o
+p
+q
+r
+s
+t
+u
+v
+w
+x
+y
+z
+A
+B
+C
+D
+E
+F
+G
+H
+I
+J
+K
+L
+M
+N
+O
+P
+Q
+R
+S
+T
+U
+V
+W
+X
+Y
+Z
+!
+"
+#
+$
+%
+&
+'
+(
+)
+*
++
+,
+-
+.
+/
+:
+;
+<
+=
+>
+?
+@
+[
+\
+]
+^
+_
+`
+{
+|
+}
+~

openocr/tools/utils/ckpt.py

@@ -0,0 +1,87 @@
+import os
+
+import torch
+
+from tools.utils.logging import get_logger
+
+
+def save_ckpt(
+    model,
+    cfg,
+    optimizer,
+    lr_scheduler,
+    epoch,
+    global_step,
+    metrics,
+    is_best=False,
+    logger=None,
+    prefix=None,
+):
+    """
+    Saving checkpoints
+
+    :param epoch: current epoch number
+    :param log: logging information of the epoch
+    :param save_best: if True, rename the saved checkpoint to 'model_best.pth.tar'
+    """
+    if logger is None:
+        logger = get_logger()
+    if prefix is None:
+        if is_best:
+            save_path = os.path.join(cfg["Global"]["output_dir"], "best.pth")
+        else:
+            save_path = os.path.join(cfg["Global"]["output_dir"], "latest.pth")
+    else:
+        save_path = os.path.join(cfg["Global"]["output_dir"], prefix + ".pth")
+    state_dict = model.module.state_dict() if cfg["Global"]["distributed"] else model.state_dict()
+    state = {
+        "epoch": epoch,
+        "global_step": global_step,
+        "state_dict": state_dict,
+        "optimizer": None if is_best else optimizer.state_dict(),
+        "scheduler": None if is_best else lr_scheduler.state_dict(),
+        "config": cfg,
+        "metrics": metrics,
+    }
+    torch.save(state, save_path)
+    logger.info(f"save ckpt to {save_path}")
+
+
+def load_ckpt(model, cfg, optimizer=None, lr_scheduler=None, logger=None):
+    """
+    Resume from saved checkpoints
+    :param checkpoint_path: Checkpoint path to be resumed
+    """
+    if logger is None:
+        logger = get_logger()
+    checkpoints = cfg["Global"].get("checkpoints")
+    pretrained_model = cfg["Global"].get("pretrained_model")
+
+    status = {}
+    if checkpoints and os.path.exists(checkpoints):
+        checkpoint = torch.load(checkpoints, map_location=torch.device("cpu"))
+        model.load_state_dict(checkpoint["state_dict"], strict=True)
+        if optimizer is not None:
+            optimizer.load_state_dict(checkpoint["optimizer"])
+        if lr_scheduler is not None:
+            lr_scheduler.load_state_dict(checkpoint["scheduler"])
+        logger.info(f"resume from checkpoint {checkpoints} (epoch {checkpoint['epoch']})")
+
+        status["global_step"] = checkpoint["global_step"]
+        status["epoch"] = checkpoint["epoch"] + 1
+        status["metrics"] = checkpoint["metrics"]
+    elif pretrained_model and os.path.exists(pretrained_model):
+        load_pretrained_params(model, pretrained_model, logger)
+        logger.info(f"finetune from checkpoint {pretrained_model}")
+    else:
+        logger.info("train from scratch")
+    return status
+
+
+def load_pretrained_params(model, pretrained_model, logger):
+    checkpoint = torch.load(pretrained_model, map_location=torch.device("cpu"))
+    model.load_state_dict(checkpoint["state_dict"], strict=False)
+    for name in model.state_dict().keys():
+        if name not in checkpoint["state_dict"]:
+            logger.info(f"{name} is not in pretrained model")
+