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

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: