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.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)