File size: 2,456 Bytes
37170d6 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 |
import onnxruntime
import numpy as np
import cv2
from abc import ABC, abstractmethod
from typing import Any, Tuple, Union, List
class BaseONNX(ABC):
def __init__(self, model_path: str, input_size: Tuple[int, int]):
"""初始化ONNX模型基类
Args:
model_path (str): ONNX模型路径
input_size (tuple): 模型输入尺寸 (width, height)
"""
self.session = onnxruntime.InferenceSession(model_path)
self.input_name = self.session.get_inputs()[0].name
self.input_size = input_size
def load_image(self, image: Union[cv2.UMat, str]) -> cv2.UMat:
"""加载图像
Args:
image (Union[cv2.UMat, str]): 图像路径或cv2图像对象
Returns:
cv2.UMat: 加载的图像
"""
if isinstance(image, str):
return cv2.imread(image)
return image.copy()
@abstractmethod
def preprocess_image(self, img_bgr: cv2.UMat, *args, **kwargs) -> np.ndarray:
"""图像预处理抽象方法
Args:
img_bgr (cv2.UMat): BGR格式的输入图像
Returns:
np.ndarray: 预处理后的图像
"""
pass
@abstractmethod
def run_inference(self, image: np.ndarray) -> Any:
"""运行推理的抽象方法
Args:
image (np.ndarray): 预处理后的输入图像
Returns:
Any: 模型输出结果
"""
pass
@abstractmethod
def pred(self, image: Union[cv2.UMat, str], *args, **kwargs) -> Any:
"""预测的抽象方法
Args:
image (Union[cv2.UMat, str]): 输入图像或图像路径
Returns:
Any: 预测结果
"""
pass
@abstractmethod
def draw_pred(self, img: cv2.UMat, *args, **kwargs) -> cv2.UMat:
"""绘制预测结果的抽象方法
Args:
img (cv2.UMat): 要绘制的图像
Returns:
cv2.UMat: 绘制结果后的图像
"""
pass
def check_images_list(self, images: List[Union[cv2.UMat, str, np.ndarray]]):
"""
检查图像列表是否有效
"""
for image in images:
if not isinstance(image, cv2.UMat) and not isinstance(image, str) and not isinstance(image, np.ndarray):
raise ValueError("The images must be a list of cv2.UMat or str or np.ndarray.")
|