initial commit
This commit is contained in:
Ankle committed
commit
0bafbee31f
54 files changed
+6317
No files matched your search
Whitespace-only changes.
@@ -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,
|
||||
}
|
||||
@@ -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
|
||||
@@ -0,0 +1,216 @@
|
||||
from typing import List, Dict
|
||||
import cv2
|
||||
import numpy as np
|
||||
import onnxruntime as ort
|
||||
|
||||
|
||||
def letterbox(
|
||||
img: np.ndarray,
|
||||
new_shape=(960, 960),
|
||||
color=(114, 114, 114),
|
||||
auto=False,
|
||||
scale_fill=False,
|
||||
scale_up=True,
|
||||
stride=32,
|
||||
):
|
||||
shape = img.shape[:2] # h, w
|
||||
|
||||
if isinstance(new_shape, int):
|
||||
new_shape = (new_shape, new_shape)
|
||||
|
||||
r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])
|
||||
if not scale_up:
|
||||
r = min(r, 1.0)
|
||||
|
||||
ratio = (r, r)
|
||||
new_unpad = (int(round(shape[1] * r)), int(round(shape[0] * r)))
|
||||
|
||||
dw = new_shape[1] - new_unpad[0]
|
||||
dh = new_shape[0] - new_unpad[1]
|
||||
|
||||
if auto:
|
||||
dw, dh = np.mod(dw, stride), np.mod(dh, stride)
|
||||
elif scale_fill:
|
||||
dw, dh = 0.0, 0.0
|
||||
new_unpad = (new_shape[1], new_shape[0])
|
||||
ratio = (new_shape[1] / shape[1], new_shape[0] / shape[0])
|
||||
|
||||
dw /= 2
|
||||
dh /= 2
|
||||
|
||||
if shape[::-1] != new_unpad:
|
||||
img = cv2.resize(img, new_unpad, interpolation=cv2.INTER_LINEAR)
|
||||
|
||||
top = int(round(dh - 0.1))
|
||||
bottom = int(round(dh + 0.1))
|
||||
left = int(round(dw - 0.1))
|
||||
right = int(round(dw + 0.1))
|
||||
|
||||
img = cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color)
|
||||
return img, ratio, (dw, dh)
|
||||
|
||||
|
||||
def xywh2xyxy(x: np.ndarray) -> np.ndarray:
|
||||
y = np.copy(x)
|
||||
y[:, 0] = x[:, 0] - x[:, 2] / 2
|
||||
y[:, 1] = x[:, 1] - x[:, 3] / 2
|
||||
y[:, 2] = x[:, 0] + x[:, 2] / 2
|
||||
y[:, 3] = x[:, 1] + x[:, 3] / 2
|
||||
return y
|
||||
|
||||
|
||||
def box_iou(box: np.ndarray, boxes: np.ndarray) -> np.ndarray:
|
||||
x1 = np.maximum(box[0], boxes[:, 0])
|
||||
y1 = np.maximum(box[1], boxes[:, 1])
|
||||
x2 = np.minimum(box[2], boxes[:, 2])
|
||||
y2 = np.minimum(box[3], boxes[:, 3])
|
||||
|
||||
inter = np.maximum(0, x2 - x1) * np.maximum(0, y2 - y1)
|
||||
area1 = max(0, box[2] - box[0]) * max(0, box[3] - box[1])
|
||||
area2 = np.maximum(0, boxes[:, 2] - boxes[:, 0]) * np.maximum(0, boxes[:, 3] - boxes[:, 1])
|
||||
union = area1 + area2 - inter + 1e-7
|
||||
return inter / union
|
||||
|
||||
|
||||
def nms_numpy(boxes: np.ndarray, scores: np.ndarray, iou_thres: float) -> List[int]:
|
||||
if len(boxes) == 0:
|
||||
return []
|
||||
|
||||
order = scores.argsort()[::-1]
|
||||
keep = []
|
||||
|
||||
while order.size > 0:
|
||||
i = order[0]
|
||||
keep.append(int(i))
|
||||
if order.size == 1:
|
||||
break
|
||||
ious = box_iou(boxes[i], boxes[order[1:]])
|
||||
inds = np.where(ious <= iou_thres)[0]
|
||||
order = order[inds + 1]
|
||||
|
||||
return keep
|
||||
|
||||
|
||||
def scale_boxes_from_letterbox(boxes: np.ndarray, ratio, pad, original_shape) -> np.ndarray:
|
||||
# boxes: xyxy in letterboxed image
|
||||
gain = ratio[0]
|
||||
pad_x, pad_y = pad
|
||||
|
||||
boxes[:, [0, 2]] -= pad_x
|
||||
boxes[:, [1, 3]] -= pad_y
|
||||
boxes[:, :4] /= gain
|
||||
|
||||
h, w = original_shape[:2]
|
||||
boxes[:, [0, 2]] = boxes[:, [0, 2]].clip(0, w - 1)
|
||||
boxes[:, [1, 3]] = boxes[:, [1, 3]].clip(0, h - 1)
|
||||
return boxes
|
||||
|
||||
|
||||
class YOLOv5ONNXDetector:
|
||||
def __init__(
|
||||
self,
|
||||
model_path: str,
|
||||
class_names: List[str],
|
||||
img_size: int = 960,
|
||||
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("[YOLO] providers:", active_providers)
|
||||
if (
|
||||
device.lower() in ["cuda", "gpu"]
|
||||
and "CUDAExecutionProvider" not in active_providers
|
||||
):
|
||||
raise RuntimeError(
|
||||
"请求使用 CUDA 推理，但 YOLO 会话未启用 "
|
||||
"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()]
|
||||
|
||||
def preprocess(self, img_bgr: np.ndarray):
|
||||
img_lb, ratio, pad = letterbox(img_bgr, new_shape=(self.img_size, self.img_size), auto=False)
|
||||
img_rgb = cv2.cvtColor(img_lb, cv2.COLOR_BGR2RGB)
|
||||
img = img_rgb.transpose(2, 0, 1).astype(np.float32) / 255.0
|
||||
img = np.expand_dims(img, axis=0)
|
||||
return img, ratio, pad
|
||||
|
||||
def infer(
|
||||
self,
|
||||
img_bgr: np.ndarray,
|
||||
conf_thres: float = 0.25,
|
||||
iou_thres: float = 0.45,
|
||||
max_det: int = 50,
|
||||
) -> List[Dict]:
|
||||
input_tensor, ratio, pad = self.preprocess(img_bgr)
|
||||
outputs = self.session.run(self.output_names, {self.input_name: input_tensor})
|
||||
|
||||
pred = outputs[0]
|
||||
|
||||
# 常见 YOLOv5 ONNX 输出: [1, num_boxes, 5 + num_classes]
|
||||
if isinstance(pred, list):
|
||||
pred = pred[0]
|
||||
pred = np.asarray(pred)
|
||||
|
||||
if pred.ndim == 3:
|
||||
pred = pred[0]
|
||||
elif pred.ndim == 2:
|
||||
pass
|
||||
else:
|
||||
raise ValueError(f"不支持的 YOLO 输出维度: {pred.shape}")
|
||||
|
||||
if pred.shape[1] < 6:
|
||||
raise ValueError(f"YOLO 输出最后一维应至少为 6，当前为: {pred.shape}")
|
||||
|
||||
boxes_xywh = pred[:, :4]
|
||||
obj_conf = pred[:, 4]
|
||||
|
||||
if pred.shape[1] == 6:
|
||||
# 某些导出格式可能已经是 x,y,w,h,conf,class_id
|
||||
cls_ids = pred[:, 5].astype(np.int32)
|
||||
scores = obj_conf
|
||||
else:
|
||||
cls_scores = pred[:, 5:]
|
||||
cls_ids = np.argmax(cls_scores, axis=1)
|
||||
cls_conf = cls_scores[np.arange(len(cls_scores)), cls_ids]
|
||||
scores = obj_conf * cls_conf
|
||||
|
||||
mask = scores >= conf_thres
|
||||
if not np.any(mask):
|
||||
return []
|
||||
|
||||
boxes_xywh = boxes_xywh[mask]
|
||||
scores = scores[mask]
|
||||
cls_ids = cls_ids[mask]
|
||||
|
||||
boxes = xywh2xyxy(boxes_xywh)
|
||||
boxes = scale_boxes_from_letterbox(boxes, ratio, pad, img_bgr.shape)
|
||||
|
||||
final_results = []
|
||||
|
||||
for class_id in np.unique(cls_ids):
|
||||
idxs = np.where(cls_ids == class_id)[0]
|
||||
keep = nms_numpy(boxes[idxs], scores[idxs], iou_thres)
|
||||
|
||||
for k in keep:
|
||||
original_idx = idxs[k]
|
||||
cid = int(class_id)
|
||||
cname = self.class_names[cid] if 0 <= cid < len(self.class_names) else str(cid)
|
||||
final_results.append({
|
||||
"bbox_xyxy": boxes[original_idx].astype(float).tolist(),
|
||||
"det_confidence": float(scores[original_idx]),
|
||||
"det_class_id": cid,
|
||||
"det_class_name": cname,
|
||||
})
|
||||
|
||||
final_results.sort(key=lambda x: x["det_confidence"], reverse=True)
|
||||
return final_results[:max_det]
|
||||
Reference in new issue
Block a user