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/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 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) 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: logger.info('GPU is not available, using CPU.') 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 if not os.path.exists(config['Global']['pretrained_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 if not os.path.exists(config['Global']['pretrained_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 else: if config['Architecture']['algorithm'] == 'SVTRv2_mobile': if not os.path.exists(config['Global']['pretrained_model']): config['Global'][ 'pretrained_model'] = check_and_download_model( MODEL_NAME_REC, DOWNLOAD_URL_REC) config['Global']['character_dict_path'] = DEFAULT_DICT_PATH_REC elif config['Architecture']['algorithm'] == 'SVTRv2_server': if not os.path.exists(config['Global']['pretrained_model']): config['Global'][ 'pretrained_model'] = check_and_download_model( MODEL_NAME_REC_SERVER, DOWNLOAD_URL_REC_SERVER) 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) 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)