From 3a5859a6dfc9d2502a8afcec928332f271272668 Mon Sep 17 00:00:00 2001 From: Borong Yuan Date: Mon, 1 Jan 2024 19:39:29 +0800 Subject: [PATCH] reimplement NMS using morphological operations --- corelib/src/camera/CameraDepthAI.cpp | 44 ++++++++++++++++++---------- 1 file changed, 29 insertions(+), 15 deletions(-) diff --git a/corelib/src/camera/CameraDepthAI.cpp b/corelib/src/camera/CameraDepthAI.cpp index 0d475fc8a9..8162c14ebf 100644 --- a/corelib/src/camera/CameraDepthAI.cpp +++ b/corelib/src/camera/CameraDepthAI.cpp @@ -679,25 +679,39 @@ SensorData CameraDepthAI::captureImage(CameraInfo * info) auto heatmap = features->getLayerFp16("heatmap"); auto desc = features->getLayerFp16("desc"); - cv::Mat prob(200, 320, CV_32FC1, heatmap.data()); - cv::resize(prob, prob, targetSize_, 0, 0, cv::INTER_CUBIC); - std::vector kpts; - cv::findNonZero(prob > threshold_, kpts); - std::vector keypoints_no_nms, keypoints; - for(auto& kpt : kpts) - { - float response = prob.at(kpt); - keypoints_no_nms.emplace_back(cv::KeyPoint(kpt, 8, -1, response)); - } + cv::Mat scores(200, 320, CV_32FC1, heatmap.data()); + cv::resize(scores, scores, targetSize_, 0, 0, cv::INTER_CUBIC); - if(nms_ && !keypoints_no_nms.empty()) + if(nms_) { - cv::Mat descEmpty; - util2d::NMS(keypoints_no_nms, descEmpty, keypoints, descEmpty, 0, nmsRadius_, targetSize_.width, targetSize_.height); + cv::Mat dilated_scores(targetSize_, CV_32FC1); + cv::dilate(scores, dilated_scores, cv::getStructuringElement(cv::MORPH_RECT, cv::Size(nmsRadius_*2+1, nmsRadius_*2+1))); + cv::Mat max_mask = scores == dilated_scores; + cv::dilate(scores, dilated_scores, cv::Mat()); + cv::Mat max_mask_r1 = scores == dilated_scores; + cv::Mat supp_mask(targetSize_, CV_8UC1); + for(size_t i=0; i<2; i++) + { + cv::dilate(max_mask, supp_mask, cv::getStructuringElement(cv::MORPH_RECT, cv::Size(nmsRadius_*2+1, nmsRadius_*2+1))); + cv::Mat supp_scores = scores.clone(); + supp_scores.setTo(0, supp_mask); + cv::dilate(supp_scores, dilated_scores, cv::getStructuringElement(cv::MORPH_RECT, cv::Size(nmsRadius_*2+1, nmsRadius_*2+1))); + cv::Mat new_max_mask = cv::Mat::zeros(targetSize_, CV_8UC1); + cv::bitwise_not(supp_mask, supp_mask); + cv::bitwise_and(supp_scores == dilated_scores, supp_mask, new_max_mask, max_mask_r1); + cv::bitwise_or(max_mask, new_max_mask, max_mask); + } + cv::bitwise_not(max_mask, supp_mask); + scores.setTo(0, supp_mask); } - else if(!keypoints_no_nms.empty()) + + std::vector kpts; + cv::findNonZero(scores > threshold_, kpts); + std::vector keypoints; + for(auto& kpt : kpts) { - keypoints = keypoints_no_nms; + float response = scores.at(kpt); + keypoints.emplace_back(cv::KeyPoint(kpt, 8, -1, response)); } cv::Mat coarse_desc(25, 40, CV_32FC(256), desc.data());