initial commit

app/inference/density_onnx.py

@@ -0,0 +1,90 @@
+import cv2
+import numpy as np
+import onnxruntime as ort
+
+
+class DensityEstimatorONNX:
+    def __init__(self, model_path, input_size=None, device="cpu", providers=None):
+        self.model_path = model_path
+        self.input_size = input_size
+
+        if providers is None:
+            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("[Density] providers:", active_providers)
+        if (
+            device.lower() in ["cuda", "gpu"]
+            and "CUDAExecutionProvider" not in active_providers
+        ):
+            raise RuntimeError(
+                "请求使用 CUDA 推理，但密度模型会话未启用 "
+                "CUDAExecutionProvider。请安装匹配 CUDA/cuDNN 的 "
+                "onnxruntime-gpu。"
+            )
+
+        self.input_name = self.session.get_inputs()[0].name
+        self.input_shape = self.session.get_inputs()[0].shape
+        self.output_names = [o.name for o in self.session.get_outputs()]
+
+        print("[Density] input_name:", self.input_name)
+        print("[Density] input_shape:", self.input_shape)
+        print("[Density] output_names:", self.output_names)
+
+        if self.input_size is None:
+            h = self.input_shape[2]
+            w = self.input_shape[3]
+            if isinstance(h, int) and isinstance(w, int):
+                self.input_size = (w, h)
+            else:
+                self.input_size = (400, 400)
+
+    def preprocess(self, image_bgr):
+        original_h, original_w = image_bgr.shape[:2]
+
+        image_rgb = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)
+        resized = cv2.resize(image_rgb, self.input_size)
+
+        img = resized.astype(np.float32) / 255.0
+
+        mean = np.array([0.485, 0.456, 0.406], dtype=np.float32)
+        std = np.array([0.229, 0.224, 0.225], dtype=np.float32)
+        img = (img - mean) / std
+
+        img = img.transpose(2, 0, 1)
+        img = np.expand_dims(img, axis=0).astype(np.float32)
+
+        return img, original_w, original_h
+
+    def infer(self, image_bgr):
+        input_tensor, original_w, original_h = self.preprocess(image_bgr)
+
+        outputs = self.session.run(self.output_names, {self.input_name: input_tensor})
+        density = outputs[0]
+
+        density = np.squeeze(density).astype(np.float32)
+
+        if density.ndim != 2:
+            density = density.reshape(density.shape[-2], density.shape[-1])
+
+        count = float(density.sum())
+
+        density_map = cv2.resize(density, (original_w, original_h))
+
+        return {
+            "count": count,
+            "density_map": density_map
+        }
+
+    @staticmethod
+    def colorize_density_map(density_map):
+        d = density_map.copy()
+        d = d - d.min()
+        if d.max() > 0:
+            d = d / d.max()
+        d = (d * 255).astype(np.uint8)
+        heatmap = cv2.applyColorMap(d, cv2.COLORMAP_JET)
+        return heatmap