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

YOLOv8 OBB win10+ visual 2022移植部署

article 2024/9/21 22:37:14

前言

想做一个目标旋转角度检测的工程，但是网上多少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 size
    angle = 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 = angle
    x = 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 YOLO

model = 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() {

    //    //----initial
    auto engine = YoloOBB::create_infer(
        "E:***.trt",            // engine file
        0,                                      // gpu id
        0.25f,                                  // confidence threshold
        0.45f,                                  // nms threshold
        YoloOBB::NMSMethod::FastGPU,            // NMS method, fast GPU / CPU
        1024,                                   // max objects
        false                                   // 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;
}