OpenOCR/tools/infer_det.py

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from pathlib import Path
import time

import numpy as np

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__, '..')))

os.environ['FLAGS_allocator_strategy'] = 'auto_growth'

import cv2
import json
from tools.engine.config import Config
from tools.utility import ArgsParser
from tools.utils.logging import get_logger
from tools.utils.utility import get_image_file_list

logger = get_logger()

root_dir = Path(__file__).resolve().parent
DEFAULT_CFG_PATH_DET = str(root_dir / '../configs/det/dbnet/repvit_db.yml')

MODEL_NAME_DET = './openocr_det_repvit_ch.pth'  # 模型文件名称
DOWNLOAD_URL_DET = 'https://github.com/Topdu/OpenOCR/releases/download/develop0.0.1/openocr_det_repvit_ch.pth'  # 模型文件 URL
MODEL_NAME_DET_ONNX = './openocr_det_model.onnx'  # 模型文件名称
DOWNLOAD_URL_DET_ONNX = 'https://github.com/Topdu/OpenOCR/releases/download/develop0.0.1/openocr_det_model.onnx'  # 模型文件 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.error(f'Error downloading the model: {e}')
        # 提示用户手动下载
        logger.error(
            f'Unable to download the model automatically. '
            f'Please download the model manually from the following URL:\n{url}\n'
            f'and save it to: {model_name} or {model_path}')
        raise RuntimeError(
            f'Failed to download the model. Please download it manually from {url} '
            f'and save it to {model_path}') from e


def replace_batchnorm(net):
    import torch
    for child_name, child in net.named_children():
        if hasattr(child, 'fuse'):
            fused = child.fuse()
            setattr(net, child_name, fused)
            replace_batchnorm(fused)
        elif isinstance(child, torch.nn.BatchNorm2d):
            setattr(net, child_name, torch.nn.Identity())
        else:
            replace_batchnorm(child)


def draw_det_res(dt_boxes, img, img_name, save_path):
    src_im = img
    for box in dt_boxes:
        box = np.array(box).astype(np.int32).reshape((-1, 1, 2))
        cv2.polylines(src_im, [box], True, color=(255, 255, 0), thickness=2)
    if not os.path.exists(save_path):
        os.makedirs(save_path)
    save_path = os.path.join(save_path, os.path.basename(img_name))
    cv2.imwrite(save_path, src_im)


def set_device(device, numId=0):
    import torch
    if device == 'gpu' and torch.cuda.is_available():
        device = torch.device(f'cuda:{numId}')
    else:
        logger.info('GPU is not available, using CPU.')
        device = torch.device('cpu')
    return device


class OpenDetector(object):

    def __init__(self,
                 config=None,
                 backend='torch',
                 onnx_model_path=None,
                 numId=0):
        """
        Args:
            config (dict, optional): 配置信息。默认为None。
            backend (str): 'torch' 或 'onnx'
            onnx_model_path (str): ONNX模型路径（仅当backend='onnx'时需要）
            numId (int, optional): 设备编号。默认为0。
        """

        if config is None:
            config = Config(DEFAULT_CFG_PATH_DET).cfg

        self._init_common(config)
        backend = backend if config['Global'].get(
            'backend', None) is None else config['Global']['backend']
        self.backend = backend
        if backend == 'torch':
            import torch
            self.torch = torch
            if config['Architecture']['algorithm'] == 'DB_mobile':
                if not os.path.exists(config['Global']['pretrained_model']):
                    config['Global'][
                        'pretrained_model'] = check_and_download_model(
                            MODEL_NAME_DET, DOWNLOAD_URL_DET)
            self._init_torch_model(config, numId)
        elif backend == 'onnx':
            from tools.infer.onnx_engine import ONNXEngine
            onnx_model_path = onnx_model_path if config['Global'].get(
                'onnx_model_path',
                None) is None else config['Global']['onnx_model_path']
            if onnx_model_path is None:
                if config['Architecture']['algorithm'] == 'DB_mobile':
                    onnx_model_path = check_and_download_model(
                        MODEL_NAME_DET_ONNX, DOWNLOAD_URL_DET_ONNX)
                else:
                    raise ValueError('ONNX模式需要指定onnx_model_path参数')
            self.onnx_det_engine = ONNXEngine(
                onnx_model_path, use_gpu=config['Global']['device'] == 'gpu')
        else:
            raise ValueError("backend参数必须是'torch'或'onnx'")

    def _init_common(self, config):
        from opendet.postprocess import build_post_process
        from opendet.preprocess import create_operators, transform
        global_config = config['Global']
        # create data ops
        self.transform = transform
        transforms = []
        for op in config['Eval']['dataset']['transforms']:
            op_name = list(op)[0]
            if 'Label' in op_name:
                continue
            elif op_name == 'KeepKeys':
                op[op_name]['keep_keys'] = ['image', 'shape']
            transforms.append(op)

        self.ops = create_operators(transforms, global_config)
        # build post process
        self.post_process_class = build_post_process(config['PostProcess'],
                                                     global_config)

    def _init_torch_model(self, config, numId=0):

        from opendet.modeling import build_model as build_det_model
        from tools.utils.ckpt import load_ckpt

        # build model
        self.model = build_det_model(config['Architecture'])
        self.model.eval()
        load_ckpt(self.model, config)
        if config['Architecture']['algorithm'] == 'DB_mobile':
            replace_batchnorm(self.model.backbone)
        self.device = set_device(config['Global']['device'], numId=numId)
        self.model.to(device=self.device)

    def _inference_onnx(self, images):
        # ONNX输入需要为numpy数组
        return self.onnx_det_engine.run(images)

    def __call__(self,
                 img_path=None,
                 img_numpy_list=None,
                 img_numpy=None,
                 return_mask=False,
                 **kwargs):
        """
        对输入图像进行处理，并返回处理结果。

        Args:
            img_path (str, optional): 图像文件路径。默认为 None。
            img_numpy_list (list, optional): 图像数据列表，每个元素为 numpy 数组。默认为 None。
            img_numpy (numpy.ndarray, optional): 图像数据，numpy 数组格式。默认为 None。

        Returns:
            list: 包含处理结果的列表。每个元素为一个字典，包含 'boxes' 和 'elapse' 两个键。
                'boxes' 的值为检测到的目标框点集，'elapse' 的值为处理时间。

        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 img_idx in range(num_img):
            if img_numpy_list is not None:
                img = img_numpy_list[img_idx]
                data = {'image': img}
            elif img_path is not None:
                with open(img_path[img_idx], 'rb') as f:
                    img = f.read()
                    data = {'image': img}
                data = self.transform(data, self.ops[:1])
            if kwargs.get('det_input_size', None) is not None:
                data['max_sile_len'] = kwargs['det_input_size']
            batch = self.transform(data, self.ops[1:])

            images = np.expand_dims(batch[0], axis=0)
            shape_list = np.expand_dims(batch[1], axis=0)
            t_start = time.time()

            if self.backend == 'torch':
                images = self.torch.from_numpy(images).to(device=self.device)
                with self.torch.no_grad():
                    preds = self.model(images)
                kwargs['torch_tensor'] = True
            elif self.backend == 'onnx':
                preds_det = self._inference_onnx(images)
                preds = {'maps': preds_det[0]}
                kwargs['torch_tensor'] = False

            t_cost = time.time() - t_start
            post_result = self.post_process_class(preds, [None, shape_list],
                                                  **kwargs)

            info = {'boxes': post_result[0]['points'], 'elapse': t_cost}
            if return_mask:
                if isinstance(preds['maps'], self.torch.Tensor):
                    mask = preds['maps'].detach().cpu().numpy()
                else:
                    mask = preds['maps']
                info['mask'] = mask
            results.append(info)
        return results


def main(cfg):
    is_visualize = cfg['Global'].get('is_visualize', False)
    model = OpenDetector(cfg)

    save_res_path = './det_results/'
    if not os.path.exists(save_res_path):
        os.makedirs(save_res_path)
    sample_num = 0
    with open(save_res_path + '/det_results.txt', 'wb') as fout:
        for file in get_image_file_list(cfg['Global']['infer_img']):
            preds_result = model(img_path=file)[0]
            logger.info('{} infer_img: {}, time cost: {}'.format(
                sample_num, file, preds_result['elapse']))
            boxes = preds_result['boxes']
            dt_boxes_json = []
            for box in boxes:
                tmp_json = {}
                tmp_json['points'] = np.array(box).tolist()
                dt_boxes_json.append(tmp_json)
            if is_visualize:
                src_img = cv2.imread(file)
                draw_det_res(boxes, src_img, file, save_res_path)
                logger.info('The detected Image saved in {}'.format(
                    os.path.join(save_res_path, os.path.basename(file))))
            otstr = file + '\t' + json.dumps(dt_boxes_json) + '\n'
            logger.info('results: {}'.format(json.dumps(dt_boxes_json)))
            fout.write(otstr.encode())
            sample_num += 1
        logger.info(
            f"Results saved to {os.path.join(save_res_path, 'det_results.txt')}.)"
        )

    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)