YOLOv8 人体姿态估计(关键点检测) python推理 && ONNX RUNTIME C++部署

目录


 

1、下载权重

我这里之前在做实例分割的时候,项目已经下载到本地,环境也安装好了,只需要下载pose的权重就可以



2、python 推理

yolo task=pose mode=predict model=yolov8n-pose.pt  source=0  show=true


3、转ONNX格式

yolo export model=yolov8n-pose.pt format=onnx 

输出:
 

(yolo) jason@honor:~/PycharmProjects/pytorch_learn/yolo/ultralytics-main-yolov8$ yolo export model=yolov8n-pose.pt format=onnx
Ultralytics YOLOv8.0.94 🚀 Python-3.8.13 torch-2.0.0+cu117 CPU
YOLOv8n-pose summary (fused): 187 layers, 3289964 parameters, 0 gradients, 9.2 GFLOPs

PyTorch: starting from yolov8n-pose.pt with input shape (1, 3, 640, 640) BCHW and output shape(s) (1, 56, 8400) (6.5 MB)

ONNX: starting export with onnx 1.13.1 opset 17...
============= Diagnostic Run torch.onnx.export version 2.0.0+cu117 =============
verbose: False, log level: Level.ERROR
======================= 0 NONE 0 NOTE 0 WARNING 0 ERROR ========================

ONNX: export success ✅ 0.8s, saved as yolov8n-pose.onnx (12.9 MB)

Export complete (1.4s)
Results saved to /home/jason/PycharmProjects/pytorch_learn/yolo/ultralytics-main-yolov8
Predict:         yolo predict task=pose model=yolov8n-pose.onnx imgsz=640 
Validate:        yolo val task=pose model=yolov8n-pose.onnx imgsz=640 data=/usr/src/app/ultralytics/datasets/coco-pose.yaml 
Visualize:       https://netron.app

用netron查看一下:

 如上图所是,YOLOv8n-pose只有一个输出:
output0: float32[1,56,8400]。这里的8400,表示有8400个检测框,56为4边界框坐标信息+人这个类别预测分数,17*3关键点信息。每个关键点由x,y,v组成,v代表该点是否可见,v小于 0.5 时,表示这个关键点可能在图外,可以考虑去除掉。

COCO的annotation一共有17个关节点。

分别是:“nose”,“left_eye”, “right_eye”,“left_ear”, “right_ear”,“left_shoulder”, “right_shoulder”,“left_elbow”, “right_elbow”,“left_wrist”, “right_wrist”,“left_hip”, “right_hip”,“left_knee”, “right_knee”,“left_ankle”, “right_ankle”。示例图如下:

4、ONNX RUNTIME C++ 部署

第二篇参考文章的github项目,以此为参考,实现ONNX RUNTIME C++部署

视频输入,效果如下:

utils.h

#pragma once
#include <iostream>
#include <opencv2/opencv.hpp>




struct  OutputPose {

    cv::Rect_<float> box;
    int label =0;
    float confidence =0.0;
    std::vector<float> kps;

};

void DrawPred(cv::Mat& img, std::vector<OutputPose>& results,
              const std::vector<std::vector<unsigned int>> &SKELLTON,
              const std::vector<std::vector<unsigned int>> &KPS_COLORS,
              const std::vector<std::vector<unsigned int>> &LIMB_COLORS);
void LetterBox(const cv::Mat& image, cv::Mat& outImage,
               cv::Vec4d& params,
               const cv::Size& newShape = cv::Size(640, 640),
               bool autoShape = false,
               bool scaleFill=false,
               bool scaleUp=true,
               int stride= 32,
               const cv::Scalar& color = cv::Scalar(114,114,114));

utils.cpp

#pragma once
#include "utils.h"
using namespace cv;
using namespace std;

void LetterBox(const cv::Mat& image, cv::Mat& outImage,
               cv::Vec4d& params,
               const cv::Size& newShape,
               bool autoShape,
               bool scaleFill,
               bool scaleUp,
               int stride,
               const cv::Scalar& color)
{
    if (false) {
        int maxLen = MAX(image.rows, image.cols);
        outImage = Mat::zeros(Size(maxLen, maxLen), CV_8UC3);
        image.copyTo(outImage(Rect(0, 0, image.cols, image.rows)));
        params[0] = 1;
        params[1] = 1;
        params[3] = 0;
        params[2] = 0;
    }

    // 取较小的缩放比例
    cv::Size shape = image.size();
    float r = std::min((float)newShape.height / (float)shape.height,
                       (float)newShape.width / (float)shape.width);
    if (!scaleUp)
        r = std::min(r, 1.0f);
    printf("原图尺寸:w:%d * h:%d, 要求尺寸:w:%d * h:%d, 即将采用的缩放比:%f\n",
           shape.width, shape.height, newShape.width, newShape.height, r);

    // 依据前面的缩放比例后,原图的尺寸
    float ratio[2]{r,r};
    int new_un_pad[2] = { (int)std::round((float)shape.width  * r), (int)std::round((float)shape.height * r)};
    printf("等比例缩放后的尺寸该为:w:%d * h:%d\n", new_un_pad[0], new_un_pad[1]);

    // 计算距离目标尺寸的padding像素数
    auto dw = (float)(newShape.width - new_un_pad[0]);
    auto dh = (float)(newShape.height - new_un_pad[1]);
    if (autoShape)
    {
        dw = (float)((int)dw % stride);
        dh = (float)((int)dh % stride);
    }
    else if (scaleFill)
    {
        dw = 0.0f;
        dh = 0.0f;
        new_un_pad[0] = newShape.width;
        new_un_pad[1] = newShape.height;
        ratio[0] = (float)newShape.width / (float)shape.width;
        ratio[1] = (float)newShape.height / (float)shape.height;
    }

    dw /= 2.0f;
    dh /= 2.0f;
    printf("填充padding: dw=%f , dh=%f\n", dw, dh);

    // 等比例缩放
    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();
    }

    // 图像四周padding填充,至此原图与目标尺寸一致
    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));
    params[0] = ratio[0]; // width的缩放比例
    params[1] = ratio[1]; // height的缩放比例
    params[2] = left; // 水平方向两边的padding像素数
    params[3] = top; //垂直方向两边的padding像素数
    cv::copyMakeBorder(outImage, outImage, top, bottom, left, right, cv::BORDER_CONSTANT, color);
}

void DrawPred(cv::Mat& img, std::vector<OutputPose>& results,
              const std::vector<std::vector<unsigned int>> &SKELLTON,
              const std::vector<std::vector<unsigned int>> &KPS_COLORS,
              const std::vector<std::vector<unsigned int>> &LIMB_COLORS)
{
    const int num_point =17;
    for (auto &result:results){
        int  left,top,width, height;
        left = result.box.x;
        top = result.box.y;
        width = result.box.width;
        height = result.box.height;


//        printf("x: %d  y:%d  w:%d  h%d\n",(int)left, (int)top, (int)result.box.width, (int)result.box.height);

        // 框出目标
        rectangle(img, result.box,Scalar(0,0,255), 2, 8);

        // 在目标框左上角标识目标类别以及概率
        string label = "person:" + to_string(result.confidence) ;
        int baseLine;
        Size labelSize = getTextSize(label, FONT_HERSHEY_SIMPLEX, 0.5, 1, &baseLine);
        top = max(top, labelSize.height);
        putText(img, label, Point(left, top), FONT_HERSHEY_SIMPLEX, 0.5, Scalar(0,0,255), 2);

        // 连线
        auto &kps = result.kps;
//        cout << "该目标关键点:" << kps.size() << endl;
        for (int k=0; k<num_point+2; k++){// 不要设置为>0.5f ,>0.0f显示效果比较好
            // 关键点绘制
            if (k<num_point){
                int kps_x = std::round(kps[k*3]);
                int kps_y = std::round(kps[k*3 + 1]);
                float kps_s = kps[k*3 + 2];

//                printf("x:%d y:%d s:%f\n", kps_x, kps_y, kps_s);

                if (kps_s > 0.0f){
                    cv::Scalar kps_color = Scalar(KPS_COLORS[k][0],KPS_COLORS[k][1],KPS_COLORS[k][2]);
                    cv::circle(img, {kps_x, kps_y}, 5, kps_color, -1);
                }
            }

            auto &ske = SKELLTON[k];
            int pos1_x = std::round(kps[(ske[0] -1) * 3]);
            int pos1_y = std::round(kps[(ske[0] -1) * 3 + 1]);

            int pos2_x = std::round(kps[(ske[1] -1) * 3]);
            int pos2_y = std::round(kps[(ske[1] -1) * 3 + 1]);

            float pos1_s = kps[(ske[0] -1) * 3 + 2];
            float pos2_s = kps[(ske[1] -1) * 3 + 2];

            if (pos1_s > 0.0f && pos2_s >0.0f){// 不要设置为>0.5f ,>0.0f显示效果比较好
                cv::Scalar limb_color = cv::Scalar(LIMB_COLORS[k][0], LIMB_COLORS[k][1], LIMB_COLORS[k][3]);
                cv::line(img, {pos1_x, pos1_y}, {pos2_x, pos2_y}, limb_color);
            }

        // 跌倒检测
            float pt5_x = kps[5*3];
            float pt5_y = kps[5*3 + 1];
            float pt6_x = kps[6*3];
            float pt6_y = kps[6*3+1];
            float center_up_x = (pt5_x + pt6_x) /2.0f ;
            float center_up_y = (pt5_y + pt6_y) / 2.0f;
            Point center_up = Point((int)center_up_x, (int)center_up_y);

            float pt11_x = kps[11*3];
            float pt11_y = kps[11*3 + 1];
            float pt12_x = kps[12*3];
            float pt12_y = kps[12*3 + 1];
            float center_down_x = (pt11_x + pt12_x) / 2.0f;
            float center_down_y = (pt11_y + pt12_y) / 2.0f;
            Point center_down = Point((int)center_down_x, (int)center_down_y);


            float right_angle_point_x = center_down_x;
            float righ_angle_point_y = center_up_y;
            Point right_angl_point = Point((int)right_angle_point_x, (int)righ_angle_point_y);


            float a = abs(right_angle_point_x - center_up_x);
            float b = abs(center_down_y - righ_angle_point_y);

            float tan_value = a / b;
            float Pi = acos(-1);
            float angle = atan(tan_value) * 180.0f/ Pi;
            string angel_label = "angle: " + to_string(angle);
            putText(img, angel_label, Point(left, top-40), FONT_HERSHEY_SIMPLEX, 0.5, Scalar(0,0,255), 2);

            if (angle > 60.0f || center_down_y <= center_up_y || (double)width/ height > 5.0f/3.0f) // 宽高比小于0.6为站立,大于5/3为跌倒
            {
                string fall_down_label = "person fall down!!!!";
                putText(img, fall_down_label , Point(left, top-20), FONT_HERSHEY_SIMPLEX, 0.5, Scalar(0,0,255), 2);

                printf("angel:%f width/height:%f\n",angle, (double)width/ height );
            }









            cv::line(img, center_up, center_down,
                     Scalar(0,0,255), 2, 8);
            cv::line(img, center_up, right_angl_point,
                     Scalar(0,0,255), 2, 8);
            cv::line(img, right_angl_point, center_down,
                     Scalar(0,0,255), 2, 8);





        }
    }
}

detect.h

#pragma onece
#include <iostream>
#include <memory>
#include <opencv2/opencv.hpp>
#include "utils.h"
#include <onnxruntime_cxx_api.h>
#include <numeric>

class Yolov8Onnx{
private:
    template<typename T>
    T VectorProduct(const std::vector<T>& v)
    {
        return std::accumulate(v.begin(), v.end(), 1, std::multiplies<T>());
    }

    int Preprocessing(const std::vector<cv::Mat>& SrcImgs,
                      std::vector<cv::Mat>& OutSrcImgs,
                      std::vector<cv::Vec4d>& params);


    const int _netWidth = 640;   //ONNX-net-input-width
    const int _netHeight = 640;  //ONNX-net-input-height

    int _batchSize = 1; //if multi-batch,set this
    bool _isDynamicShape = false;//onnx support dynamic shape

    int _anchorLength=56;// pose一个框的信息56个数

    float _classThreshold = 0.25;
    float _nmsThrehold= 0.45;



    //ONNXRUNTIME
    Ort::Env _OrtEnv = Ort::Env(OrtLoggingLevel::ORT_LOGGING_LEVEL_ERROR, "Yolov5-Seg");
    Ort::SessionOptions _OrtSessionOptions = Ort::SessionOptions();
    Ort::Session* _OrtSession = nullptr;
    Ort::MemoryInfo _OrtMemoryInfo;

    std::shared_ptr<char> _inputName, _output_name0;
    std::vector<char*> _inputNodeNames; //输入节点名
    std::vector<char*> _outputNodeNames; // 输出节点名

    size_t _inputNodesNum = 0;        // 输入节点数
    size_t _outputNodesNum = 0;      // 输出节点数

    ONNXTensorElementDataType _inputNodeDataType;  //数据类型
    ONNXTensorElementDataType _outputNodeDataType;

    std::vector<int64_t> _inputTensorShape;  // 输入张量shape
    std::vector<int64_t> _outputTensorShape;

public:
    Yolov8Onnx():_OrtMemoryInfo(Ort::MemoryInfo::CreateCpu(OrtAllocatorType::OrtDeviceAllocator, OrtMemType::OrtMemTypeCPUOutput)) {};
    ~Yolov8Onnx() {};// delete _OrtMemoryInfo;

public:
    /** \brief Read onnx-model
    * \param[in] modelPath:onnx-model path
    * \param[in] isCuda:if true,use Ort-GPU,else run it on cpu.
    * \param[in] cudaID:if isCuda==true,run Ort-GPU on cudaID.
    * \param[in] warmUp:if isCuda==true,warm up GPU-model.
    */
    bool ReadModel(const std::string& modelPath, bool isCuda=false, int cudaId=0, bool warmUp=true);

    /** \brief  detect.
    * \param[in] srcImg:a 3-channels image.
    * \param[out] output:detection results of input image.
    */
    bool OnnxDetect(cv::Mat& srcImg, std::vector<OutputPose>& output);

    /** \brief  detect,batch size= _batchSize
    * \param[in] srcImg:A batch of images.
    * \param[out] output:detection results of input images.
    */
    bool OnnxBatchDetect(std::vector<cv::Mat>& srcImgs, std::vector<std::vector<OutputPose>>& output);

//public:
//    std::vector<std::string> _className = {
//        "person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck", "boat", "traffic light",
//        "fire hydrant", "stop sign", "parking meter", "bench", "bird", "cat", "dog", "horse", "sheep", "cow",
//        "elephant", "bear", "zebra", "giraffe", "backpack", "umbrella", "handbag", "tie", "suitcase", "frisbee",
//        "skis", "snowboard", "sports ball", "kite", "baseball bat", "baseball glove", "skateboard", "surfboard",
//        "tennis racket", "bottle", "wine glass", "cup", "fork", "knife", "spoon", "bowl", "banana", "apple",
//        "sandwich", "orange", "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", "chair", "couch",
//        "potted plant", "bed", "dining table", "toilet", "tv", "laptop", "mouse", "remote", "keyboard", "cell phone",
//        "microwave", "oven", "toaster", "sink", "refrigerator", "book", "clock", "vase", "scissors", "teddy bear",
//        "hair drier", "toothbrush"
//    };



};

detect.cpp

#include "detect.h"
using namespace std;
using namespace cv;
using namespace cv::dnn;
using namespace Ort;

bool Yolov8Onnx::ReadModel(const std::string &modelPath, bool isCuda, int cudaId, bool warmUp){
    if (_batchSize < 1) _batchSize =1;
    try
    {
        std::vector<std::string> available_providers = GetAvailableProviders();
        auto cuda_available = std::find(available_providers.begin(), available_providers.end(), "CUDAExecutionProvider");


        if (isCuda && (cuda_available == available_providers.end()))
        {
            std::cout << "Your ORT build without GPU. Change to CPU." << std::endl;
            std::cout << "************* Infer model on CPU! *************" << std::endl;
        }
        else if (isCuda && (cuda_available != available_providers.end()))
        {
            std::cout << "************* Infer model on GPU! *************" << std::endl;
//#if ORT_API_VERSION < ORT_OLD_VISON
//			OrtCUDAProviderOptions cudaOption;
//			cudaOption.device_id = cudaID;
//            _OrtSessionOptions.AppendExecutionProvider_CUDA(cudaOption);
//#else
//			OrtStatus* status = OrtSessionOptionsAppendExecutionProvider_CUDA(_OrtSessionOptions, cudaID);
//#endif
        }
        else
        {
            std::cout << "************* Infer model on CPU! *************" << std::endl;
        }
        //

        _OrtSessionOptions.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_EXTENDED);

#ifdef _WIN32
        std::wstring model_path(modelPath.begin(), modelPath.end());
        _OrtSession = new Ort::Session(_OrtEnv, model_path.c_str(), _OrtSessionOptions);
#else
        _OrtSession = new Ort::Session(_OrtEnv, modelPath.c_str(), _OrtSessionOptions);
#endif

        Ort::AllocatorWithDefaultOptions allocator;
        //init input
        _inputNodesNum = _OrtSession->GetInputCount();
#if ORT_API_VERSION < ORT_OLD_VISON
        _inputName = _OrtSession->GetInputName(0, allocator);
        _inputNodeNames.push_back(_inputName);
#else
        _inputName = std::move(_OrtSession->GetInputNameAllocated(0, allocator));
        _inputNodeNames.push_back(_inputName.get());
#endif
        //cout << _inputNodeNames[0] << endl;
        Ort::TypeInfo inputTypeInfo = _OrtSession->GetInputTypeInfo(0);
        auto input_tensor_info = inputTypeInfo.GetTensorTypeAndShapeInfo();
        _inputNodeDataType = input_tensor_info.GetElementType();
        _inputTensorShape = input_tensor_info.GetShape();

        if (_inputTensorShape[0] == -1)
        {
            _isDynamicShape = true;
            _inputTensorShape[0] = _batchSize;

        }
        if (_inputTensorShape[2] == -1 || _inputTensorShape[3] == -1) {
            _isDynamicShape = true;
            _inputTensorShape[2] = _netHeight;
            _inputTensorShape[3] = _netWidth;
        }
        //init output
        _outputNodesNum = _OrtSession->GetOutputCount();
#if ORT_API_VERSION < ORT_OLD_VISON
        _output_name0 = _OrtSession->GetOutputName(0, allocator);
        _outputNodeNames.push_back(_output_name0);
#else
        _output_name0 = std::move(_OrtSession->GetOutputNameAllocated(0, allocator));
        _outputNodeNames.push_back(_output_name0.get());
#endif
        Ort::TypeInfo type_info_output0(nullptr);
        type_info_output0 = _OrtSession->GetOutputTypeInfo(0);  //output0

        auto tensor_info_output0 = type_info_output0.GetTensorTypeAndShapeInfo();
        _outputNodeDataType = tensor_info_output0.GetElementType();
        _outputTensorShape = tensor_info_output0.GetShape();

        //_outputMaskNodeDataType = tensor_info_output1.GetElementType(); //the same as output0
        //_outputMaskTensorShape = tensor_info_output1.GetShape();
        //if (_outputTensorShape[0] == -1)
        //{
        //	_outputTensorShape[0] = _batchSize;
        //	_outputMaskTensorShape[0] = _batchSize;
        //}
        //if (_outputMaskTensorShape[2] == -1) {
        //	//size_t ouput_rows = 0;
        //	//for (int i = 0; i < _strideSize; ++i) {
        //	//	ouput_rows += 3 * (_netWidth / _netStride[i]) * _netHeight / _netStride[i];
        //	//}
        //	//_outputTensorShape[1] = ouput_rows;

        //	_outputMaskTensorShape[2] = _segHeight;
        //	_outputMaskTensorShape[3] = _segWidth;
        //}

        //warm up
        if (isCuda && warmUp) {
            //draw run
            cout << "Start warming up" << endl;
            size_t input_tensor_length = VectorProduct(_inputTensorShape);
            float* temp = new float[input_tensor_length];
            std::vector<Ort::Value> input_tensors;
            std::vector<Ort::Value> output_tensors;
            input_tensors.push_back(Ort::Value::CreateTensor<float>(
                _OrtMemoryInfo, temp, input_tensor_length, _inputTensorShape.data(),
                _inputTensorShape.size()));
            for (int i = 0; i < 3; ++i) {
                output_tensors = _OrtSession->Run(Ort::RunOptions{ nullptr },
                    _inputNodeNames.data(),
                    input_tensors.data(),
                    _inputNodeNames.size(),
                    _outputNodeNames.data(),
                    _outputNodeNames.size());
            }

            delete[]temp;
        }
    }
    catch (const std::exception&) {
        return false;
    }
    return true;

}

int Yolov8Onnx::Preprocessing(const std::vector<cv::Mat> &SrcImgs,
                              std::vector<cv::Mat> &OutSrcImgs,
                              std::vector<cv::Vec4d> &params){
    OutSrcImgs.clear();
    Size input_size = Size(_netWidth, _netHeight);

    // 信封处理
    for (size_t i=0; i<SrcImgs.size(); ++i){
        Mat temp_img = SrcImgs[i];
        Vec4d temp_param = {1,1,0,0};
        if (temp_img.size() != input_size){
            Mat borderImg;
            LetterBox(temp_img, borderImg, temp_param, input_size, false, false, true, 32);
            OutSrcImgs.push_back(borderImg);
            params.push_back(temp_param);
        }
        else {
            OutSrcImgs.push_back(temp_img);
            params.push_back(temp_param);
        }
    }

    int lack_num = _batchSize - SrcImgs.size();
    if (lack_num > 0){
        Mat temp_img = Mat::zeros(input_size, CV_8UC3);
        Vec4d temp_param = {1,1,0,0};
        OutSrcImgs.push_back(temp_img);
        params.push_back(temp_param);
    }
    return 0;
}

bool Yolov8Onnx::OnnxBatchDetect(std::vector<cv::Mat> &srcImgs, std::vector<std::vector<OutputPose>> &output)
{
    vector<Vec4d> params;
    vector<Mat> input_images;
    cv::Size input_size(_netWidth, _netHeight);

    //preprocessing (信封处理)
    Preprocessing(srcImgs, input_images, params);
    // [0~255] --> [0~1]; BGR2RGB
    Mat blob = cv::dnn::blobFromImages(input_images, 1 / 255.0, input_size, Scalar(0,0,0), true, false);

    // 前向传播得到推理结果
    int64_t input_tensor_length = VectorProduct(_inputTensorShape);// ?
    std::vector<Ort::Value> input_tensors;
    std::vector<Ort::Value> output_tensors;
    input_tensors.push_back(Ort::Value::CreateTensor<float>(_OrtMemoryInfo, (float*)blob.data,
                                                            input_tensor_length, _inputTensorShape.data(),
                                                            _inputTensorShape.size()));

    output_tensors = _OrtSession->Run(Ort::RunOptions{ nullptr },
        _inputNodeNames.data(),
        input_tensors.data(),
        _inputNodeNames.size(),
        _outputNodeNames.data(),
        _outputNodeNames.size()
    );

    //post-process

    float* all_data = output_tensors[0].GetTensorMutableData<float>(); // 第一张图片的输出

    _outputTensorShape = output_tensors[0].GetTensorTypeAndShapeInfo().GetShape(); // 一张图片输出的维度信息 [1, 84, 8400]

    int64_t one_output_length = VectorProduct(_outputTensorShape) / _outputTensorShape[0]; // 一张图片输出所占内存长度 8400*84

    for (int img_index = 0; img_index < srcImgs.size(); ++img_index){
        Mat output0 = Mat(Size((int)_outputTensorShape[2], (int)_outputTensorShape[1]), CV_32F, all_data).t(); // [1, 56 ,8400] -> [1, 8400, 56]

        all_data += one_output_length; //指针指向下一个图片的地址

        float* pdata = (float*)output0.data; // [classid,x,y,w,h,x,y,...21个点]
        int rows = output0.rows; // 预测框的数量 8400

        // 一张图片的预测框

        vector<float> confidences;
        vector<Rect> boxes;
        vector<int> labels;
        vector<vector<float>> kpss;
        for (int r=0; r<rows; ++r){

            // 得到人类别概率
            auto kps_ptr = pdata + 5;


            // 预测框坐标映射到原图上
            float score = pdata[4];
            if (score > _classThreshold){

                // rect [x,y,w,h]
                float x = (pdata[0] - params[img_index][2]) / params[img_index][0]; //x
                float y = (pdata[1] - params[img_index][3]) / params[img_index][1]; //y
                float w = pdata[2] / params[img_index][0]; //w
                float h = pdata[3] / params[img_index][1]; //h

                int left = MAX(int(x - 0.5 *w +0.5), 0);
                int top = MAX(int(y - 0.5*h + 0.5), 0);

                std::vector<float> kps;
                for (int k=0; k< 17; k++){
                    float kps_x = (*(kps_ptr + 3*k)   - params[img_index][2]) / params[img_index][0];
                    float kps_y = (*(kps_ptr + 3*k + 1)  - params[img_index][3]) / params[img_index][1];
                    float kps_s = *(kps_ptr + 3*k +2);

//                    cout << *(kps_ptr + 3*k) << endl;

                    kps.push_back(kps_x);
                    kps.push_back(kps_y);
                    kps.push_back(kps_s);
                }



                confidences.push_back(score);
                labels.push_back(0);
                kpss.push_back(kps);
                boxes.push_back(Rect(left, top, int(w + 0.5), int(h + 0.5)));
            }
            pdata += _anchorLength; //下一个预测框
        }

        // 对一张图的预测框执行NMS处理
        vector<int> nms_result;
        cv::dnn::NMSBoxes(boxes, confidences, _classThreshold, _nmsThrehold, nms_result); // 还需要classThreshold?

        // 对一张图片:依据NMS处理得到的索引,得到类别id、confidence、box,并置于结构体OutputDet的容器中
        vector<OutputPose> temp_output;
        for (size_t i=0; i<nms_result.size(); ++i){
            int idx = nms_result[i];
            OutputPose result;

            result.confidence = confidences[idx];
            result.box = boxes[idx];
            result.label = labels[idx];
            result.kps = kpss[idx];
            temp_output.push_back(result);
        }
        output.push_back(temp_output); // 多张图片的输出;添加一张图片的输出置于此容器中

    }
    if (output.size())
        return true;
    else
        return false;

}


bool Yolov8Onnx::OnnxDetect(cv::Mat &srcImg, std::vector<OutputPose> &output){
    vector<Mat> input_data = {srcImg};
    vector<vector<OutputPose>> temp_output;

    if(OnnxBatchDetect(input_data, temp_output)){
        output = temp_output[0];
        return true;
    }
    else return false;
}

main.cpp

#include <iostream>
#include <opencv2/opencv.hpp>
#include "detect.h"
#include <sys/time.h>

#include <vector>

using namespace std;
using namespace cv;
using namespace cv::dnn;

const std::vector<std::vector<unsigned int>> KPS_COLORS =
        {{0,   255, 0},
         {0,   255, 0},
         {0,   255, 0},
         {0,   255, 0},
         {0,   255, 0},
         {255, 128, 0},
         {255, 128, 0},
         {255, 128, 0},
         {255, 128, 0},
         {255, 128, 0},
         {255, 128, 0},
         {51,  153, 255},
         {51,  153, 255},
         {51,  153, 255},
         {51,  153, 255},
         {51,  153, 255},
         {51,  153, 255}};

const std::vector<std::vector<unsigned int>> SKELETON = {{16, 14},
                                                         {14, 12},
                                                         {17, 15},
                                                         {15, 13},
                                                         {12, 13},
                                                         {6,  12},
                                                         {7,  13},
                                                         {6,  7},
                                                         {6,  8},
                                                         {7,  9},
                                                         {8,  10},
                                                         {9,  11},
                                                         {2,  3},
                                                         {1,  2},
                                                         {1,  3},
                                                         {2,  4},
                                                         {3,  5},
                                                         {4,  6},
                                                         {5,  7}};

const std::vector<std::vector<unsigned int>> LIMB_COLORS = {{51,  153, 255},
                                                            {51,  153, 255},
                                                            {51,  153, 255},
                                                            {51,  153, 255},
                                                            {255, 51,  255},
                                                            {255, 51,  255},
                                                            {255, 51,  255},
                                                            {255, 128, 0},
                                                            {255, 128, 0},
                                                            {255, 128, 0},
                                                            {255, 128, 0},
                                                            {255, 128, 0},
                                                            {0,   255, 0},
                                                            {0,   255, 0},
                                                            {0,   255, 0},
                                                            {0,   255, 0},
                                                            {0,   255, 0},
                                                            {0,   255, 0},
                                                            {0,   255, 0}};



int main(){

    //  读取模型
    string detect_model_path = "/home/jason/PycharmProjects/pytorch_learn/yolo/ultralytics-main-yolov8/yolov8n-pose.onnx";
    Yolov8Onnx yolov8;
    if (yolov8.ReadModel(detect_model_path))
        cout << "read Net ok!\n";
    else {
        return -1;
    }



    VideoCapture capture;
    capture.open("/home/jason/work/01-img/fall-down3.mp4");
    if (capture.isOpened())
        cout << "read video ok!\n";
    else
        cout << "read video err!\n";
    int width = capture.get(CAP_PROP_FRAME_WIDTH);
    int height = capture.get(CAP_PROP_FRAME_HEIGHT);
    Size size1 = Size(width, height);
    double delay = 1000/capture.get(CAP_PROP_FPS);
    int frame_pos = 0;
    int frame_all = capture.get(CAP_PROP_FRAME_COUNT);

    VideoWriter writer;
    writer.open("/home/jason/work/01-img/fall-down-result.mp4", VideoWriter::fourcc('m', 'p', '4', 'v'),
                delay,size1);

    Mat frame;
    struct timeval t1, t2;
    double timeuse;
    while (1) {

        //
        capture>>frame;

        if (frame_pos == frame_all-1) break;


        // YOLOv8检测
        vector<OutputPose> result;
        gettimeofday(&t1, NULL);
        bool  find = yolov8.OnnxDetect(frame, result);
        gettimeofday(&t2, NULL);
        frame_pos+=1;
        printf("%d/%d:find %d person!\n",frame_pos, frame_all, (int)result.size());


        if(find)
        {
            DrawPred(frame, result, SKELETON, KPS_COLORS, LIMB_COLORS);
            }
        else {
            cout << "not find!\n";
        }

        timeuse = (t2.tv_sec - t1.tv_sec) + (double)(t2.tv_usec - t1.tv_usec)/1000000;
        timeuse *= 1000;
        string label = "TimeUse: " + to_string(timeuse);
        putText(frame, label, Point(30,30), FONT_HERSHEY_SIMPLEX, 1, Scalar(0,0,255), 2, 8);

        writer << frame;
        imshow("yolov8n-pose", frame);
        if(waitKey(1)=='q') break;

    }

    capture.release();
//    writer.release();


    return 0;
}

CmakeList.txt

cmake_minimum_required(VERSION 3.5)

project(05-YOLOv8-pose-onnruntime LANGUAGES CXX)

set(CMAKE_CXX_STANDARD 11)
set(CMAKE_CXX_STANDARD_REQUIRED ON)

include_directories("/home/jason/下载/onnxruntime-linux-x64-1.14.1/include")
#link_directories("/home/jason/下载/onnxruntime-linux-x64-1.14.1/lib")

include_directories(./include)
aux_source_directory(./src SOURCES)


find_package(OpenCV 4 REQUIRED)
add_executable(${PROJECT_NAME} ${SOURCES})

target_link_libraries(${PROJECT_NAME} ${OpenCV_LIBS})
target_link_libraries(${PROJECT_NAME} "/home/jason/下载/onnxruntime-linux-x64-1.14.1/lib/libonnxruntime.so")


参考:
Yolov8 姿态估计 – 知乎

YOLOv8-Pose 的 TensorRT8 推理尝试 – 知乎

文章出处登录后可见!

已经登录?立即刷新

共计人评分,平均

到目前为止还没有投票!成为第一位评论此文章。

(0)
心中带点小风骚的头像心中带点小风骚普通用户
上一篇 2023年6月6日
下一篇 2023年6月6日

相关推荐

此站出售,如需请站内私信或者邮箱!