245 lines
7.6 KiB
Python
245 lines
7.6 KiB
Python
import os
|
||
import time
|
||
from typing import Any, Dict, List, Optional, Tuple
|
||
|
||
import cv2
|
||
import numpy as np
|
||
import tensorrt as trt
|
||
import torch
|
||
from ultralytics import YOLO
|
||
from ultralytics.engine.results import Results
|
||
|
||
from config import get_config
|
||
|
||
|
||
class TensorRTEngine:
|
||
def __init__(self, engine_path: Optional[str] = None, device: int = 0):
|
||
config = get_config()
|
||
self.engine_path = engine_path or config.model.engine_path
|
||
self.device = device
|
||
self.imgsz = tuple(config.model.imgsz)
|
||
self.conf_thresh = config.model.conf_threshold
|
||
self.iou_thresh = config.model.iou_threshold
|
||
self.half = config.model.half
|
||
|
||
self.logger = trt.Logger(trt.Logger.INFO)
|
||
self.engine = None
|
||
self.context = None
|
||
self.stream = None
|
||
self.input_buffer = None
|
||
self.output_buffers = []
|
||
|
||
self._load_engine()
|
||
|
||
def _load_engine(self):
|
||
if not os.path.exists(self.engine_path):
|
||
raise FileNotFoundError(f"TensorRT引擎文件不存在: {self.engine_path}")
|
||
|
||
with open(self.engine_path, "rb") as f:
|
||
serialized_engine = f.read()
|
||
|
||
runtime = trt.Runtime(self.logger)
|
||
self.engine = runtime.deserialize_cuda_engine(serialized_engine)
|
||
|
||
self.context = self.engine.create_execution_context()
|
||
|
||
self.stream = torch.cuda.Stream(device=self.device)
|
||
|
||
for i in range(self.engine.num_io_tensors):
|
||
name = self.engine.get_tensor_name(i)
|
||
dtype = self.engine.get_tensor_dtype(name)
|
||
shape = self.engine.get_tensor_shape(name)
|
||
|
||
if self.engine.get_tensor_mode(name) == trt.TensorIOMode.INPUT:
|
||
self.context.set_tensor_address(name, None)
|
||
else:
|
||
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)
|
||
self.context.set_tensor_address(name, buffer.data_ptr())
|
||
|
||
self.context.set_optimization_profile_async(0, self.stream)
|
||
|
||
self.input_buffer = torch.zeros(
|
||
(1, 3, self.imgsz[0], self.imgsz[1]),
|
||
dtype=torch.float16 if self.half else torch.float32,
|
||
device=self.device,
|
||
)
|
||
|
||
def preprocess(self, frame: np.ndarray) -> torch.Tensor:
|
||
img = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
|
||
img = cv2.resize(img, self.imgsz)
|
||
|
||
img = img.transpose(2, 0, 1).astype(np.float32) / 255.0
|
||
|
||
if self.half:
|
||
img = img.astype(np.float16)
|
||
|
||
tensor = torch.from_numpy(img).unsqueeze(0).to(self.device)
|
||
|
||
return tensor
|
||
|
||
def inference(self, images: List[np.ndarray]) -> List[Results]:
|
||
batch_size = len(images)
|
||
if batch_size == 0:
|
||
return []
|
||
|
||
input_tensor = self.preprocess(images[0])
|
||
|
||
if batch_size > 1:
|
||
for i in range(1, batch_size):
|
||
input_tensor = torch.cat(
|
||
[input_tensor, self.preprocess(images[i])], dim=0
|
||
)
|
||
|
||
self.context.set_tensor_address(
|
||
"input", input_tensor.contiguous().data_ptr()
|
||
)
|
||
|
||
torch.cuda.synchronize(self.stream)
|
||
self.context.execute_async_v3(self.stream.handle)
|
||
torch.cuda.synchronize(self.stream)
|
||
|
||
results = []
|
||
for i in range(batch_size):
|
||
pred = self.output_buffers[0][i].cpu().numpy()
|
||
boxes = pred[:, :4]
|
||
scores = pred[:, 4]
|
||
classes = pred[:, 5].astype(np.int32)
|
||
|
||
mask = scores > self.conf_thresh
|
||
boxes = boxes[mask]
|
||
scores = scores[mask]
|
||
classes = classes[mask]
|
||
|
||
indices = cv2.dnn.NMSBoxes(
|
||
boxes.tolist(),
|
||
scores.tolist(),
|
||
self.conf_thresh,
|
||
self.iou_thresh,
|
||
)
|
||
|
||
if len(indices) > 0:
|
||
for idx in indices:
|
||
box = boxes[idx]
|
||
x1, y1, x2, y2 = box
|
||
w, h = x2 - x1, y2 - y1
|
||
conf = scores[idx]
|
||
cls = classes[idx]
|
||
|
||
orig_h, orig_w = images[i].shape[:2]
|
||
scale_x, scale_y = orig_w / self.imgsz[1], orig_h / self.imgsz[0]
|
||
box_orig = [
|
||
int(x1 * scale_x),
|
||
int(y1 * scale_y),
|
||
int(w * scale_x),
|
||
int(h * scale_y),
|
||
]
|
||
|
||
result = Results(
|
||
orig_img=images[i],
|
||
path="",
|
||
names={0: "person"},
|
||
boxes=Boxes(
|
||
torch.tensor([box_orig + [conf, cls]]),
|
||
orig_shape=(orig_h, orig_w),
|
||
),
|
||
)
|
||
results.append(result)
|
||
|
||
return results
|
||
|
||
def inference_single(self, frame: np.ndarray) -> List[Results]:
|
||
return self.inference([frame])
|
||
|
||
def warmup(self, num_warmup: int = 10):
|
||
dummy_frame = np.zeros((480, 640, 3), dtype=np.uint8)
|
||
for _ in range(num_warmup):
|
||
self.inference_single(dummy_frame)
|
||
|
||
def __del__(self):
|
||
if self.context:
|
||
self.context.synchronize()
|
||
if self.stream:
|
||
self.stream.synchronize()
|
||
|
||
|
||
class Boxes:
|
||
def __init__(
|
||
self,
|
||
data: torch.Tensor,
|
||
orig_shape: Tuple[int, int],
|
||
is_track: bool = False,
|
||
):
|
||
self.data = data
|
||
self.orig_shape = orig_shape
|
||
self.is_track = is_track
|
||
|
||
@property
|
||
def xyxy(self):
|
||
if self.is_track:
|
||
return self.data[:, :4]
|
||
return self.data[:, :4]
|
||
|
||
@property
|
||
def conf(self):
|
||
if self.is_track:
|
||
return self.data[:, 4]
|
||
return self.data[:, 4]
|
||
|
||
@property
|
||
def cls(self):
|
||
if self.is_track:
|
||
return self.data[:, 5]
|
||
return self.data[:, 5]
|
||
|
||
|
||
class YOLOEngine:
|
||
def __init__(
|
||
self,
|
||
model_path: Optional[str] = None,
|
||
device: int = 0,
|
||
use_trt: bool = True,
|
||
):
|
||
self.use_trt = use_trt
|
||
self.device = device
|
||
self.trt_engine = None
|
||
|
||
if not use_trt:
|
||
if model_path:
|
||
pt_path = model_path
|
||
elif hasattr(get_config().model, 'pt_model_path'):
|
||
pt_path = get_config().model.pt_model_path
|
||
else:
|
||
pt_path = get_config().model.engine_path.replace(".engine", ".pt")
|
||
self.model = YOLO(pt_path)
|
||
self.model.to(device)
|
||
else:
|
||
try:
|
||
self.trt_engine = TensorRTEngine(device=device)
|
||
self.trt_engine.warmup()
|
||
except Exception as e:
|
||
print(f"TensorRT加载失败，回退到PyTorch: {e}")
|
||
self.use_trt = False
|
||
if model_path:
|
||
pt_path = model_path
|
||
elif hasattr(get_config().model, 'pt_model_path'):
|
||
pt_path = get_config().model.pt_model_path
|
||
else:
|
||
pt_path = get_config().model.engine_path.replace(".engine", ".pt")
|
||
self.model = YOLO(pt_path)
|
||
self.model.to(device)
|
||
|
||
def __call__(self, frame: np.ndarray, **kwargs) -> List[Results]:
|
||
if self.use_trt:
|
||
return self.trt_engine.inference_single(frame)
|
||
else:
|
||
results = self.model(frame, imgsz=get_config().model.imgsz, **kwargs)
|
||
return results
|
||
|
||
def __del__(self):
|
||
if self.trt_engine:
|
||
del self.trt_engine
|