YOLO-TensorRT 使用教程

项目介绍

YOLO-TensorRT 是一个基于 NVIDIA TensorRT 的 YOLO 目标检测模型的优化和加速项目。该项目旨在通过 TensorRT 的优化技术，提高 YOLO 模型的推理速度，同时保持高准确率。YOLO-TensorRT 支持多种 YOLO 模型，包括 YOLOv3、YOLOv4 等，并提供了 C++ 和 Python 的推理接口。

项目快速启动

环境准备

1. 安装 CUDA：确保系统中已安装 CUDA，推荐版本 >= 11.6。
2. 安装 TensorRT：下载并安装 TensorRT，推荐版本 >= 8.6。
3. 克隆项目：

   git clone https://github.com/enazoe/yolo-tensorrt.git
   cd yolo-tensorrt
   

编译与安装

1. 编译项目：

   mkdir build
   cd build
   cmake ..
   make
   

2. 模型转换：

   ./yolo-tensorrt --model /path/to/yolov3.cfg --weights /path/to/yolov3.weights --output /path/to/yolov3.trt
   

模型推理

1. C++ 推理示例：

   #include "yolo_tensorrt.h"
   
   int main() {
       YoloTensorRT yolo;
       yolo.loadModel("/path/to/yolov3.trt");
       yolo.detect("/path/to/image.jpg");
       return 0;
   }
   

2. Python 推理示例：

   from yolo_tensorrt import YoloTensorRT
   
   yolo = YoloTensorRT()
   yolo.load_model("/path/to/yolov3.trt")
   yolo.detect("/path/to/image.jpg")
   

应用案例和最佳实践

案例一：实时视频检测

在视频流中实时检测目标，可以应用于安防监控、自动驾驶等领域。

import cv2
from yolo_tensorrt import YoloTensorRT

yolo = YoloTensorRT()
yolo.load_model("/path/to/yolov3.trt")

cap = cv2.VideoCapture(0)
while True:
    ret, frame = cap.read()
    if not ret:
        break
    results = yolo.detect(frame)
    for result in results:
        x, y, w, h, label = result
        cv2.rectangle(frame, (x, y), (x+w, y+h), (0, 255, 0), 2)
        cv2.putText(frame, label, (x, y-10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)
    cv2.imshow('YOLO Detection', frame)
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

cap.release()
cv2.destroyAllWindows()

案例二：批量图片检测

批量处理图片，适用于大规模图像数据集的检测任务。

import os
from yolo_tensorrt import YoloTensorRT

yolo = YoloTensorRT()
yolo.load_model("/path/to/yolov3.trt")

image_dir = "/path/to/images"
for filename in os.listdir(image_dir):
    if filename.endswith(".jpg"):
        image_path = os.path.join(image_dir, filename)
        results = yolo.detect(image_path)
        for result in results:
            x, y, w, h, label = result
            print(f"Detected {label} at ({x}, {y}) with size ({w}, {h}) in {filename}")

典型生态项目

TensorRT-YOLO

TensorRT-YOLO 是一个基于 TensorRT 的 YOLO 模型加速项目，支持多种 YOLO