TensorRT推理
介绍
TensorRT是NVIDIA 推出的高性能深度学习推理优化器（Optimizer）和运行时（Runtime）库，用于在 NVIDIA GPU 上对训练好的模型做加速推理
使用条件：
- ✅ 部署硬件是 NVIDIA GPU
- ✅ 模型已定型，即将上线（不再频繁改动结构）
- ✅ 对延迟或吞吐有要求（实时检测、多路视频、高并发服务）
- ✅ 模型由常规算子组成（CNN / Transformer 等主流结构，TensorRT 支持好）
- ✅ 输入尺寸固定或变化有规律（TensorRT 支持动态 shape，但静态 shape 优化最充分）
深度学习网络往往由几百上千个层组成。每个层都是一个独立的 CUDA kernel（核函数），每执行一个kernel 都要：
启动 kernel + 把数据从显存读出来 + 算完再写回显存。
显存读写（访存）才是瓶颈，而不是计算本身
TensorRT 把相邻的、可以合并的层"焊"在一起，一个 kernel 算完：
例如
优化前：Conv → BN → ReLU （3 个 kernel，读写显存 3 次）
优化后：Conv+BN+ReLU 融合为 CBR （1 个 kernel，只读写 1 次）
模型训练时用 FP32 保证精度；
推理时把计算精度降下来，推理速度大幅提升
yolov8n.onnx 12.8 MB ← ONNX 原始模型
yolov8n_fp32.engine 18.8 MB
yolov8n_fp16.engine 18.9 MB
yolov8n_int8.engine 4.6 MB ← INT8 体积约为 FP32 的 1/4
yolov8推理代码
yolov8-tensort.py（环境为tensorRT的环境）
"""
YOLOv8 TensorRT 推理脚本
环境:
# 1. 建 Python 3.10 环境（TRT 8.5 只支持到 cp310）
conda create -n trt85 python=3.10 -y
conda activate trt85
# 2. 装 TRT 8.5.1.7 自带 wheel + 依赖
pip install "D:\PycharmProjects\20260819\TensorRT-8.5.1.7\python\tensorrt-8.5.1.7-cp310-none-win_amd64.whl"
pip install "numpy<2" opencv-python
用法: python yolov8-tensort.py
"""
import ctypes
import os
import cv2
import numpy as np
import tensorrt as trt
# 消除 TRT 的 "CUDA lazy loading is not enabled" 警告（同时减少显存占用）
os.environ.setdefault("CUDA_MODULE_LOADING", "LAZY")
# ==================== CUDA Runtime Helper (ctypes) ====================
class _CudaRuntime:
"""
通过 ctypes 调用 CUDA Runtime API — 同步拷贝，不用流(stream)
为什么需要这个类？
TensorRT 的 Python 包只负责"推理计算"，不负责"显存管理"。
GPU 显存的分配/释放/拷贝必须通过 CUDA Runtime 完成，
这里用 ctypes 直接加载 cudart64_110.dll（CUDA 11.8 的运行时库），
从而免去安装 pycuda。
两个容易踩坑的关键点：
1. 每个函数都要声明 argtypes / restype：
ctypes 默认把参数当 32 位 int 传递，64 位指针会被截断导致崩溃。
2. _check 检查返回值：
CUDA API 出错时不抛异常、只返回错误码，必须自己检查。
"""
H2D = 1 # cudaMemcpyHostToDevice: CPU → GPU
D2H = 2 # cudaMemcpyDeviceToHost: GPU → CPU
def __init__(self):
# 加载 CUDA 运行时库：优先从 CUDA_PATH 找，其次依赖 PATH
cuda_bin = os.path.join(os.environ.get('CUDA_PATH', ''), 'bin')
path = os.path.join(cuda_bin, 'cudart64_110.dll')
self._dll = ctypes.CDLL(path if os.path.exists(path) else 'cudart64_110.dll')
# 函数签名声明（必须！原因见类说明第 1 点）
self._dll.cudaMalloc.restype = ctypes.c_int
self._dll.cudaMalloc.argtypes = [ctypes.POINTER(ctypes.c_void_p), ctypes.c_size_t]
self._dll.cudaFree.argtypes = [ctypes.c_void_p]
self._dll.cudaMemcpy.restype = ctypes.c_int
self._dll.cudaMemcpy.argtypes = [ctypes.c_void_p, ctypes.c_void_p,
ctypes.c_size_t, ctypes.c_int]
def _check(self, err, msg):
"""CUDA API 返回非 0 就是出错，抛出带说明的异常"""
if err != 0:
raise RuntimeError(f'{msg}: CUDA error {err}')
def mem_alloc(self, nbytes):
"""在 GPU 上分配 nbytes 字节显存，返回显存地址（裸指针 int）"""
ptr = ctypes.c_void_p()
self._check(self._dll.cudaMalloc(ctypes.byref(ptr), nbytes), 'cudaMalloc')
return ptr.value
def mem_free(self, ptr):
"""释放 GPU 显存"""
self._dll.cudaFree(ctypes.c_void_p(ptr))
def memcpy_htod(self, dst, src_np):
"""把 numpy 数组拷到 GPU (Host → Device)，返回时拷贝已完成"""
self._check(self._dll.cudaMemcpy(ctypes.c_void_p(dst), src_np.ctypes.data,
src_np.nbytes, self.H2D), 'cudaMemcpy H2D')
def memcpy_dtoh(self, dst_np, src):
"""把 GPU 数据拷回 numpy 数组 (Device → Host)，返回时拷贝已完成"""
self._check(self._dll.cudaMemcpy(dst_np.ctypes.data, ctypes.c_void_p(src),
dst_np.nbytes, self.D2H), 'cudaMemcpy D2H')
cuda = _CudaRuntime()
# ============ 配置 ============
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
# 三个引擎任选（都经过 check_engine.py 体检）：
# yolov8n_fp32.engine / yolov8n_fp16.engine —— trtexec 直接转
# yolov8n_int8_v2.engine —— build_int8.py 正确校准后构建（分数≈FP32）
# 旧的 yolov8n_int8.engine 校准坏了（分数被压到 0.001），不要用
ENGINE_PATH = os.path.join(BASE_DIR, "yolov8n_fp32.engine")
IMAGE_PATH = os.path.join(BASE_DIR, "dog.png")
CONF_THRES = 0.5
IOU_THRES = 0.45
# COCO 80 类
NAMES = [
'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat',
'traffic light', 'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat',
'dog', 'horse', 'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe', 'backpack',
'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee', 'skis', 'snowboard', 'sports ball',
'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard', 'tennis racket',
'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple',
'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair',
'couch', 'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse',
'remote', 'keyboard', 'cell phone', 'microwave', 'oven', 'toaster', 'sink',
'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear', 'hair drier',
'toothbrush'
]
TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
import time
class YOLOv8EngineDetector:
def __init__(self, engine_path, conf_thres=0.5, iou_thres=0.45):
self.conf_thres = conf_thres
self.iou_thres = iou_thres
self.engine = self._load_engine(engine_path)
# ---- 1. 加载 engine ----
def _load_engine(self, engine_path):
with open(engine_path, "rb") as f, trt.Runtime(TRT_LOGGER) as runtime:
return runtime.deserialize_cuda_engine(f.read())
# ---- 2. 预处理 ----
def _preprocess(self, img, size=640):
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
h, w = img.shape[:2]
s = max(h, w)
bg = np.zeros((s, s, 3), np.uint8) # 正方化填充
bg[:h, :w] = img
blob = np.expand_dims(cv2.resize(bg, (size, size))
.transpose(2, 0, 1).astype(np.float32) / 255.0, 0)
return blob, bg
# ---- 3. GPU 推理 ----
def _infer(self, input_data):
context = self.engine.create_execution_context()
input_name = "images"
output_name = "output0"
dtype = np.float32 # FP32 engine：输入输出都是 float32
if self.engine.num_optimization_profiles > 0: # 动态 shape 引擎才需要
context.set_input_shape(input_name, input_data.shape)
output_shape = tuple(context.get_tensor_shape(output_name))
h_input = np.ascontiguousarray(input_data.ravel(), dtype=dtype)
h_output = np.empty(output_shape, np.float32)
d_input = cuda.mem_alloc(h_input.nbytes)
d_output = cuda.mem_alloc(h_output.nbytes)
context.set_tensor_address(input_name, d_input)
context.set_tensor_address(output_name, d_output)
# CPU → GPU → 推理 → GPU → CPU（0 = 默认流，按顺序执行）
cuda.memcpy_htod(d_input, h_input)
context.execute_async_v3(0)
cuda.memcpy_dtoh(h_output, d_output)
cuda.mem_free(d_input)
cuda.mem_free(d_output)
return h_output.reshape(output_shape)
# ---- 4. 后处理 (NMS) ----
def _postprocess(self, output, bg):
ratio = bg.shape[0] / 640
detections = np.transpose(output[0], (1, 0)) # (1,84,8400) → (8400,84)
boxes, confidences, cls_ids = [], [], []
for cx, cy, w, h, *scores in detections:
cls_id = np.argmax(scores)
score = scores[cls_id]
if score > self.conf_thres:
# NMSBoxes 要求 [x, y, w, h]（左上角 + 宽高），先存宽高
boxes.append([int((cx - w / 2) * ratio), int((cy - h / 2) * ratio),
int(w * ratio), int(h * ratio)])
confidences.append(float(score))
cls_ids.append(int(cls_id))
if not boxes:
return np.array([])
indices = cv2.dnn.NMSBoxes(boxes, confidences, self.conf_thres, self.iou_thres)
if len(indices) == 0:
return np.array([])
# NMS 之后转成 [xmin, ymin, xmax, ymax]，方便 cv2.rectangle 直接画框
return np.array([[boxes[i][0], boxes[i][1],
boxes[i][0] + boxes[i][2], boxes[i][1] + boxes[i][3],
confidences[i], cls_ids[i]] for i in indices.flatten()])
# ---- 5. 绘制 ----
def _draw(self, img, detections):
for xmin, ymin, xmax, ymax, conf, cls_id in detections:
label = f"{NAMES[int(cls_id)]} {conf:.2f}"
cv2.rectangle(img, (int(xmin), int(ymin)), (int(xmax), int(ymax)),
(0, 0, 255), 2)
cv2.putText(img, label, (int(xmin), int(ymin) - 5),
cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1)
return img
# ---- 6. 入口 ----
def run(self, img_path):
img = cv2.imread(img_path)
blob, bg = self._preprocess(img)
output = self._infer(blob)
detections = self._postprocess(output, bg)
self._draw(img, detections)
cv2.imshow("YOLOv8 TensorRT", img)
cv2.waitKey(0)
cv2.destroyAllWindows()
if __name__ == "__main__":
detector = YOLOv8EngineDetector(ENGINE_PATH, CONF_THRES, IOU_THRES)
st_time=time.time()
for i in range(1):
detector.run(IMAGE_PATH)
ed_time = time.time()
print("cost time:",(ed_time-st_time)/1)
结果如图
yolov8n_fp32.engine 和yolov8n_fp16.engine
int8推理如图
（量化后需要校准）
校准
build_int8.py
# -*- coding: utf-8 -*-
"""
用 IInt8EntropyCalibrator2 正确构建 yolov8n INT8 引擎
关键点（）：
校准图片必须和推理脚本用完全一样的预处理：
BGR→RGB、左上角正方形填充、resize 640、/255、CHW、float32
校准过程由 get_batch 返回"显存指针"实现（复用 ctypes 版 CUDA runtime，不装 pycuda）
用法: python build_int8.py
"""
import os
import numpy as np
import tensorrt as trt
import cv2
# 复用推理脚本里的 _preprocess，保证校准预处理和推理 100% 一致
import importlib.util
spec = importlib.util.spec_from_file_location(
"m8", r"D:\vscodefile\20260818\yolov8-tensort.py"
)
m8 = importlib.util.module_from_spec(spec)
spec.loader.exec_module(m8)
cuda = m8.cuda # 推理脚本里的 ctypes CUDA runtime
BASE = r"D:\vscodefile\20260818"
ONNX_PATH = os.path.join(BASE, "yolov8n.onnx")
ENGINE_PATH = os.path.join(BASE, "yolov8n_int8_v2.engine")
CACHE_PATH = os.path.join(BASE, "yolov8n_int8_v2.cache")
CALIB_IMAGES = [
os.path.join(BASE, "dog.png"),
] # 校准图越多越准，生产建议几百张
# os.path.join(BASE, "dance1.jpg")] # 校准图越多越准，生产建议几百张
class EntropyCalibrator(trt.IInt8EntropyCalibrator2):
"""把校准图预先拷到显存，get_batch 按 TRT 要求返回显存指针"""
def __init__(self, blobs, cache_file):
trt.IInt8EntropyCalibrator2.__init__(self)
self.cache_file = cache_file
self.dev_bufs = []
for b in blobs:
buf = cuda.mem_alloc(b.nbytes)
cuda.memcpy_htod(buf, np.ascontiguousarray(b.ravel()))
self.dev_bufs.append(buf)
self.idx = 0
def get_batch_size(self):
return 1
def get_batch(self, names):
if self.idx >= len(self.dev_bufs):
return None # None = 校准数据结束
ptr = self.dev_bufs[self.idx]
self.idx += 1
return [int(ptr)]
def read_calibration_cache(self):
if os.path.exists(self.cache_file):
with open(self.cache_file, "rb") as f:
return f.read()
return None
def write_calibration_cache(self, cache):
with open(self.cache_file, "wb") as f:
f.write(cache)
def free(self):
for p in self.dev_bufs:
cuda.mem_free(p)
def main():
# 1) 用推理脚本自己的预处理生成校准 blob
blobs = []
for p in CALIB_IMAGES:
img = cv2.imread(p)
assert img is not None, f"读图失败: {p}"
blob, _ = m8.YOLOv8EngineDetector._preprocess(None, img)
blobs.append(blob)
print(f"校准图 {os.path.basename(p)}: {img.shape} -> blob {blob.shape}")
# 2) 构建 INT8 引擎
logger = trt.Logger(trt.Logger.WARNING)
builder = trt.Builder(logger)
network = builder.create_network(1)
parser = trt.OnnxParser(network, logger)
with open(ONNX_PATH, "rb") as f:
if not parser.parse(f.read()):
for i in range(parser.num_errors):
print(f"[ERROR] {parser.get_error(i)}")
raise RuntimeError("ONNX 解析失败")
config = builder.create_builder_config()
config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 30)
config.set_flag(trt.BuilderFlag.INT8)
config.int8_calibrator = EntropyCalibrator(blobs, CACHE_PATH)
serialized = builder.build_serialized_network(network, config)
if serialized is None:
raise RuntimeError("engine 构建失败（serialized 为 None）")
with open(ENGINE_PATH, "wb") as f:
f.write(serialized)
print(f"✅ INT8 engine 保存成功: {ENGINE_PATH}")
if __name__ == "__main__":
main()
校准后推理如图
yolov5推理
训练环境
导出onnx (yolov5环境)
python export.py --weights yolov5s.pt --include onnx
拷贝yolov5s.onnx到甲方环境
甲方环境
trtexec.exe --onnx=yolov5s.onnx --saveEngine=yolov5s_fp16.engine --fp16
推理代码
"""
YOLOv5 TensorRT 推理脚本
环境:
conda activate trt85 （TRT 8.5.1.7 / CUDA 11.8 / Python 3.10）
engine 构建（GTX 1080 用 FP32，什么都不加就是 FP32）:
trtexec.exe --onnx=yolov5s.onnx --saveEngine=yolov5s.engine
用法: python yolov5-tensort.py
"""
import ctypes
import os
import cv2
import numpy as np
import tensorrt as trt
#
os.environ.setdefault("CUDA_MODULE_LOADING", "LAZY")
# ==================== CUDA Runtime Helper (ctypes) ====================
class _CudaRuntime:
"""
通过 ctypes 调用 CUDA Runtime API — 同步拷贝，不用流(stream)
为什么需要这个类？
TensorRT 的 Python 包只负责"推理计算"，不负责"显存管理"。
GPU 显存的分配/释放/拷贝必须通过 CUDA Runtime 完成，
这里用 ctypes 直接加载 cudart64_110.dll（CUDA 11.8 的运行时库），
从而免去安装 pycuda。
两个容易踩坑的关键点：
1. 每个函数都要声明 argtypes / restype：
ctypes 默认把参数当 32 位 int 传递，64 位指针会被截断导致崩溃。
2. _check 检查返回值：
CUDA API 出错时不抛异常、只返回错误码，必须自己检查。
"""
H2D = 1 # cudaMemcpyHostToDevice: CPU → GPU
D2H = 2 # cudaMemcpyDeviceToHost: GPU → CPU
def __init__(self):
# 加载 CUDA 运行时库：优先从 CUDA_PATH 找，其次依赖 PATH
cuda_bin = os.path.join(os.environ.get("CUDA_PATH", ""), "bin")
path = os.path.join(cuda_bin, "cudart64_110.dll")
self._dll = ctypes.CDLL(path if os.path.exists(path) else "cudart64_110.dll")
# 函数签名声明（必须！原因见类说明第 1 点）
self._dll.cudaMalloc.restype = ctypes.c_int
self._dll.cudaMalloc.argtypes = [
ctypes.POINTER(ctypes.c_void_p),
ctypes.c_size_t,
]
self._dll.cudaFree.argtypes = [ctypes.c_void_p]
self._dll.cudaMemcpy.restype = ctypes.c_int
self._dll.cudaMemcpy.argtypes = [
ctypes.c_void_p,
ctypes.c_void_p,
ctypes.c_size_t,
ctypes.c_int,
]
def _check(self, err, msg):
"""CUDA API 返回非 0 就是出错，抛出带说明的异常"""
if err != 0:
raise RuntimeError(f"{msg}: CUDA error {err}")
def mem_alloc(self, nbytes):
"""在 GPU 上分配 nbytes 字节显存，返回显存地址（裸指针 int）"""
ptr = ctypes.c_void_p()
self._check(self._dll.cudaMalloc(ctypes.byref(ptr), nbytes), "cudaMalloc")
return ptr.value
def mem_free(self, ptr):
"""释放 GPU 显存"""
self._dll.cudaFree(ctypes.c_void_p(ptr))
def memcpy_htod(self, dst, src_np):
"""把 numpy 数组拷到 GPU (Host → Device)，返回时拷贝已完成"""
self._check(
self._dll.cudaMemcpy(
ctypes.c_void_p(dst), src_np.ctypes.data, src_np.nbytes, self.H2D
),
"cudaMemcpy H2D",
)
def memcpy_dtoh(self, dst_np, src):
"""把 GPU 数据拷回 numpy 数组 (Device → Host)，返回时拷贝已完成"""
self._check(
self._dll.cudaMemcpy(
dst_np.ctypes.data, ctypes.c_void_p(src), dst_np.nbytes, self.D2H
),
"cudaMemcpy D2H",
)
cuda = _CudaRuntime()
# ============ 配置 ============
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
# yolov5s_int8.engine 校准也有问题（分数被压到 0.586 以下），改用 FP32 引擎
ENGINE_PATH = os.path.join(BASE_DIR, "yolov5s_fp16.engine")
IMAGE_PATH = os.path.join(BASE_DIR, "dog.png")
CONF_THRES = 0.5
IOU_THRES = 0.45
# COCO 80 类
NAMES = [
"person",
"bicycle",
"car",
"motorcycle",
"airplane",
"bus",
"train",
"truck",
"boat",
"traffic light",
"fire hydrant",
"stop sign",
"parking meter",
"bench",
"bird",
"cat",
"dog",
"horse",
"sheep",
"cow",
"elephant",
"bear",
"zebra",
"giraffe",
"backpack",
"umbrella",
"handbag",
"tie",
"suitcase",
"frisbee",
"skis",
"snowboard",
"sports ball",
"kite",
"baseball bat",
"baseball glove",
"skateboard",
"surfboard",
"tennis racket",
"bottle",
"wine glass",
"cup",
"fork",
"knife",
"spoon",
"bowl",
"banana",
"apple",
"sandwich",
"orange",
"broccoli",
"carrot",
"hot dog",
"pizza",
"donut",
"cake",
"chair",
"couch",
"potted plant",
"bed",
"dining table",
"toilet",
"tv",
"laptop",
"mouse",
"remote",
"keyboard",
"cell phone",
"microwave",
"oven",
"toaster",
"sink",
"refrigerator",
"book",
"clock",
"vase",
"scissors",
"teddy bear",
"hair drier",
"toothbrush",
]
TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
class YOLOv5EngineDetector:
def __init__(self, engine_path, conf_thres=0.5, iou_thres=0.45):
self.conf_thres = conf_thres
self.iou_thres = iou_thres
self.engine = self._load_engine(engine_path)
# ---- 1. 加载 engine ----
def _load_engine(self, engine_path):
with open(engine_path, "rb") as f, trt.Runtime(TRT_LOGGER) as runtime:
return runtime.deserialize_cuda_engine(f.read())
# ---- 2. 预处理 ----
def _preprocess(self, img, size=640):
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
h, w = img.shape[:2]
s = max(h, w)
bg = np.zeros((s, s, 3), np.uint8) # 正方化填充
bg[:h, :w] = img
blob = np.expand_dims(
cv2.resize(bg, (size, size)).transpose(2, 0, 1).astype(np.float32) / 255.0,
0,
)
return blob, bg
# ---- 3. GPU 推理 ----
def _infer(self, input_data):
context = self.engine.create_execution_context()
# YOLOv5 的 ONNX 输出名不统一（"output" / "output0"），按索引取最稳：
# 索引 0 是输入，索引 1 是输出（单输入单输出的模型）
input_name = self.engine.get_tensor_name(0)
output_name = self.engine.get_tensor_name(1)
dtype = np.float32 # FP32 engine：输入输出都是 float32
if self.engine.num_optimization_profiles > 0: # 动态 shape 引擎才需要
context.set_input_shape(input_name, input_data.shape)
output_shape = tuple(context.get_tensor_shape(output_name))
h_input = np.ascontiguousarray(input_data.ravel(), dtype=dtype)
h_output = np.empty(output_shape, np.float32)
d_input = cuda.mem_alloc(h_input.nbytes)
d_output = cuda.mem_alloc(h_output.nbytes)
context.set_tensor_address(input_name, d_input)
context.set_tensor_address(output_name, d_output)
# CPU → GPU → 推理 → GPU → CPU（0 = 默认流，按顺序执行）
cuda.memcpy_htod(d_input, h_input)
context.execute_async_v3(0)
cuda.memcpy_dtoh(h_output, d_output)
cuda.mem_free(d_input)
cuda.mem_free(d_output)
return h_output.reshape(output_shape)
# ---- 4. 后处理 (NMS) ----
def _postprocess(self, output, bg):
ratio = bg.shape[0] / 640
detections = output[0]
if detections.shape[0] == 85: # 兼容 (85,25200) 布局
detections = detections.T # 统一为 (25200,85)
boxes, confidences, cls_ids = [], [], []
for cx, cy, w, h, obj, *scores in detections:
cls_id = np.argmax(scores)
score = obj * scores[cls_id] # 最终得分 = objectness × 类别分数
if score > self.conf_thres:
boxes.append(
[
int((cx - w / 2) * ratio),
int((cy - h / 2) * ratio),
int(w * ratio),
int(h * ratio),
]
)
confidences.append(float(score))
cls_ids.append(int(cls_id))
if not boxes:
return np.array([])
indices = cv2.dnn.NMSBoxes(boxes, confidences, self.conf_thres, self.iou_thres)
if len(indices) == 0:
return np.array([])
# NMS 之后转回 [xmin, ymin, xmax, ymax]，方便 cv2.rectangle 画框
return np.array(
[
[
boxes[i][0],
boxes[i][1],
boxes[i][0] + boxes[i][2],
boxes[i][1] + boxes[i][3],
confidences[i],
cls_ids[i],
]
for i in indices.flatten()
]
)
# ---- 5. 绘制 ----
def _draw(self, img, detections):
for xmin, ymin, xmax, ymax, conf, cls_id in detections:
label = f"{NAMES[int(cls_id)]} {conf:.2f}"
cv2.rectangle(
img, (int(xmin), int(ymin)), (int(xmax), int(ymax)), (0, 0, 255), 2
)
cv2.putText(
img,
label,
(int(xmin), int(ymin) - 5),
cv2.FONT_HERSHEY_SIMPLEX,
0.5,
(0, 255, 0),
1,
)
return img
# ---- 6. 入口 ----
def run(self, img_path):
img = cv2.imread(img_path)
blob, bg = self._preprocess(img)
output = self._infer(blob)
detections = self._postprocess(output, bg)
self._draw(img, detections)
cv2.imshow("YOLOv5 TensorRT", img)
cv2.waitKey(0)
cv2.destroyAllWindows()
if __name__ == "__main__":
detector = YOLOv5EngineDetector(ENGINE_PATH, CONF_THRES, IOU_THRES)
detector.run(IMAGE_PATH)
结果如图
扫一扫
5065
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
