c++实现NMS

使用C++和OpenCV库实现nms, 等我闲的时候再来优化一下

#include 
#include 
#include 
#include 

struct BoundingBox
{
    cv::Rect box;
    float confidence;
    int index;
};

static bool sort_score(BoundingBox bbox1, BoundingBox bbox2)
{
    return bbox1.confidence > bbox2.confidence ? true : false;
}

static float get_iou_value(cv::Rect rect1, cv::Rect rect2)
{
    int xx1, yy1, xx2, yy2;
    xx1 = std::max(rect1.x, rect2.x);
    yy1 = std::max(rect1.y, rect2.y);
    xx2 = std::min(rect1.x + rect1.width - 1, rect2.x + rect2.width - 1);
    yy2 = std::min(rect1.y + rect1.height - 1, rect2.y + rect2.height - 1);
    int inter_width, inter_height;
    inter_width = std::max(0, xx2 - xx1 + 1);
    inter_height = std::max(0, yy2 - yy1 + 1);
    float inter_area, iou, union_area;
    inter_area = float(inter_width) * inter_height;
    union_area = float(rect1.width * rect1.height + rect2.width * rect2.height - inter_area);
    iou = inter_area / union_area;
    return iou;
}

void nms_boxes(std::vector<cv::Rect> &boxes, std::vector<float> confidences_, float confthreshold, float nmsthreshold, std::vector<int> indices)
{
    BoundingBox bbox;
    std::vector<BoundingBox> bboxes;
    int i, j;
    for (i = 0; i < boxes.size(); i++)
    {
        bbox.box = boxes[i];
        bbox.confidence = confidences_[i];
        bbox.index = i;
        bboxes.push_back(bbox);
    }
    std::sort(bboxes.begin(), bboxes.end(), sort_score);
    int k = bboxes.size();

    for (i = 0; i < k; i++)
    {
        if (bboxes[i].confidence < confthreshold)
            continue;
        indices.push_back(bboxes[i].index);
        for (j = i + 1; j < k; j++)
        {
            float iou = get_iou_value(bboxes[i].box, bboxes[j].box);
            // std::cout << "iou:" << iou << std::endl;
            if (iou > nmsthreshold)
            {
                bboxes.erase(bboxes.begin() + j);
                k = bboxes.size();
            }
        }
    }
    for (int i = 0; i < bboxes.size(); i++)
    {
        std::cout << "bboxes:" << bboxes[i].box << std::endl;
    }
}

int main()
{
    std::vector<cv::Rect> input;
    input.push_back(cv::Rect(0, 0, 100, 100));
    input.push_back(cv::Rect(0, 0, 50, 50));
    input.push_back(cv::Rect(20, 20, 50, 50));
    input.push_back(cv::Rect(2, 2, 100, 100));
    input.push_back(cv::Rect(200, 200, 100, 50));

    std::vector<float> confidences;
    confidences.push_back(0.9f);
    confidences.push_back(0.8f);
    confidences.push_back(0.7f);
    confidences.push_back(0.6f);
    confidences.push_back(0.5f);

    float nmsthresh = 0.4f;
    float confthresh = 0.7f;
    std::vector<cv::Rect> output;
    std::vector<int> indice;
    nms_boxes(input, confidences, confthresh, nmsthresh, indice);
}

你可能感兴趣的:(深度学习,目标检测,c++,c++,目标检测,人工智能)