C++打包成动态库给python调用

以yolov5模型推理c++工程封装为例：

C++打包成动态库

common.h

#ifndef COMMON_H
#define COMMON_H

#include <iostream>
#include <fstream>
#include <sstream>
#include <vector>
#include <dlfcn.h>
#include <time.h>
#include <cuda_runtime.h>
#include <NvInfer.h>
#include <NvInferRuntime.h>
#include <opencv2/opencv.hpp>

inline const char* severity_string(nvinfer1::ILogger::Severity t);

class TRTLogger : public nvinfer1::ILogger {
public:
void log(Severity severity, const char* msg) noexcept override;
};

extern TRTLogger logger;

std::vector<unsigned char> load_file(const std::string& file);

#endif

common.cpp

#include "common.h"

// 保持原有的TRTLogger、load_file实现…
inline const char* severity_string(nvinfer1::ILogger::Severity t) {
switch (t) {
case nvinfer1::ILogger::Severity::kINTERNAL_ERROR: return "INTERNAL_ERROR";
case nvinfer1::ILogger::Severity::kERROR: return "ERROR";
case nvinfer1::ILogger::Severity::kWARNING: return "WARNING";
case nvinfer1::ILogger::Severity::kINFO: return "INFO";
case nvinfer1::ILogger::Severity::kVERBOSE: return "VERBOSE";
default: return "UNKNOWN";
}
}

TRTLogger logger;

void TRTLogger::log(Severity severity, const char* msg) noexcept {
if (severity <= Severity::kWARNING) {
std::cerr << "TRT " << severity_string(severity) << ": " << msg << std::endl;
}
}

std::vector<unsigned char> load_file(const std::string& file) {
std::ifstream in(file, std::ios::binary);
if (!in.is_open()) return {};
in.seekg(0, std::ios::end);
size_t size = in.tellg();
std::vector<unsigned char> data(size);
in.seekg(0, std::ios::beg);
in.read((char*)data.data(), size);
in.close();
return data;
}

yolo.h

#ifndef YOLO_H
#define YOLO_H

#include "common.h"

// 定义推理结果结构体（供C++内部使用）
struct InferResult {
std::vector<float> output0;
};

// 关键：用extern "C"导出函数，避免C++名称修饰
#ifdef __cplusplus
extern "C" {
#endif

// 初始化函数：加载模型和插件
__attribute__((visibility("default"))) // 显式导出符号
bool yolo_init(const char* engine_path);

// 推理函数：接收输入指针和长度，返回结果指针（避免传递std::vector给C接口）
__attribute__((visibility("default")))
void yolo_infer(
const unsigned char* input0, int input0_shape0, int input0_shape1,
float** output0, int* output0_len
);

// 释放推理结果内存
__attribute__((visibility("default")))
void yolo_free_result(
float* output0
);

// 释放资源
__attribute__((visibility("default")))
void yolo_destroy();

#ifdef __cplusplus
}
#endif

#endif

yolo.cpp

#include "yolo.h"

// 全局变量保持引擎和上下文
static nvinfer1::ICudaEngine* engine = nullptr;
static nvinfer1::IExecutionContext* execution_context = nullptr;
static cudaStream_t stream = nullptr;

// 输出维度（根据你的模型修改，确保和之前一致）
static const int OUTPUT0_LEN = 4;
const int input_width = 640;
const int input_height = 640;
const int input_numel = 1 * 3 * input_width * input_height;
const float score_threshold = 0.2;
const float nms_threshold = 0.5;
const float confidence_threshold = 0.2;
const int num_classes = 80;
const int output_numprob = 4 + num_classes;
const int output_numbox = input_width / 8 * input_height / 8 + input_width / 16 * input_height / 16 + input_width / 32 * input_height / 32;
const int output_numel = 1 * output_numprob * output_numbox;

// 初始化函数（C风格接口）
bool yolo_init(const char* engine_path) {
if (!engine_path) {
std::cerr << "yolo引擎路径或插件路径为空" << std::endl;
return false;
}

// 加载引擎
std::string engine_path_str(engine_path);
auto engine_data = load_file(engine_path_str);
nvinfer1::IRuntime* runtime = nvinfer1::createInferRuntime(logger);
if (!runtime) {
std::cerr << "yolo创建TRT运行时失败" << std::endl;
return false;
}

engine = runtime->deserializeCudaEngine(engine_data.data(), engine_data.size());
if (!engine) {
std::cerr << "yolo反序列化引擎失败" << std::endl;
return false;
}

// 创建上下文和流
execution_context = engine->createExecutionContext();
if (!execution_context) {
std::cerr << "yolo创建执行上下文失败" << std::endl;
return false;
}

cudaError_t cuda_ret = cudaStreamCreate(&stream);
if (cuda_ret != cudaSuccess) {
std::cerr << "yolo创建CUDA流失败: " << cudaGetErrorString(cuda_ret) << std::endl;
return false;
}

std::cout << "yolo TRT引擎初始化成功" << std::endl;
return true;
}

//LetterBox处理
void LetterBox(const cv::Mat& image, cv::Mat& outImage, const cv::Size& newShape = cv::Size(640, 640), const cv::Scalar& color = cv::Scalar(114, 114, 114))
{
cv::Size shape = image.size();
float r = std::min((float)newShape.height / (float)shape.height, (float)newShape.width / (float)shape.width);
float ratio[2]{ r, r };
int new_un_pad[2] = { (int)std::round((float)shape.width * r),(int)std::round((float)shape.height * r) };

auto dw = (float)(newShape.width – new_un_pad[0]) / 2;
auto dh = (float)(newShape.height – new_un_pad[1]) / 2;

if (shape.width != new_un_pad[0] && shape.height != new_un_pad[1])
cv::resize(image, outImage, cv::Size(new_un_pad[0], new_un_pad[1]));
else
outImage = image.clone();

int top = int(std::round(dh – 0.1f));
int bottom = int(std::round(dh + 0.1f));
int left = int(std::round(dw – 0.1f));
int right = int(std::round(dw + 0.1f));

cv::Vec4d params;
params[0] = ratio[0];
params[1] = ratio[1];
params[2] = left;
params[3] = top;

cv::copyMakeBorder(outImage, outImage, top, bottom, left, right, cv::BORDER_CONSTANT, color);
}

//预处理
void pre_process(const cv::Mat& image, float* input_data_host)
{
cv::Mat letterbox;
LetterBox(image, letterbox, cv::Size(input_width, input_height));
letterbox.convertTo(letterbox, CV_32FC3, 1.0f / 255.0f);

int image_area = letterbox.cols * letterbox.rows;
float* pimage = (float*)letterbox.data;
float* phost_b = input_data_host + image_area * 0;
float* phost_g = input_data_host + image_area * 1;
float* phost_r = input_data_host + image_area * 2;
for (int i = 0; i < image_area; ++i, pimage += 3)
{
*phost_r++ = pimage[0];
*phost_g++ = pimage[1];
*phost_b++ = pimage[2];
}
}

//非最大值抑制
void nms(std::vector<cv::Rect>& boxes, std::vector<float>& scores, float score_threshold, float nms_threshold, std::vector<int>& indices)
{
assert(boxes.size() == scores.size());

struct BoxScore
{
cv::Rect box;
float score;
int id;
};
std::vector<BoxScore> boxes_scores;
for (size_t i = 0; i < boxes.size(); i++)
{
BoxScore box_conf;
box_conf.box = boxes[i];
box_conf.score = scores[i];
box_conf.id = i;
if (scores[i] > score_threshold)boxes_scores.push_back(box_conf);
}

std::sort(boxes_scores.begin(), boxes_scores.end(), [](BoxScore a, BoxScore b) { return a.score > b.score; });

std::vector<float> area(boxes_scores.size());
for (size_t i = 0; i < boxes_scores.size(); ++i)
{
area[i] = boxes_scores[i].box.width * boxes_scores[i].box.height;
}

std::vector<bool> isSuppressed(boxes_scores.size(), false);
for (size_t i = 0; i < boxes_scores.size(); ++i)
{
if (isSuppressed[i]) continue;
for (size_t j = i + 1; j < boxes_scores.size(); ++j)
{
if (isSuppressed[j]) continue;

float x1 = (std::max)(boxes_scores[i].box.x, boxes_scores[j].box.x);
float y1 = (std::max)(boxes_scores[i].box.y, boxes_scores[j].box.y);
float x2 = (std::min)(boxes_scores[i].box.x + boxes_scores[i].box.width, boxes_scores[j].box.x + boxes_scores[j].box.width);
float y2 = (std::min)(boxes_scores[i].box.y + boxes_scores[i].box.height, boxes_scores[j].box.y + boxes_scores[j].box.height);
float w = (std::max)(0.0f, x2 – x1);
float h = (std::max)(0.0f, y2 – y1);
float inter = w * h;
float ovr = inter / (area[i] + area[j] – inter);

if (ovr >= nms_threshold) isSuppressed[j] = true;
}
}

for (int i = 0; i < boxes_scores.size(); ++i)
{
if (!isSuppressed[i])indices.push_back(boxes_scores[i].id);
}
}

//box缩放到原图尺寸
void scale_boxes(cv::Rect& box, cv::Size size)
{
float gain = std::min(input_width * 1.0 / size.width, input_height * 1.0 / size.height);
int pad_w = (input_width – size.width * gain) / 2;
int pad_h = (input_height – size.height * gain) / 2;
box.x -= pad_w;
box.y -= pad_h;
box.x /= gain;
box.y /= gain;
box.width /= gain;
box.height /= gain;
}

//后处理
std::vector<float> post_process(const cv::Mat& image, float* output_data_host)
{
std::vector<cv::Rect> boxes;
std::vector<float> scores;
std::vector<int> class_ids;

for (int i = 0; i < output_numbox; ++i)
{
float* ptr = output_data_host + i * output_numprob;
float* classes_scores = 4 + ptr;
int class_id = std::max_element(classes_scores, classes_scores + num_classes) – classes_scores;
float score = classes_scores[class_id];
if (score < score_threshold)
continue;

float x = ptr[0];
float y = ptr[1];
float w = ptr[2];
float h = ptr[3];
int left = int(x – 0.5 * w);
int top = int(y – 0.5 * h);
int width = int(w);
int height = int(h);

cv::Rect box = cv::Rect(left, top, width, height);
scale_boxes(box, image.size());
boxes.push_back(box);
scores.push_back(score);
class_ids.push_back(class_id);
}

std::vector<int> indices;
nms(boxes, scores, score_threshold, nms_threshold, indices);

if(indices.size() > 0)
{
cv::Rect box = boxes[indices[0]];
return {float(box.x), float(box.y), float(box.x + box.width), float(box.y + box.height)};
}
return {0.0f, 0.0f, 0.0f, 0.0f};
}

// 推理函数（C风格接口：返回堆内存，需要调用yolo_free_result释放）
void yolo_infer(const unsigned char* input0, int input0_shape0, int input0_shape1,
float** output0, int* output0_len) {
// 初始化输出参数
*output0 = nullptr; *output0_len = 0;

// 检查输入有效性
if (!input0 || !engine || !execution_context) {
std::cerr << "yolo输入为空或引擎未初始化" << std::endl;
return;
}

// 分配主机端输出内存（堆内存，需要外部释放）
*output0 = new float[OUTPUT0_LEN];
*output0_len = OUTPUT0_LEN;

cv::Mat image(input0_shape0, input0_shape1, CV_8UC3, (unsigned char*)input0);

float* input_h = nullptr;
float* output_h = nullptr;
cudaMallocHost(&input_h, sizeof(float) * input_numel);
cudaMallocHost(&output_h, sizeof(float) * output_numel);

pre_process(image, input_h);

float* input_d = nullptr;
cudaMalloc(&input_d, sizeof(float) * input_numel);
cudaMemcpyAsync(input_d, input_h, sizeof(float) * input_numel, cudaMemcpyHostToDevice, stream);

float* output_d = nullptr;
cudaMalloc(&output_d, sizeof(float) * output_numel);

float* bindings[] = { input_d, output_d };
execution_context->executeV2((void**)bindings);
cudaMemcpyAsync(output_h, output_d, sizeof(float) * output_numel, cudaMemcpyDeviceToHost, stream);

std::vector<float> box = post_process(image, output_h);

cudaMemcpyAsync(*output0, box.data(), OUTPUT0_LEN * sizeof(float), cudaMemcpyHostToHost, stream);
cudaStreamSynchronize(stream);
cudaFree(input_d);
cudaFree(output_d);
}

// 释放推理结果内存（必须调用，避免内存泄漏）
void yolo_free_result(float* output0) {
delete[] output0;
}

// 释放资源
void yolo_destroy() {
if (stream) {
cudaStreamDestroy(stream);
stream = nullptr;
}
std::cout << "yolo TRT资源已释放" << std::endl;
}

CMakeLists.txt

cmake_minimum_required(VERSION 3.20)
project(trt_inference)
set(CMAKE_CXX_STANDARD 14)

find_package(CUDA REQUIRED)
include_directories(${CUDA_INCLUDE_DIRS})

set(OpenCV_DIR /home/tfy/document/HybrIK/cpp/opencv-4.12.0/lib/cmake/opencv4)
find_package(OpenCV REQUIRED)

set(TENSORRT_INCLUDE_DIRS /home/tfy/docker_share/TensorRT-10.6.0.26/include)
set(TENSORRT_LIBRARY_DIRS /home/tfy/docker_share/TensorRT-10.6.0.26/targets/x86_64-linux-gnu/lib)
set(TRT_ENABLE_ITT OFF CACHE BOOL "Disable TensorRT ITT tracing")
include_directories(${TENSORRT_INCLUDE_DIRS} ${OpenCV_INCLUDE_DIRS})
link_directories(${TENSORRT_LIBRARY_DIRS})

add_library(trt_infer SHARED common.cpp yolo.cpp)
target_link_libraries(trt_infer PRIVATE ${OpenCV_LIBS} ${CUDA_LIBRARIES} ${TENSORRT_LIBRARY_DIRS}/libnvinfer.so ${TENSORRT_LIBRARY_DIRS}/libnvinfer_plugin.so)

编译得到动态库libtrt_infer.so。

python调用

trt_infer.py

import ctypes
import os
import numpy as np
import cv2
import time

# 加载共享库（确保路径正确）
MODEL_PATH = "/home/tfy/document/HybrIK/yolov12n_fp32.engine"
LIB_PATH = "/home/tfy/document/HybrIK/cpp/build/libtrt_infer.so"

if not os.path.exists(LIB_PATH):
raise FileNotFoundError(f"共享库不存在：{LIB_PATH}")
trt_lib = ctypes.CDLL(LIB_PATH)

# ===================== 声明函数原型（必须严格匹配C接口）=====================
# 1. trt_init
trt_lib.yolo_init.argtypes = [ctypes.c_char_p]
trt_lib.yolo_init.restype = ctypes.c_bool

# 2. trt_infer：接收输入指针+长度，输出结果指针+长度（通过指针返回）
trt_lib.yolo_infer.argtypes = [
# 输入参数
ctypes.POINTER(ctypes.c_uint8), ctypes.c_int, ctypes.c_int,
# 输出参数（通过指针返回结果地址和长度）
ctypes.POINTER(ctypes.POINTER(ctypes.c_float)), ctypes.POINTER(ctypes.c_int),
]
trt_lib.yolo_infer.restype = None

# 3. trt_free_result：释放推理结果内存
trt_lib.yolo_free_result.argtypes = [
ctypes.POINTER(ctypes.c_float),
]
trt_lib.yolo_free_result.restype = None

# 4. trt_destroy：释放TRT资源
trt_lib.yolo_destroy.argtypes = []
trt_lib.yolo_destroy.restype = None

# ===================== 初始化TRT引擎 =====================
def init_trt(MODEL_PATH):
# 转换为C风格字符串（bytes类型）
MODEL_PATH_bytes = MODEL_PATH.encode('utf-8')
success = trt_lib.yolo_init(MODEL_PATH_bytes)
if not success:
raise RuntimeError("yolo TRT引擎初始化失败!")

# ===================== 执行推理 =====================
def infer_trt(input):
# 检查输入类型
if input.dtype != np.uint8:
input = input.astype(np.uint8)

# 准备输出参数（指针变量）
output_ptr = ctypes.POINTER(ctypes.c_float)()
output_len = ctypes.c_int()

# 调用C接口推理
trt_lib.yolo_infer(
# 输入：数据指针 + 长度
input.ctypes.data_as(ctypes.POINTER(ctypes.c_uint8)), input.shape[0], input.shape[1],
# 输出：结果指针的指针 + 长度指针
ctypes.byref(output_ptr), ctypes.byref(output_len)
)
output = np.ctypeslib.as_array(output_ptr, shape=(output_len.value,)).copy()
return output

# ===================== 测试代码 =====================
if __name__ == "__main__":
try:
# 1. 初始化
init_trt(MODEL_PATH)

# 2. 准备输入数据（和你的模型输入维度一致）
input = cv2.imread("/home/tfy/document/HybrIK/examples/000000000431.jpg")

# 3. 执行推理
for i in range(10):
start = time.time()
output = infer_trt(input)
print(output)
print(f"第{i+1}次推理耗时：{time.time() – start}秒")

except Exception as e:
print(f"错误：{e}")
finally:
# 5. 释放TRT资源（必须调用）
trt_lib.yolo_destroy()

运行python trt_infer.py
输出：

yolo TRT引擎初始化成功
[154. 67. 311. 321.]
第1次推理耗时：0.07363486289978027秒
[154. 67. 311. 321.]
第2次推理耗时：0.01118779182434082秒
[154. 67. 311. 321.]
第3次推理耗时：0.009070634841918945秒
[154. 67. 311. 321.]
第4次推理耗时：0.009114503860473633秒
[154. 67. 311. 321.]
第5次推理耗时：0.008969306945800781秒
[154. 67. 311. 321.]
第6次推理耗时：0.008978843688964844秒
[154. 67. 311. 321.]
第7次推理耗时：0.008900880813598633秒
[154. 67. 311. 321.]
第8次推理耗时：0.009151458740234375秒
[154. 67. 311. 321.]
第9次推理耗时：0.009235858917236328秒
[154. 67. 311. 321.]
第10次推理耗时：0.01047658920288086秒
yolo TRT资源已释放