|
|
@@ -523,18 +523,17 @@ class ObjectDetector:
|
|
|
raise ImportError("未安装 rknnlite,请运行: pip install rknnlite2")
|
|
|
|
|
|
def _load_onnx_model(self):
|
|
|
- """加载 ONNX 模型"""
|
|
|
+ """加载 ONNX 模型(使用 Ultralytics 推理,保证后处理正确)"""
|
|
|
if not self.model_path:
|
|
|
raise ValueError("ONNX 模型需要指定 model_path")
|
|
|
-
|
|
|
+
|
|
|
try:
|
|
|
- import onnxruntime as ort
|
|
|
- self.session = ort.InferenceSession(self.model_path)
|
|
|
- self.input_name = self.session.get_inputs()[0].name
|
|
|
- self.output_names = [o.name for o in self.session.get_outputs()]
|
|
|
- print(f"ONNX 模型加载成功: {self.model_path}, 输出数量={len(self.output_names)}")
|
|
|
+ from ultralytics import YOLO
|
|
|
+ self.model = YOLO(self.model_path, task='detect')
|
|
|
+ self.device = 'cpu' # RK3588 无 CUDA,强制 CPU
|
|
|
+ print(f"ONNX 模型加载成功: {self.model_path} (via Ultralytics, device=cpu)")
|
|
|
except ImportError:
|
|
|
- raise ImportError("未安装 onnxruntime,请运行: pip install onnxruntime")
|
|
|
+ raise ImportError("未安装 ultralytics")
|
|
|
|
|
|
def _load_yolo_model(self):
|
|
|
"""加载YOLO11检测模型"""
|
|
|
@@ -929,7 +928,7 @@ class ObjectDetector:
|
|
|
|
|
|
class_probs = output[4:, i]
|
|
|
best_class = int(np.argmax(class_probs))
|
|
|
- confidence = float(class_probs[best_class])
|
|
|
+ confidence = 1.0 / (1.0 + np.exp(-float(class_probs[best_class])))
|
|
|
|
|
|
if confidence < conf_threshold:
|
|
|
continue
|
|
|
@@ -972,6 +971,7 @@ class ObjectDetector:
|
|
|
"""分区域检测 - 将宽幅全景图分成多个重叠区域分别检测,提高远处目标识别率"""
|
|
|
results = []
|
|
|
h0, w0 = frame.shape[:2]
|
|
|
+ need_nms = not self.is_end2end
|
|
|
conf_threshold = self.config['confidence_threshold']
|
|
|
class_map = self.config.get('class_map', {0: 'person'})
|
|
|
|
|
|
@@ -1038,7 +1038,7 @@ class ObjectDetector:
|
|
|
for i in range(output.shape[1]):
|
|
|
class_probs = output[4:, i]
|
|
|
best_class = int(np.argmax(class_probs))
|
|
|
- confidence = float(class_probs[best_class])
|
|
|
+ confidence = 1.0 / (1.0 + np.exp(-float(class_probs[best_class])))
|
|
|
if confidence < conf_threshold:
|
|
|
continue
|
|
|
x1 = output[0, i] - output[2, i] / 2
|
|
|
@@ -1086,6 +1086,9 @@ class ObjectDetector:
|
|
|
)
|
|
|
results.append(obj)
|
|
|
|
|
|
+ if need_nms and results:
|
|
|
+ results = nms(results, iou_threshold=0.45)
|
|
|
+
|
|
|
return results
|
|
|
|
|
|
def detect(self, frame: np.ndarray) -> List[DetectedObject]:
|