mirror of
https://github.com/RapidAI/RapidOCR.git
synced 2026-11-11 05:20:49 +08:00
66 lines
2.0 KiB
Python
66 lines
2.0 KiB
Python
# -*- encoding: utf-8 -*-
|
|
# @Author: SWHL
|
|
# @Contact: liekkaskono@163.com
|
|
from pathlib import Path
|
|
from typing import Dict, Optional, Union
|
|
|
|
import numpy as np
|
|
import torch
|
|
import yaml
|
|
|
|
root_dir = Path(__file__).resolve().parent.parent
|
|
DEFAULT_CFG_PATH = root_dir / "arch_config.yaml"
|
|
|
|
|
|
def read_yaml(yaml_path: Union[str, Path]) -> Dict[str, Dict]:
|
|
with open(yaml_path, "rb") as f:
|
|
data = yaml.load(f, Loader=yaml.Loader)
|
|
return data
|
|
|
|
|
|
from rapidocr_torch.modeling.architectures.base_model import BaseModel
|
|
|
|
from .logger import get_logger
|
|
|
|
|
|
class TorchInferSession:
|
|
def __init__(self, config, mode: Optional[str] = None) -> None:
|
|
all_arch_config = read_yaml(DEFAULT_CFG_PATH)
|
|
|
|
self.logger = get_logger("TorchInferSession")
|
|
self.mode = mode
|
|
model_path = Path(config["model_path"])
|
|
self._verify_model(model_path)
|
|
file_name = model_path.stem
|
|
if file_name not in all_arch_config:
|
|
raise ValueError(f"architecture {file_name} is not in config.yaml")
|
|
arch_config = all_arch_config[file_name]
|
|
self.predictor = BaseModel(arch_config)
|
|
self.predictor.load_state_dict(torch.load(model_path, weights_only=True))
|
|
self.predictor.eval()
|
|
self.use_gpu = False
|
|
if config["use_cuda"]:
|
|
self.predictor.cuda()
|
|
self.use_gpu = True
|
|
|
|
def __call__(self, img: np.ndarray):
|
|
with torch.no_grad():
|
|
inp = torch.from_numpy(img)
|
|
if self.use_gpu:
|
|
inp = inp.cuda()
|
|
# 适配跟onnx对齐取值逻辑
|
|
outputs = self.predictor(inp).unsqueeze(0)
|
|
return outputs.cpu().numpy()
|
|
|
|
@staticmethod
|
|
def _verify_model(model_path):
|
|
model_path = Path(model_path)
|
|
if not model_path.exists():
|
|
raise FileNotFoundError(f"{model_path} does not exists.")
|
|
if not model_path.is_file():
|
|
raise FileExistsError(f"{model_path} is not a file.")
|
|
|
|
|
|
class TorchInferError(Exception):
|
|
pass
|