improve NMS implementation (#1173)

This commit is contained in:
Borong Yuan
2023-11-29 05:36:28 +08:00
committed by GitHub
parent 3a7f88dc6c
commit 0876603325

View File

@@ -2118,6 +2118,7 @@ void NMS(
cv::Mat inds = cv::Mat(cv::Size(img_width, img_height), CV_16UC1); cv::Mat inds = cv::Mat(cv::Size(img_width, img_height), CV_16UC1);
cv::Mat confidence = cv::Mat(cv::Size(img_width, img_height), CV_32FC1); cv::Mat confidence = cv::Mat(cv::Size(img_width, img_height), CV_32FC1);
cv::Mat dilated_conf = cv::Mat(cv::Size(img_width, img_height), CV_32FC1);
grid.setTo(0); grid.setTo(0);
inds.setTo(0); inds.setTo(0);
@@ -2134,11 +2135,8 @@ void NMS(
confidence.at<float>(vv, uu) = ptsIn[i].response; confidence.at<float>(vv, uu) = ptsIn[i].response;
} }
// debug cv::dilate(confidence, dilated_conf, cv::Mat());
//cv::Mat confidenceVis = confidence.clone() * 255; cv::Mat peaks = confidence == dilated_conf;
//confidenceVis.convertTo(confidenceVis, CV_8UC1);
//cv::imwrite("confidence.bmp", confidenceVis);
//cv::imwrite("grid_in.bmp", grid);
cv::copyMakeBorder(grid, grid, dist_thresh, dist_thresh, dist_thresh, dist_thresh, cv::BORDER_CONSTANT, 0); cv::copyMakeBorder(grid, grid, dist_thresh, dist_thresh, dist_thresh, dist_thresh, cv::BORDER_CONSTANT, 0);
@@ -2150,6 +2148,8 @@ void NMS(
float c = confidence.at<float>(vv-dist_thresh, uu-dist_thresh); float c = confidence.at<float>(vv-dist_thresh, uu-dist_thresh);
if (grid.at<unsigned char>(vv, uu) == 100) // If not yet suppressed. if (grid.at<unsigned char>(vv, uu) == 100) // If not yet suppressed.
{
if (peaks.at<unsigned char>(vv-dist_thresh, uu-dist_thresh) == 255)
{ {
for(int k = -dist_thresh; k < (dist_thresh+1); k++) for(int k = -dist_thresh; k < (dist_thresh+1); k++)
{ {
@@ -2159,13 +2159,16 @@ void NMS(
continue; continue;
if (confidence.at<float>(vv + k - dist_thresh, uu + j - dist_thresh) <= c) if (confidence.at<float>(vv + k - dist_thresh, uu + j - dist_thresh) <= c)
{
grid.at<unsigned char>(vv + k, uu + j) = 0; grid.at<unsigned char>(vv + k, uu + j) = 0;
} }
} }
}
grid.at<unsigned char>(vv, uu) = 255; grid.at<unsigned char>(vv, uu) = 255;
} }
else
{
grid.at<unsigned char>(vv, uu) = 0;
}
}
} }
size_t valid_cnt = 0; size_t valid_cnt = 0;
@@ -2173,9 +2176,6 @@ void NMS(
grid = cv::Mat(grid, cv::Rect(dist_thresh, dist_thresh, img_width, img_height)); grid = cv::Mat(grid, cv::Rect(dist_thresh, dist_thresh, img_width, img_height));
//debug
//cv::imwrite("grid_nms.bmp", grid);
for (int v = 0; v < img_height; v++) for (int v = 0; v < img_height; v++)
{ {
for (int u = 0; u < img_width; u++) for (int u = 0; u < img_width; u++)