91 lines
3.0 KiB
Python
91 lines
3.0 KiB
Python
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
|