当前位置：首页 > news >正文

YOLOv8 OBB win10+ visual 2022移植部署

news 来源：原创 2024/9/25 18:50:56

前言

想做一个目标旋转角度检测的工程，但是网上多少python的，或者linux的。在win10+ visual 2022移植部署，记录一下。
参考这篇文章没有C++ win10 环境下的部署教程，我相对于是对此做了补充。

1、下载工程

https://github.com/shouxieai/tensorRT_Pro

2 模型导出

在 ultralytics/engine/exporter.py 文件中改动一处

# output_names = ["output0", "output1"] if isinstance(self.model, SegmentationModel) else ["output0"]
# dynamic = self.args.dynamic
# if dynamic:
#     dynamic = {"images": {0: "batch", 2: "height", 3: "width"}}  # shape(1,3,640,640)
#     if isinstance(self.model, SegmentationModel):
#         dynamic["output0"] = {0: "batch", 2: "anchors"}  # shape(1, 116, 8400)
#         dynamic["output1"] = {0: "batch", 2: "mask_height", 3: "mask_width"}  # shape(1,32,160,160)
#     elif isinstance(self.model, DetectionModel):
#         dynamic["output0"] = {0: "batch", 2: "anchors"}  # shape(1, 84, 8400)
# ========== exporter.py ==========
output_names = ['output0', 'output1'] if isinstance(self.model, SegmentationModel) else ['output']
dynamic = self.args.dynamic
if dynamic:dynamic = {'images': {0: 'batch'}}  # shape(1,3,640,640)if isinstance(self.model, SegmentationModel):dynamic['output0'] = {0: 'batch', 2: 'anchors'}  # shape(1, 116, 8400)dynamic['output1'] = {0: 'batch', 2: 'mask_height', 3: 'mask_width'}  # shape(1,32,160,160)elif isinstance(self.model, DetectionModel):dynamic['output'] = {0: 'batch'}  # shape(1, 84, 8400)

在 ultralytics/nn/modules/head.py 文件中改动一处

def forward(self, x):"""Concatenates and returns predicted bounding boxes and class probabilities."""bs = x[0].shape[0]  # batch sizeangle = torch.cat([self.cv4[i](x[i]).view(bs, self.ne, -1) for i in range(self.nl)], 2)  # OBB theta logits# NOTE: set `angle` as an attribute so that `decode_bboxes` could use it.angle = (angle.sigmoid() - 0.25) * math.pi  # [-pi/4, 3pi/4]# angle = angle.sigmoid() * math.pi / 2  # [0, pi/2]if not self.training:self.angle = anglex = Detect.forward(self, x)if self.training:return x, angle#return torch.cat([x, angle], 1) if self.export else (torch.cat([x[0], angle], 1), (x[1], angle))return torch.cat([x, angle], 1).permute(0, 2, 1) if self.export else (torch.cat([x[0], angle], 1), (x[1], angle))

在终端执行如下指令即可完成 onnx 导出：

from ultralytics import YOLOmodel = YOLO("yolov8_obb_zwc_0918.pt")success = model.export(format="onnx", dynamic=True, simplify=True)

在这里插入图片描述

3 C++

修改一个简单的测试测序，在原工程app_yolo_obb.cpp 的基础上修改如下

#include "trt_builder.hpp"
#include <trt_infer.hpp>
#include <ilogger.hpp>
#include "yolo_obb.hpp"
static const char* dotalabels[] = {"0topleft", "1topright", "2downleft", "3downright", "4top","5right", "6down", "7left", "bridge", "large vehicle","small vehicle", "helicopter", "roundabout", "soccer ball field", "swimming pool"
};
using namespace std;
static vector<cv::Point> xywhr2xyxyxyxy(const YoloOBB::Box& box) {float cos_value = std::cos(box.angle);float sin_value = std::sin(box.angle);float w_2 = box.width / 2, h_2 = box.height / 2;float vec1_x = w_2 * cos_value, vec1_y = w_2 * sin_value;float vec2_x = -h_2 * sin_value, vec2_y = h_2 * cos_value;vector<cv::Point> corners;corners.push_back(cv::Point(box.center_x + vec1_x + vec2_x, box.center_y + vec1_y + vec2_y));corners.push_back(cv::Point(box.center_x + vec1_x - vec2_x, box.center_y + vec1_y - vec2_y));corners.push_back(cv::Point(box.center_x - vec1_x - vec2_x, box.center_y - vec1_y - vec2_y));corners.push_back(cv::Point(box.center_x - vec1_x + vec2_x, box.center_y - vec1_y + vec2_y));return corners;
}
static void test_single_image() {//    //----initialauto engine = YoloOBB::create_infer("E:***.trt",            // engine file0,                                      // gpu id0.25f,                                  // confidence threshold0.45f,                                  // nms thresholdYoloOBB::NMSMethod::FastGPU,            // NMS method, fast GPU / CPU1024,                                   // max objectsfalse                                   // preprocess use multi stream);if (engine == nullptr) {INFOE("Engine is nullptr");return;}//auto files = iLogger::find_files("IMG", "*.jpg;*.jpeg;*.png;*.gif;*.tif");vector<cv::Mat> images;for (int i = 0; i < files.size(); ++i) {auto image = cv::imread(files[i]);images.emplace_back(image);}//for (auto& img : images)//{//    cv::imshow("src", img);//    cv::waitKey(0);//}for(auto &img:images){cv::Mat image=img ;if (image.empty()) {INFOE("Image is empty");return;}auto boxes = engine->commit(image).get();for (auto& obj : boxes) {uint8_t b, g, r;tie(b, g, r) = iLogger::random_color(obj.class_label);auto corners = xywhr2xyxyxyxy(obj);cv::polylines(image, vector<vector<cv::Point>>{corners}, true, cv::Scalar(b, g, r), 2, 16);auto name = dotalabels[obj.class_label];auto caption = iLogger::format("%s %.2f", name, obj.confidence);int width = cv::getTextSize(caption, 0, 1, 2, nullptr).width + 10;cv::rectangle(image, cv::Point(corners[0].x - 3, corners[0].y - 33), cv::Point(corners[0].x - 3 + width, corners[0].y), cv::Scalar(b, g, r), -1);cv::putText(image, caption, cv::Point(corners[0].x - 3, corners[0].y - 5), 0, 1, cv::Scalar::all(0), 2, 16);}INFO("Save to Result.jpg, %d objects", boxes.size());cv::imwrite("Result.jpg", image);cv::imshow("result", image);cv::waitKey(0);}engine.reset();  
}
int main() {test_single_image();return 0;
}