yolo,c++目标识别

yolo,c++目标识别

• yolo 官网 https://pjreddie.com/darknet/yolo/
• yolo github https://github.com/pjreddie/darknet/

下载编译yolo源码

git clone https://github.com/pjreddie/darknet
cd darknet
make

下载预训练权值文件

wget https://pjreddie.com/media/files/yolov3.weights

预览效果

./darknet detect cfg/yolov3.cfg yolov3.weights data/dog.jpg

ubuntu 环境，确保系统安装opencv

sudo apt-get install libopencv-dev

CMakeLists.txt

cmake_minimum_required(VERSION 2.8)

enable_language(C CXX)
set(CXX_STANDARD 11)

project(test_yolo)

find_package(OpenCV REQUIRED)

set(TARGET test_yolo)

add_executable(${TARGET} main.cpp)
target_link_libraries(${TARGET} ${OpenCV_LIBS})

main.cpp

#include 
#include 
#include 
#include 
#include 
using namespace std;

/// @brief 
/// @param net 
/// @return 
std::vector<std::string> getOutputsNames(const cv::dnn::Net& net) {
    std::vector<std::string> names;
    std::vector<int> outLayers = net.getUnconnectedOutLayers();
    std::vector<cv::String> layerNames = net.getLayerNames();
    
    for (int i = 0; i < outLayers.size(); i++) {
        int layerIndex = outLayers[i] - 1;
        if (layerIndex >= 0 && layerIndex < layerNames.size()) {
            names.push_back(layerNames[layerIndex]);
        }
    }
    return names;
}

/// @brief 
/// @param frame 
/// @param outs 
/// @param classes 
void postprocess(cv::Mat& frame, const std::vector<cv::Mat>& outs, const std::vector<std::string>& classes) {
    float confThreshold = 0.5;  // You can adjust this threshold based on your needs
    std::vector<int> classIds;
    std::vector<float> confidences;
    std::vector<cv::Rect> boxes;

    // Define colors for different classes
    std::vector<cv::Scalar> classColors(classes.size());
    for (size_t i = 0; i < classes.size(); ++i) {
        // Generate random color for each class
        classColors[i] = cv::Scalar(rand() % 255, rand() % 255, rand() % 255);
    }

    for (size_t i = 0; i < outs.size(); ++i) {
        float* data = (float*)outs[i].data;
        for (int j = 0; j < outs[i].rows; ++j, data += outs[i].cols) {
            cv::Mat scores = outs[i].row(j).colRange(5, outs[i].cols);
            cv::Point classIdPoint;
            double confidence;
            cv::minMaxLoc(scores, 0, &confidence, 0, &classIdPoint);
            if (confidence > confThreshold) {
                int centerX = (int)(data[0] * frame.cols);
                int centerY = (int)(data[1] * frame.rows);
                int width = (int)(data[2] * frame.cols);
                int height = (int)(data[3] * frame.rows);
                int left = centerX - width / 2;
                int top = centerY - height / 2;

                classIds.push_back(classIdPoint.x);
                confidences.push_back((float)confidence);
                boxes.push_back(cv::Rect(left, top, width, height));
            }
        }
    }

    // Apply non-maximum suppression to filter overlapping boxes
    std::vector<int> indices;
    cv::dnn::NMSBoxes(boxes, confidences, confThreshold, 0.4, indices);

    for (size_t i = 0; i < indices.size(); ++i) {
        int idx = indices[i];
        cv::Rect box = boxes[idx];
        int classId = classIds[idx];
        float confidence = confidences[idx];

        // Draw bounding box with color for the specific class
        cv::rectangle(frame, box, classColors[classId], 2);

        // Display class name and confidence with color
        std::string label = classes[classId] + ": " + std::to_string(confidence);
        int baseline;
        cv::Size labelSize = cv::getTextSize(label, cv::FONT_HERSHEY_SIMPLEX, 0.5, 1, &baseline);
        cv::putText(frame, label, cv::Point(box.x, box.y - 5),
                    cv::FONT_HERSHEY_SIMPLEX, 0.5, classColors[classId], 1, cv::LINE_AA);
    }
}


int main() {

    //--- use you config ---
    std::string prefix = "/home/a/workspace/ai/darknet/" ;

    std::string cfg = prefix+"cfg/yolov3.cfg";
    std::string weights = prefix+"yolov3.weights" ;
    std::string classNames = prefix+"data/coco.names";
    std::string target_image = prefix + "data/dog.jpg";

    std::cout<<"prefix          : "<<prefix <<std::endl;
    std::cout<<"weights         : "<<weights <<std::endl;
    std::cout<<"classNames      : "<<classNames <<std::endl;
    std::cout<<"target_image    : "<<target_image <<std::endl;

    //--------------------------------------------------------------


    // 加载 YOLO 模型
    cv::dnn::Net net = cv::dnn::readNetFromDarknet(cfg, weights);
    if (net.empty()) {
        std::cerr << "Failed to load YOLO model." << std::endl;
        return -1;
    }

    // 加载类别文件
    // std::string c_names = prefix + "data/coco.names";
    std::ifstream classesFile(classNames);
    if (!classesFile.is_open()) {
        std::cerr << "Failed to open class names file." << std::endl;
        return -1;
    }

    std::vector<std::string> classes;
    std::string className;
    while (std::getline(classesFile, className)) {
        classes.push_back(className);
    }

    // 读取图像
    cv::Mat frame = cv::imread(target_image);
    if (frame.empty()) {
        std::cerr << "Failed to load image." << std::endl;
        return -1;
    }

    // YOLO 输入图像预处理
    cv::Mat blob = cv::dnn::blobFromImage(frame, 1 / 255.0, cv::Size(416, 416), cv::Scalar(0, 0, 0), true, false);
    net.setInput(blob);

    // 进行目标检测
    std::vector<cv::Mat> outs;
    net.forward(outs, getOutputsNames(net));

    // 后处理
    postprocess(frame, outs, classes);

    // 显示结果
    cv::imshow("Object Detection", frame);
    cv::waitKey(0);

    return 0;
}