initial commit

app/inference/convnext_onnx.py

@@ -0,0 +1,101 @@
+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,
+        }