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 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_REC_SERVER = str(root_dir / '../configs/rec/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 MODEL_NAME_REC_ONNX = './openocr_rec_model.onnx' # 模型文件名称 DOWNLOAD_URL_REC_ONNX = 'https://github.com/Topdu/OpenOCR/releases/download/develop0.0.1/openocr_rec_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 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) from torchvision import transforms as T from torchvision.transforms import functional as F self.F = F 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 = self.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): 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 OpenRecognizer: def __init__(self, config=None, mode='mobile', backend='torch', onnx_model_path=None, numId=0): """ Args: config (dict, optional): 配置信息。默认为None。 mode (str, optional): 模式，'server' 或 'mobile'。默认为'mobile'。 backend (str): 'torch' 或 'onnx' onnx_model_path (str): ONNX模型路径（仅当backend='onnx'时需要） numId (int, optional): 设备编号。默认为0。 """ if config is None: config_file = DEFAULT_CFG_PATH_REC_SERVER if mode == 'server' else DEFAULT_CFG_PATH_REC config = Config(config_file).cfg self.cfg = config # 公共初始化 self._init_common() 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 self._init_torch_model(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 not onnx_model_path: if self.cfg['Architecture']['algorithm'] == 'SVTRv2_mobile': onnx_model_path = check_and_download_model( MODEL_NAME_REC_ONNX, DOWNLOAD_URL_REC_ONNX) else: raise ValueError('ONNX模式需要指定onnx_model_path参数') self.onnx_rec_engine = ONNXEngine( onnx_model_path, use_gpu=config['Global']['device'] == 'gpu') else: raise ValueError("backend参数必须是'torch'或'onnx'") def _init_common(self): # 初始化公共组件 from openrec.postprocess import build_post_process from openrec.preprocess import create_operators, transform self.transform = transform # 构建预处理流程 algorithm_name = self.cfg['Architecture']['algorithm'] if algorithm_name in ['SVTRv2_mobile', 'SVTRv2_server']: self.cfg['Global']['character_dict_path'] = DEFAULT_DICT_PATH_REC self.post_process_class = build_post_process(self.cfg['PostProcess'], self.cfg['Global']) char_num = self.post_process_class.get_character_num() self.cfg['Architecture']['Decoder']['out_channels'] = char_num transforms, ratio_resize_flag = build_rec_process(self.cfg) self.ops = create_operators(transforms, self.cfg['Global']) if ratio_resize_flag: ratio_resize = RatioRecTVReisze(cfg=self.cfg) self.ops.insert(-1, ratio_resize) def _init_torch_model(self, numId): from tools.utils.ckpt import load_ckpt from tools.infer_det import replace_batchnorm # PyTorch专用初始化 algorithm_name = self.cfg['Architecture']['algorithm'] if algorithm_name in ['SVTRv2_mobile', 'SVTRv2_server']: if not os.path.exists(self.cfg['Global']['pretrained_model']): pretrained_model = check_and_download_model( MODEL_NAME_REC, DOWNLOAD_URL_REC ) if algorithm_name == 'SVTRv2_mobile' else check_and_download_model( MODEL_NAME_REC_SERVER, DOWNLOAD_URL_REC_SERVER) self.cfg['Global']['pretrained_model'] = pretrained_model from openrec.modeling import build_model as build_rec_model self.model = build_rec_model(self.cfg['Architecture']) load_ckpt(self.model, self.cfg) self.device = set_device(self.cfg['Global']['device'], numId) self.model.to(self.device) self.model.eval() if algorithm_name == 'SVTRv2_mobile': replace_batchnorm(self.model.encoder) def _inference_onnx(self, images): # ONNX输入需要为numpy数组 return self.onnx_rec_engine.run(images) 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) resized_image = batch[0] if isinstance( batch[0], np.ndarray) else batch[0].numpy() 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 = np.zeros( (len(batch_data), 3, max_height, max_width), dtype=np.float32) for i, img in enumerate(batch_data): h, w = img.shape[-2:] padded_batch[i, :, :h, :w] = img if batch_others: others = np.concatenate(batch_others, axis=0) else: others = None t_start = time.time() if self.backend == 'torch': images = self.torch.from_numpy(padded_batch).to( device=self.device) with self.torch.no_grad(): preds = self.model(images, others) # bs, len, num_classes torch_tensor = True elif self.backend == 'onnx': # ONNX推理 preds = self._inference_onnx(padded_batch) preds = preds[0] # bs, len, num_classes torch_tensor = False t_cost = time.time() - t_start post_results = self.post_process_class(preds, torch_tensor=torch_tensor) 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 = './rec_results/' 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 logger.info( f"Results saved to {os.path.join(save_res_path, 'rec_results.txt')}.)" ) 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)