【OpenCV】OpenCV4调用darknet yolov3模型进行目标检测

可以用自己训练的模型后缀名是.weights，对应的配置文件后缀名是.cfg。如果没有自己训练的模型可以到YOLO官网下载预训练好的模型。

自己训练模型可以参考darknet-YOLO系列博客。

int main()
{
    std::string modelPath = "yolov3.weights";
    std::string configPath = "yolov3.cfg";
    std::string labelPath = "coco.names";
    std::string imagePath = "test.jpg"
    int networkW = 416;
    int networkH = 416;
    float thresh = 0.5;
    float nms = 0.45;

    dnn::Net net= dnn::readNet(modelPath , configPath );//我用的OpenCV4.1.1版本，调用dnn::readNetFromDarknet会报错，调用dnn::readNet可以正常识别
    vector labels;
    int err = readClassNamesFromFile(labelPath , labels);

    std::vector outNames = net.getUnconnectedOutLayersNames();

    Mat srcImg = imread(imagePath);
    Mat inputBlob = dnn::blobFromImage(srcImg, 1.0/255, Size(networkW, networkH), Scalar(), false, false);//这里因为训练时对图像做了归一化，所以在推理的时候也要对图像进行归一化

    net.setInput(blob);
    vector outs;
    net.forward(outs, outNames);

    postprocess(srcImg, outs, net, thresh, nms);

    return 0;
}

int readClassNamesFromFile(String fileName, vector& classNames)
{
	ifstream fp(fileName, ios::in);
	if (!fp.is_open())
	{
		cout << "can not open file " << fileName << endl;
		return -1;
	}

	String name;
	while (!fp.eof())
	{
		getline(fp, name);
		if (name.length())
		{
			classNames.push_back(name);
		}
	}
	fp.close();
	return 0;
}

void postprocess(Mat& frame, const std::vector& outs, Net& net, float thresh, float nms)
{
    static std::vector outLayers = net.getUnconnectedOutLayers();
    static std::string outLayerType = net.getLayer(outLayers[0])->type;
    std::vector classIds;
    std::vector confidences;
    std::vector boxes;
    if (outLayerType == "Region")
    {
        for (size_t i = 0; i < outs.size(); ++i)
        {
            //网络输出的数据是一个NxC矩阵向量，N是检测到的目标数量，C的类别数 + 4
            //开始的4个数据是[center_x, center_y, width, height]
            float* data = (float*)outs[i].data;
            for (int j = 0; j < outs[i].rows; ++j, data += outs[i].cols)
            {
                Mat scores = outs[i].row(j).colRange(5, outs[i].cols);
                Point classIdPoint;
                double confidence;
                minMaxLoc(scores, 0, &confidence, 0, &classIdPoint);
                if (confidence > thresh)
                {
                    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(Rect(left, top, width, height));
                }
            }
        }
        std::vector indices;
        NMSBoxes(boxes, confidences, thresh, nms, indices);
        for (size_t i = 0; i < indices.size(); ++i)
        {
            int idx = indices[i];
            Rect box = boxes[idx];
            drawPrediction(classIds[idx], confidences[idx], box.x, box.y,
                 box.x + box.width, box.y + box.height, frame);
        }
    }

int drawPrediction(const vector& labelNames, int classId, float conf, int left, int top, int right, int bottom, Mat& img)
{
	rectangle(img, Point(left, top), Point(right, bottom), cv::Scalar(255, 0, 0), 2);


	if (labelNames.empty())
	{
		cout << "labelNames is empty!" << endl;
		return -1;
	}
	if (classId >= (int)labelNames.size())
	{
		cout << "classId is out of boundary!" << endl;
		return -1;
	}

	String label = labelNames[classId];
	int baseLine;
	Size labelSize = getTextSize(label, FONT_HERSHEY_SIMPLEX, 0.5, 1, &baseLine);
	top = max(top, labelSize.height);
	putText(img, label, Point(left, top), cv::FONT_HERSHEY_COMPLEX_SMALL, 1, cv::Scalar(255, 0, 0), 1);

	return 0;
}