C++实现非极大抑制(NMS)算法

1、NMS算法

1.1、算法原理

1、将同一类的检测结果按照得分排序。
2、计算得分最高的检测框与其他检测框的重叠度(IOU),删除大于设定的重叠度阈值的检测框。
3、对于小于重叠度阈值的检测框重复1、2的操作,直到遍历完所有的检测框

1.2、算法实现(只有一类)

std::vector<int> cpu_nms(Eigen::Matrix<float, Eigen::Dynamic, Eigen::Dynamic>&bboxes, tensorflow::Tensor &scores, float nms_thresh)
{

	auto m_scores = Eigen::Map<Eigen::Matrix<float, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>>(scores.flat<float>().data(), 300, 2);

	Eigen::VectorXf class_vec = m_scores.col(1);

	//按照得分排序
	Eigen::VectorXi score_order(300);
	argsort(class_vec, score_order);

	Eigen::VectorXf x1 = bboxes.col(4);
	Eigen::VectorXf y1 = bboxes.col(5);
	Eigen::VectorXf x2 = bboxes.col(6);
	Eigen::VectorXf y2 = bboxes.col(7);

	Eigen::VectorXf offset(300);
	offset.setConstant(1.0);

	std::vector<int> keep;

	//计算bounding box面积
	Eigen::VectorXf areas = (x2 - x1 + offset).cwiseProduct(y2 - y1 + offset);

	while (score_order.size()>0)
	{
		//std::cout << "order:" << score_order.transpose() << std::endl;
		int i = score_order[0];
		keep.push_back(i);
		//std::cout <<"keep:"<< i << std::endl;
		int order_size = score_order.size();
		//计算得分最高的检测框与其他检测框的重合区域面积
		Eigen::VectorXf xx1 = x1(score_order.segment(1, order_size - 1)).cwiseMax(x1[i]);
		Eigen::VectorXf yy1 = y1(score_order.segment(1, order_size - 1)).cwiseMax(y1[i]);
		Eigen::VectorXf xx2 = x2(score_order.segment(1, order_size - 1)).cwiseMin(x2[i]);
		Eigen::VectorXf yy2 = y2(score_order.segment(1, order_size - 1)).cwiseMin(y2[i]);

		Eigen::VectorXf w = (xx2 - xx1 + Eigen::VectorXf::Ones(order_size - 1)).cwiseMax(0);
		Eigen::VectorXf h = (yy2 - yy1 + Eigen::VectorXf::Ones(order_size - 1)).cwiseMax(0);
		Eigen::VectorXf inter = w.cwiseProduct(h);

		Eigen::VectorXf area_score_max(order_size - 1);
		area_score_max.setConstant(areas[i]);
		
		//计算得分最高的检测框与其他检测框的面积和
		Eigen::VectorXf ovr = inter.cwiseQuotient(area_score_max + areas(score_order.segment(1, order_size - 1)) - inter);
		
		Eigen::VectorXi index_cond = (ovr.array() < nms_thresh).cast<int>();
		//std::cout << "index_cond:" << index_cond.transpose() << std::endl;
		int cond_sum = index_cond.sum();

		Eigen::VectorXi inds(cond_sum);
		inds.setZero();
		select_where(index_cond, inds);
		//std::cout << "inds:" << inds.transpose() << std::endl;
		Eigen::VectorXi offset2(cond_sum);
		offset2.setConstant(1);

		Eigen::VectorXi selcet_elem = score_order(inds+offset2);
		score_order = selcet_elem;
		//std::cout << "selcet_elem:" << selcet_elem.transpose() << std::endl;
	}
	return keep;
}
發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章