yolo_nms

最新推荐文章于 2024-09-30 10:59:24 发布

谢欣燕

最新推荐文章于 2024-09-30 10:59:24 发布

阅读量374

点赞数

分类专栏：笔记文章标签：计算机视觉

本文链接：https://blog.csdn.net/weixin_43941538/article/details/124265779

版权

笔记专栏收录该内容

70 篇文章 6 订阅

订阅专栏

该代码实现了一个用于目标检测的非极大值抑制（NMS）算法，通过计算两个边界框的IoU（交并比）来过滤掉重叠的检测结果。算法首先根据置信度对检测结果排序，然后逐步删除与当前最高置信度边界框IoU超过阈值的其他边界框。

摘要由CSDN通过智能技术生成

struct aligns(float) Detection
{
    float bbox[4];
    float conf;
    float class_id;
};
float iou(float lbox[4], float rbox[4])
{
    float inner_box[] = {
        (std::max)(lbox[0]-lbox[2]/2.f,rbox[0]-rbox[2]/2.f),//left
        (std::min)(lbox[0]+lbox[2]/2.f,rbox[0]+rbox[2]/2.f),//right
        (std::max)(lbox[1]-lbox[3]/2.f,rbox[1]-rbox[3]/2.f),//top
        (std::min)(lbox[1]+lbox[3]/2.f,rbox[1]+rbox[3]/2.f),//bottom
    };
    if(inner_box[2]>inner_box[3] || inner_box[0]>inner_box[1])
        return 0.0f;
    float interBoxS = (inner_box[1]-inner_box[0])*(inner_box[3]-inner_box[2]);
    return interBoxS/(lbox[2]*lbox[3] + rbox[2]*rbox[3] - interBoxS);
}
bool cmp(const Detection&a, const Detection& b)
{
    return a.conf>b.conf;
}
void nms(std::vector<Detection>&res, float* output, float conf_thresh, float nms_thresh=0.5)
{
    int det_size = sizeof(Detection)/sizeof (float);
    std::map<float, std::vector<Detection>> m;
    for(int i =0; i<output[0];i++)
    {
        if(output[1+det_size*i+4]<=conf_thresh)
            continue;
        Detection det;
        memcpy(&det, &output[1+det_size*i],det_size*sizeof (float));
        if(m.count(det.class_id)==0)
            m.emplace(det.class_id, std::vector<Detection>());
        m[det.class_id].push_back(det);
    }
    for(auto it = m.begin(); it!=m.end();it++)
    {
        auto& dets = it->second;
        std::sort(dets.begin(),dets.end(),cmp);
        for(size_t m=0;m<dets.size();++m)
        {
            auto& item = dets[m];
            res.push_back(item);
            for(size_t n=m+1;n<dets.size();++n)
            {
                if(iou(item.bbox, dets[n].bbox)>nms_thresh)
                {
                    dets.erase(dets.begin()+n);
                    --n;
                }
            }
        }
    }
}