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;
}
發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章