XW-AIOT/Security_AI_integrated
4
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(inference): resolve multiple YOLO inference and API issues

Browse Source
This commit is contained in:
16337
2026-01-21 14:48:01 +08:00
parent 1b344aeb2e
commit 1c7190bbb0
5 changed files with 146 additions and 80 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
config.py
View File

@@ -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):

119
inference/engine.py
View File

@@ -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:

1
inference/pipeline.py
View File

@@ -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

85
logs/app.log
View File

@@ -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

10
main.py
View File

@@ -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