102 lines
3.5 KiB
Python
102 lines
3.5 KiB
Python
from typing import List, Dict
|
||
import cv2
|
||
import numpy as np
|
||
import onnxruntime as ort
|
||
|
||
|
||
def softmax(x: np.ndarray) -> np.ndarray:
|
||
x = x.astype(np.float32)
|
||
x = x - np.max(x)
|
||
exp = np.exp(x)
|
||
return exp / (np.sum(exp) + 1e-12)
|
||
|
||
|
||
def to_probabilities(output: np.ndarray) -> np.ndarray:
|
||
"""
|
||
兼容两种常见 ONNX 导出：
|
||
1. 模型直接输出 logits，需要 softmax；
|
||
2. 模型图中已经包含 softmax，直接输出概率。
|
||
|
||
对概率再次做 softmax 会把原本很高的置信度压到接近均匀分布。
|
||
"""
|
||
values = np.asarray(output, dtype=np.float32).reshape(-1)
|
||
total = float(np.sum(values))
|
||
is_probability_output = (
|
||
values.size > 0
|
||
and np.all(np.isfinite(values))
|
||
and float(np.min(values)) >= -1e-5
|
||
and float(np.max(values)) <= 1.0 + 1e-5
|
||
and abs(total - 1.0) <= 1e-3
|
||
)
|
||
if is_probability_output:
|
||
return values / max(total, 1e-12)
|
||
return softmax(values)
|
||
|
||
|
||
class ConvNeXtONNXClassifier:
|
||
def __init__(
|
||
self,
|
||
model_path: str,
|
||
class_names: List[str],
|
||
img_size: int = 224,
|
||
device: str = "cpu",
|
||
):
|
||
self.model_path = model_path
|
||
self.class_names = class_names
|
||
self.img_size = img_size
|
||
|
||
providers = ["CPUExecutionProvider"]
|
||
if device.lower() in ["cuda", "gpu"]:
|
||
providers = ["CUDAExecutionProvider", "CPUExecutionProvider"]
|
||
|
||
self.session = ort.InferenceSession(model_path, providers=providers)
|
||
active_providers = self.session.get_providers()
|
||
print("[ConvNeXt] providers:", active_providers)
|
||
if (
|
||
device.lower() in ["cuda", "gpu"]
|
||
and "CUDAExecutionProvider" not in active_providers
|
||
):
|
||
raise RuntimeError(
|
||
"请求使用 CUDA 推理，但 ConvNeXt 会话未启用 "
|
||
"CUDAExecutionProvider。请安装匹配 CUDA/cuDNN 的 "
|
||
"onnxruntime-gpu。"
|
||
)
|
||
self.input_name = self.session.get_inputs()[0].name
|
||
self.output_names = [o.name for o in self.session.get_outputs()]
|
||
|
||
# 常见 ImageNet 归一化参数。若训练 ConvNeXt 时使用了其他 mean/std，需要保持一致。
|
||
self.mean = np.array([0.485, 0.456, 0.406], dtype=np.float32)
|
||
self.std = np.array([0.229, 0.224, 0.225], dtype=np.float32)
|
||
|
||
def preprocess(self, img_bgr: np.ndarray) -> np.ndarray:
|
||
img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
|
||
img_rgb = cv2.resize(img_rgb, (self.img_size, self.img_size), interpolation=cv2.INTER_LINEAR)
|
||
img = img_rgb.astype(np.float32) / 255.0
|
||
img = (img - self.mean) / self.std
|
||
img = img.transpose(2, 0, 1)
|
||
img = np.expand_dims(img, axis=0).astype(np.float32)
|
||
return img
|
||
|
||
def infer(self, img_bgr: np.ndarray) -> Dict:
|
||
input_tensor = self.preprocess(img_bgr)
|
||
outputs = self.session.run(self.output_names, {self.input_name: input_tensor})
|
||
logits = np.asarray(outputs[0])
|
||
|
||
if logits.ndim == 2:
|
||
logits = logits[0]
|
||
elif logits.ndim == 1:
|
||
pass
|
||
else:
|
||
logits = logits.reshape(-1)
|
||
|
||
probs = to_probabilities(logits)
|
||
class_id = int(np.argmax(probs))
|
||
confidence = float(probs[class_id])
|
||
class_name = self.class_names[class_id] if 0 <= class_id < len(self.class_names) else str(class_id)
|
||
|
||
return {
|
||
"cls_class_id": class_id,
|
||
"cls_class_name": class_name,
|
||
"cls_confidence": confidence,
|
||
}
|