diff --git a/config.py b/config.py index 5f4b965..420e4c8 100644 --- a/config.py +++ b/config.py @@ -23,16 +23,17 @@ class DatabaseConfig(BaseModel): class ModelConfig(BaseModel): - engine_path: str = "models/yolo11s.engine" - onnx_path: str = "models/yolo11s.onnx" - pt_model_path: str = "models/yolo11s.pt" + engine_path: str = "models/yolo11n.engine" + onnx_path: str = "models/yolo11n.onnx" + pt_model_path: str = "models/yolo11n.pt" imgsz: List[int] = [640, 640] conf_threshold: float = 0.5 iou_threshold: float = 0.45 device: int = 0 batch_size: int = 8 - half: bool = True - use_onnx: bool = True + half: bool = False + use_onnx: bool = False + use_trt: bool = False class StreamConfig(BaseModel): diff --git a/inference/engine.py b/inference/engine.py index e125f90..9026814 100644 --- a/inference/engine.py +++ b/inference/engine.py @@ -49,6 +49,10 @@ class ONNXEngine: return img def postprocess(self, output: np.ndarray, orig_img: np.ndarray) -> List[Results]: + import torch + import numpy as np + from ultralytics.engine.results import Boxes as BoxesObj, Results + c, n = output.shape output = output.T @@ -74,6 +78,9 @@ class ONNXEngine: orig_h, orig_w = orig_img.shape[:2] scale_x, scale_y = orig_w / self.imgsz[1], orig_h / self.imgsz[0] + if len(indices) == 0: + return [Results(orig_img=orig_img, path="", names={0: "person"})] + filtered_boxes = [] for idx in indices: if idx >= len(boxes): @@ -82,30 +89,30 @@ class ONNXEngine: x1, y1, x2, y2 = box w, h = x2 - x1, y2 - y1 filtered_boxes.append([ - int(x1 * scale_x), - int(y1 * scale_y), - int(w * scale_x), - int(h * scale_y), + float(x1 * scale_x), + float(y1 * scale_y), + float(w * scale_x), + float(h * scale_y), float(scores[idx]), int(classes[idx]) ]) - from ultralytics.engine.results import Boxes as BoxesObj if filtered_boxes: - box_tensor = torch.tensor(filtered_boxes) - boxes_obj = BoxesObj( - box_tensor, - orig_shape=(orig_h, orig_w) - ) - result = Results( - orig_img=orig_img, - path="", - names={0: "person"}, - boxes=boxes_obj - ) - return [result] + box_array = np.array(filtered_boxes, dtype=np.float32) + else: + box_array = np.zeros((0, 6), dtype=np.float32) - return [Results(orig_img=orig_img, path="", names={0: "person"})] + boxes_obj = BoxesObj( + torch.from_numpy(box_array), + orig_shape=(orig_h, orig_w) + ) + result = Results( + orig_img=orig_img, + path="", + names={0: "person"}, + boxes=boxes_obj + ) + return [result] def inference(self, images: List[np.ndarray]) -> List[Results]: if not images: @@ -183,29 +190,21 @@ class TensorRTEngine: self.context = self.engine.create_execution_context() self.stream = torch.cuda.Stream(device=self.device) - self.batch_size = 1 for i in range(self.engine.num_io_tensors): name = self.engine.get_tensor_name(i) dtype = self.engine.get_tensor_dtype(name) shape = list(self.engine.get_tensor_shape(name)) + if dtype == trt.float16: + buffer = torch.zeros(shape, dtype=torch.float16, device=self.device) + else: + buffer = torch.zeros(shape, dtype=torch.float32, device=self.device) + if self.engine.get_tensor_mode(name) == trt.TensorIOMode.INPUT: - if -1 in shape: - shape = [self.batch_size if d == -1 else d for d in shape] - if dtype == trt.float16: - buffer = torch.zeros(shape, dtype=torch.float16, device=self.device) - else: - buffer = torch.zeros(shape, dtype=torch.float32, device=self.device) self.input_buffer = buffer self.input_name = name else: - if -1 in shape: - shape = [self.batch_size if d == -1 else d for d in shape] - if dtype == trt.float16: - buffer = torch.zeros(shape, dtype=torch.float16, device=self.device) - else: - buffer = torch.zeros(shape, dtype=torch.float32, device=self.device) self.output_buffers.append(buffer) if self.output_name is None: self.output_name = name @@ -215,8 +214,6 @@ class TensorRTEngine: stream_handle = torch.cuda.current_stream(self.device).cuda_stream self.context.set_optimization_profile_async(0, stream_handle) - self.batch_size = 1 - def preprocess(self, frame: np.ndarray) -> torch.Tensor: img = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) img = cv2.resize(img, self.imgsz) @@ -247,9 +244,6 @@ class TensorRTEngine: self.input_name, input_tensor.contiguous().data_ptr() ) - input_shape = list(input_tensor.shape) - self.context.set_input_shape(self.input_name, input_shape) - torch.cuda.synchronize(self.stream) self.context.execute_async_v3(self.stream.cuda_stream) torch.cuda.synchronize(self.stream) @@ -336,6 +330,10 @@ class Boxes: self.orig_shape = orig_shape self.is_track = is_track + @property + def ndim(self) -> int: + return self.data.ndim + @property def xyxy(self): if self.is_track: @@ -369,35 +367,15 @@ class YOLOEngine: self, model_path: Optional[str] = None, device: int = 0, - use_trt: bool = True, + use_trt: bool = False, ): self.use_trt = False self.onnx_engine = None self.trt_engine = None + self.model = None self.device = device config = get_config() - - if use_trt: - try: - self.trt_engine = TensorRTEngine(device=device) - self.trt_engine.warmup() - self.use_trt = True - print("TensorRT引擎加载成功") - return - except Exception as e: - print(f"TensorRT加载失败: {e}") - - try: - onnx_path = config.model.onnx_path - if os.path.exists(onnx_path): - self.onnx_engine = ONNXEngine(device=device) - self.onnx_engine.warmup() - print("ONNX引擎加载成功") - return - else: - print(f"ONNX模型不存在: {onnx_path}") - except Exception as e: - print(f"ONNX加载失败: {e}") + self.config = config try: pt_path = model_path or config.model.pt_model_path @@ -409,26 +387,17 @@ class YOLOEngine: raise FileNotFoundError(f"PT文件无效或不存在: {pt_path}") except Exception as e: print(f"PyTorch加载失败: {e}") - raise RuntimeError("所有模型加载方式均失败") + raise RuntimeError("无法加载模型") def __call__(self, frame: np.ndarray, **kwargs) -> List[Results]: - if self.use_trt and self.trt_engine: + if self.model is not None: try: - return self.trt_engine.inference_single(frame) + return self.model(frame, imgsz=self.config.model.imgsz, conf=self.config.model.conf_threshold, iou=self.config.model.iou_threshold, **kwargs) except Exception as e: - print(f"TensorRT推理失败，切换到ONNX: {e}") - self.use_trt = False - if self.onnx_engine: - return self.onnx_engine.inference_single(frame) - elif self.model: - return self.model(frame, imgsz=get_config().model.imgsz, **kwargs) - else: - return [] - elif self.onnx_engine: - return self.onnx_engine.inference_single(frame) - else: - results = self.model(frame, imgsz=get_config().model.imgsz, **kwargs) - return results + print(f"PyTorch推理失败: {e}") + + print("警告: 模型不可用，返回空结果") + return [] def __del__(self): if self.trt_engine: diff --git a/inference/pipeline.py b/inference/pipeline.py index 3ada530..b671993 100644 --- a/inference/pipeline.py +++ b/inference/pipeline.py @@ -7,6 +7,7 @@ from collections import deque from datetime import datetime from typing import Any, Dict, List, Optional, Tuple +import cv2 import numpy as np from config import get_config diff --git a/logs/app.log b/logs/app.log index 985d340..e4b27b7 100644 --- a/logs/app.log +++ b/logs/app.log @@ -133,3 +133,88 @@ 2026-01-21 13:18:55,795 - security_monitor - INFO - 启动安保异常行为识别系统 2026-01-21 13:18:55,809 - security_monitor - INFO - 数据库初始化完成 2026-01-21 13:19:08,492 - security_monitor - INFO - 推理Pipeline启动，活跃摄像头数: 2 +2026-01-21 14:01:21,015 - security_monitor - INFO - 正在关闭系统... +2026-01-21 14:01:21,257 - security_monitor - INFO - 系统已关闭 +2026-01-21 14:03:48,547 - security_monitor - INFO - 启动安保异常行为识别系统 +2026-01-21 14:03:48,563 - security_monitor - INFO - 数据库初始化完成 +2026-01-21 14:04:01,197 - security_monitor - INFO - 推理Pipeline启动，活跃摄像头数: 2 +2026-01-21 14:04:20,191 - security_monitor - INFO - 正在关闭系统... +2026-01-21 14:04:20,414 - security_monitor - INFO - 系统已关闭 +2026-01-21 14:05:48,342 - security_monitor - INFO - 启动安保异常行为识别系统 +2026-01-21 14:05:48,355 - security_monitor - INFO - 数据库初始化完成 +2026-01-21 14:06:00,984 - security_monitor - INFO - 推理Pipeline启动，活跃摄像头数: 2 +2026-01-21 14:07:24,065 - security_monitor - INFO - 正在关闭系统... +2026-01-21 14:07:24,222 - security_monitor - INFO - 系统已关闭 +2026-01-21 14:08:10,073 - security_monitor - INFO - 启动安保异常行为识别系统 +2026-01-21 14:08:10,088 - security_monitor - INFO - 数据库初始化完成 +2026-01-21 14:08:22,715 - security_monitor - INFO - 推理Pipeline启动，活跃摄像头数: 2 +2026-01-21 14:09:05,249 - security_monitor - INFO - 正在关闭系统... +2026-01-21 14:09:05,480 - security_monitor - INFO - 系统已关闭 +2026-01-21 14:11:29,491 - security_monitor - INFO - 启动安保异常行为识别系统 +2026-01-21 14:11:29,513 - security_monitor - INFO - 数据库初始化完成 +2026-01-21 14:11:42,900 - security_monitor - INFO - 推理Pipeline启动，活跃摄像头数: 2 +2026-01-21 14:14:04,974 - security_monitor - INFO - 正在关闭系统... +2026-01-21 14:14:05,161 - security_monitor - INFO - 系统已关闭 +2026-01-21 14:14:41,203 - security_monitor - INFO - 启动安保异常行为识别系统 +2026-01-21 14:14:41,220 - security_monitor - INFO - 数据库初始化完成 +2026-01-21 14:14:54,380 - security_monitor - INFO - 推理Pipeline启动，活跃摄像头数: 2 +2026-01-21 14:15:30,975 - security_monitor - INFO - 正在关闭系统... +2026-01-21 14:15:31,180 - security_monitor - INFO - 系统已关闭 +2026-01-21 14:16:24,472 - security_monitor - INFO - 启动安保异常行为识别系统 +2026-01-21 14:16:24,485 - security_monitor - INFO - 数据库初始化完成 +2026-01-21 14:16:37,611 - security_monitor - INFO - 推理Pipeline启动，活跃摄像头数: 2 +2026-01-21 14:17:01,178 - security_monitor - INFO - 正在关闭系统... +2026-01-21 14:17:01,420 - security_monitor - INFO - 系统已关闭 +2026-01-21 14:18:00,008 - security_monitor - INFO - 启动安保异常行为识别系统 +2026-01-21 14:18:00,022 - security_monitor - INFO - 数据库初始化完成 +2026-01-21 14:18:13,126 - security_monitor - INFO - 推理Pipeline启动，活跃摄像头数: 2 +2026-01-21 14:18:13,128 - security_monitor - INFO - 正在关闭系统... +2026-01-21 14:18:21,683 - security_monitor - INFO - 系统已关闭 +2026-01-21 14:20:04,985 - security_monitor - INFO - 启动安保异常行为识别系统 +2026-01-21 14:20:04,999 - security_monitor - INFO - 数据库初始化完成 +2026-01-21 14:20:18,151 - security_monitor - INFO - 推理Pipeline启动，活跃摄像头数: 2 +2026-01-21 14:21:24,782 - security_monitor - INFO - 正在关闭系统... +2026-01-21 14:21:24,927 - security_monitor - INFO - 系统已关闭 +2026-01-21 14:22:48,064 - security_monitor - INFO - 启动安保异常行为识别系统 +2026-01-21 14:22:48,078 - security_monitor - INFO - 数据库初始化完成 +2026-01-21 14:23:01,270 - security_monitor - INFO - 推理Pipeline启动，活跃摄像头数: 2 +2026-01-21 14:23:13,509 - security_monitor - INFO - 正在关闭系统... +2026-01-21 14:23:13,628 - security_monitor - INFO - 系统已关闭 +2026-01-21 14:24:16,374 - security_monitor - INFO - 启动安保异常行为识别系统 +2026-01-21 14:24:16,386 - security_monitor - INFO - 数据库初始化完成 +2026-01-21 14:24:29,425 - security_monitor - INFO - 推理Pipeline启动，活跃摄像头数: 2 +2026-01-21 14:24:42,751 - security_monitor - INFO - 正在关闭系统... +2026-01-21 14:24:42,846 - security_monitor - INFO - 系统已关闭 +2026-01-21 14:25:25,549 - security_monitor - INFO - 启动安保异常行为识别系统 +2026-01-21 14:25:25,562 - security_monitor - INFO - 数据库初始化完成 +2026-01-21 14:25:38,636 - security_monitor - INFO - 推理Pipeline启动，活跃摄像头数: 2 +2026-01-21 14:26:02,871 - security_monitor - INFO - 正在关闭系统... +2026-01-21 14:26:03,124 - security_monitor - INFO - 系统已关闭 +2026-01-21 14:26:45,885 - security_monitor - INFO - 启动安保异常行为识别系统 +2026-01-21 14:26:45,899 - security_monitor - INFO - 数据库初始化完成 +2026-01-21 14:26:59,042 - security_monitor - INFO - 推理Pipeline启动，活跃摄像头数: 2 +2026-01-21 14:27:26,873 - security_monitor - INFO - 正在关闭系统... +2026-01-21 14:27:26,980 - security_monitor - INFO - 系统已关闭 +2026-01-21 14:31:38,376 - security_monitor - INFO - 启动安保异常行为识别系统 +2026-01-21 14:31:38,390 - security_monitor - INFO - 数据库初始化完成 +2026-01-21 14:31:51,594 - security_monitor - INFO - 推理Pipeline启动，活跃摄像头数: 2 +2026-01-21 14:32:17,471 - security_monitor - INFO - 正在关闭系统... +2026-01-21 14:32:17,536 - security_monitor - INFO - 系统已关闭 +2026-01-21 14:32:53,841 - security_monitor - INFO - 启动安保异常行为识别系统 +2026-01-21 14:32:53,855 - security_monitor - INFO - 数据库初始化完成 +2026-01-21 14:33:06,946 - security_monitor - INFO - 推理Pipeline启动，活跃摄像头数: 2 +2026-01-21 14:34:30,645 - security_monitor - INFO - 正在关闭系统... +2026-01-21 14:34:30,818 - security_monitor - INFO - 系统已关闭 +2026-01-21 14:38:24,673 - security_monitor - INFO - 启动安保异常行为识别系统 +2026-01-21 14:38:24,685 - security_monitor - INFO - 数据库初始化完成 +2026-01-21 14:38:37,183 - security_monitor - INFO - 推理Pipeline启动，活跃摄像头数: 2 +2026-01-21 14:39:04,359 - security_monitor - INFO - 正在关闭系统... +2026-01-21 14:39:04,486 - security_monitor - INFO - 系统已关闭 +2026-01-21 14:40:07,246 - security_monitor - INFO - 启动安保异常行为识别系统 +2026-01-21 14:40:07,259 - security_monitor - INFO - 数据库初始化完成 +2026-01-21 14:40:19,745 - security_monitor - INFO - 推理Pipeline启动，活跃摄像头数: 2 +2026-01-21 14:40:33,742 - security_monitor - INFO - 正在关闭系统... +2026-01-21 14:40:33,863 - security_monitor - INFO - 系统已关闭 +2026-01-21 14:41:27,191 - security_monitor - INFO - 启动安保异常行为识别系统 +2026-01-21 14:41:27,205 - security_monitor - INFO - 数据库初始化完成 +2026-01-21 14:41:39,701 - security_monitor - INFO - 推理Pipeline启动，活跃摄像头数: 2 diff --git a/main.py b/main.py index 378e680..6f23400 100644 --- a/main.py +++ b/main.py @@ -18,6 +18,16 @@ from prometheus_client import start_http_server sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +from ultralytics.engine.results import Boxes as UltralyticsBoxes + +def _patch_boxes_ndim(): + if not hasattr(UltralyticsBoxes, 'ndim'): + @property + def ndim(self): + return self.data.ndim + UltralyticsBoxes.ndim = ndim +_patch_boxes_ndim() + from api.alarm import router as alarm_router from api.camera import router as camera_router from api.roi import router as roi_router