python
import cv2
import numpy as np
import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit
classNames = []
with open("coco.names", "r") as f:
for line in f:
line = line.strip()
if line:
classNames.append(line)
TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
with open("yolo26n.engine", "rb") as f:
runtime = trt.Runtime(TRT_LOGGER)
engine = runtime.deserialize_cuda_engine(f.read())
context = engine.create_execution_context()
input_shape = engine.get_tensor_shape("images")
output_shape = engine.get_tensor_shape("output0")
input_h, input_w = input_shape[2], input_shape[3]
output_h, output_w = output_shape[1], output_shape[2]
print(f"Input H: {input_h}, Input W: {input_w}")
print(f"Output data format: {output_h}x{output_w}")
host_input = np.empty(input_shape, dtype=np.float32)
device_input = cuda.mem_alloc(host_input.nbytes)
host_output = np.empty(output_shape, dtype=np.float32)
device_output = cuda.mem_alloc(host_output.nbytes)
context.set_tensor_address("images", int(device_input))
context.set_tensor_address("output0", int(device_output))
stream = cuda.Stream()
start_time = cv2.getTickCount()
frame = cv2.imread("bus.jpg")
h, w = frame.shape[:2]
_max = max(h, w)
# Letterbox 填充
image = np.zeros((_max, _max, 3), dtype=np.uint8)
image[:h, :w] = frame
x_factor = image.shape[1] / float(input_w)
y_factor = image.shape[0] / float(input_h)
blob = cv2.dnn.blobFromImage(image, scalefactor=1 / 255.0, size=(input_w, input_h), swapRB=True)
host_input = np.ascontiguousarray(blob, dtype=np.float32)
cuda.memcpy_htod_async(device_input, host_input, stream)
context.execute_async_v3(stream_handle=stream.handle)
cuda.memcpy_dtoh_async(host_output, device_output, stream)
stream.synchronize()
detection_outputs = host_output.reshape(output_h, output_w)
boxes, classIds, confidences = [], [], []
for i in range(detection_outputs.shape[0]):
confidence = detection_outputs[i, 4]
if confidence > 0.25:
x1, y1, x2, y2 = detection_outputs[i, :4]
class_id = int(detection_outputs[i, 5])
width = int((x2 - x1) * x_factor)
height = int((y2 - y1) * y_factor)
x = int(x1 * x_factor)
y = int(y1 * y_factor)
boxes.append([x, y, width, height])
classIds.append(class_id)
confidences.append(float(confidence))
for i in range(len(boxes)):
x, y, width, height = boxes[i]
cv2.rectangle(frame, (x, y), (x + width, y + height), (0, 0, 255), 2)
cv2.rectangle(frame, (x, y - 20), (x + width, y), (0, 255, 255), -1)
cv2.putText(frame, classNames[classIds[i]], (x, y),
cv2.FONT_HERSHEY_PLAIN, 2.0, (255, 0, 0), 2)
t = (cv2.getTickCount() - start_time) / cv2.getTickFrequency()
cv2.putText(frame, f"FPS: {1.0 / t:.2f}", (20, 40),
cv2.FONT_HERSHEY_PLAIN, 2.0, (255, 0, 0), 2)
cv2.imshow("YOLO26+TensorRT10.8 (Python)", frame)
cv2.waitKey(0)
cv2.destroyAllWindows()Deploy.cpp
cpp
#include <iostream>
#include <fstream>
#include <sstream>
#include <opencv2/opencv.hpp>
#include "NvInfer.h"
#include "NvOnnxParser.h"
using namespace nvinfer1;
using namespace nvonnxparser;
using namespace cv;
class Logger : public ILogger
{
void log(Severity severity, const char* msg) noexcept
{
// suppress info-level messages
if (severity != Severity::kINFO)
std::cout << msg << std::endl;
}
} gLogger;
int main(int argc, char** argv) {
std::vector<std::string> classNames;
std::string label_map = "coco.names";
std::ifstream fp(label_map);
std::string name;
while (!fp.eof()) {
getline(fp, name);
if (name.length()) {
classNames.push_back(name);
}
}
float confidence_threshold = 0.4;
float score_threshold = 0.25;
void* buffers[2] = { NULL, NULL };
std::vector<float> prob;
cudaStream_t stream;
fp.close();
std::ifstream file("yolo26n.engine", std::ios::binary);
char* trtModelStream = NULL;
int size = 0;
if (file.good()) {
file.seekg(0, file.end);
size = file.tellg();
file.seekg(0, file.beg);
trtModelStream = new char[size];
assert(trtModelStream);
file.read(trtModelStream, size);
file.close();
}
IRuntime* runtime = createInferRuntime(gLogger);
ICudaEngine* engine = runtime->deserializeCudaEngine(trtModelStream, size);
IExecutionContext* context = engine->createExecutionContext();
delete[] trtModelStream;
auto inputDims = engine->getTensorShape("images");
auto outDims = engine->getTensorShape("output0");
int input_h = inputDims.d[2];
int input_w = inputDims.d[3];
printf("inputH : %d, inputW: %d \n", input_h, input_w);
int output_h = outDims.d[1];
int output_w = outDims.d[2];
std::cout << "out data format: " << output_h << "x" << output_w << std::endl;
cudaMalloc(&buffers[0], input_h * input_w * 3 * sizeof(float));
cudaMalloc(&buffers[1], output_h * output_w * sizeof(float));
prob.resize(output_h * output_w);
cudaStreamCreate(&stream);
int64 start = cv::getTickCount();
cv::Mat frame = cv::imread("bus.jpg");
int w = frame.cols;
int h = frame.rows;
int _max = std::max(h, w);
cv::Mat image = cv::Mat::zeros(cv::Size(_max, _max), CV_8UC3);
cv::Rect roi(0, 0, w, h);
frame.copyTo(image(roi));
float x_factor = image.cols / static_cast<float>(input_w);
float y_factor = image.rows / static_cast<float>(input_h);
cv::Mat blob = cv::dnn::blobFromImage(image, 1 / 255.0, cv::Size(input_w, input_h), cv::Scalar(0, 0, 0), true, false);
cudaMemcpyAsync(buffers[0], blob.ptr<float>(), 3 * input_h * input_w * sizeof(float), cudaMemcpyHostToDevice, stream);
context->executeV2(buffers);
cudaMemcpyAsync(prob.data(), buffers[1], output_h * output_w * sizeof(float), cudaMemcpyDeviceToHost, stream);
cv::Mat detection_outputs(output_h, output_w, CV_32F, (float*)prob.data());
std::vector<cv::Rect> boxes;
std::vector<int> classIds;
std::vector<float> confidences;
for (int i = 0; i < detection_outputs.rows; ++i)
{
double confidence = detection_outputs.at<float>(i, 4);
if (confidence > 0.25)
{
const float x1 = detection_outputs.at<float>(i, 0);
const float y1 = detection_outputs.at<float>(i, 1);
const float x2 = detection_outputs.at<float>(i, 2);
const float y2 = detection_outputs.at<float>(i, 3);
int class_id = static_cast<int>(detection_outputs.at<float>(i, 5));
int width = static_cast<int>((x2 - x1) * x_factor);
int height = static_cast<int>((y2 - y1) * y_factor);
int x = static_cast<int>(x1 * x_factor);
int y = static_cast<int>(y1 * y_factor);
cv::Rect box;
box.x = x;
box.y = y;
box.width = width;
box.height = height;
boxes.push_back(box);
classIds.push_back(class_id);
confidences.push_back(confidence);
}
}
for (size_t i = 0; i < boxes.size(); i++)
{
cv::rectangle(frame, boxes[i], cv::Scalar(0, 0, 255), 2, 8);
cv::rectangle(frame, cv::Point(boxes[i].tl().x, boxes[i].tl().y - 20),
cv::Point(boxes[i].br().x, boxes[i].tl().y), cv::Scalar(0, 255, 255), -1);
cv::putText(frame, classNames[classIds[i]], cv::Point(boxes[i].tl().x, boxes[i].tl().y), cv::FONT_HERSHEY_PLAIN, 2.0, cv::Scalar(255, 0, 0), 2, 8);
}
float t = (cv::getTickCount() - start) / static_cast<float>(cv::getTickFrequency());
cv::putText(frame, cv::format("FPS: %.2f", 1.0 / t), cv::Point(20, 40), cv::FONT_HERSHEY_PLAIN, 2.0, cv::Scalar(255, 0, 0), 2, 8);
cv::imshow("YOLO26+TENSORRT10.8", frame);
cv::waitKey(0);
return 0;
}CMakeLists.txt
bash
cmake_minimum_required(VERSION 3.18)
project(yolo26)
set(OpenCV_DIR "E:\\Opencv gpu\\newbuild\\install")
set(OpenCV_INCLUDE_DIRS ${OpenCV_DIR}\\include)
set(OpenCV_LIB_DIRS ${OpenCV_DIR}\\x64\\vc17\\lib)
set(OpenCV_LIB_DEBUG ${OpenCV_DIR}\\x64\\vc17\\lib\\opencv_world470d.lib
${OpenCV_DIR}\\x64\\vc17\\lib\\opencv_img_hash470d.lib)
set(OpenCV_LIB_RELEASE ${OpenCV_DIR}\\x64\\vc17\\lib\\opencv_world470.lib
${OpenCV_DIR}\\x64\\vc17\\lib\\opencv_img_hash470.lib)
set(CMAKE_CUDA_ARCHITECTURES 86)
find_package(CUDA REQUIRED)
enable_language(CUDA)
find_package(OpenCV QUIET)
include_directories(${CUDA_INCLUDE_DIRS})
include_directories(${OpenCV_INCLUDE_DIRS})
link_directories(${OpenCV_LIB_DIRS})
set(SOURCES
main.cpp
)
add_executable(yolo26 ${SOURCES})
target_link_libraries(yolo26 "nvinfer.lib" "nvinfer_plugin.lib" "nvonnxparser.lib")
target_link_libraries(yolo26 ${CUDA_LIBRARIES})
target_link_libraries(yolo26
$<$<CONFIG:Debug>:${OpenCV_LIB_DEBUG}>
$<$<CONFIG:Release>:${OpenCV_LIB_RELEASE}>
)