superpoint, fixing slow convert after cnn (#659)

* fixing slow convert after cnn

* restoring cmake
This commit is contained in:
42Max42
2020-12-11 15:27:18 +01:00
committed by GitHub
parent da8e76ffed
commit 3271723034

View File

@@ -168,12 +168,16 @@ std::vector<cv::KeyPoint> SPDetector::detect(const cv::Mat &img, const cv::Mat &
auto kpts = (prob_ > threshold_); auto kpts = (prob_ > threshold_);
kpts = torch::nonzero(kpts); // [n_keypoints, 2] (y, x) kpts = torch::nonzero(kpts); // [n_keypoints, 2] (y, x)
//convert back to cpu if in gpu
auto kpts_cpu = kpts.to(torch::kCPU);
auto prob_cpu = prob_.to(torch::kCPU);
std::vector<cv::KeyPoint> keypoints_no_nms; std::vector<cv::KeyPoint> keypoints_no_nms;
for (int i = 0; i < kpts.size(0); i++) { for (int i = 0; i < kpts_cpu.size(0); i++) {
if(mask.empty() || mask.at<unsigned char>(kpts[i][0].item<int>(), kpts[i][1].item<int>()) != 0) if(mask.empty() || mask.at<unsigned char>(kpts_cpu[i][0].item<int>(), kpts_cpu[i][1].item<int>()) != 0)
{ {
float response = prob_[kpts[i][0]][kpts[i][1]].item<float>(); float response = prob_cpu[kpts_cpu[i][0]][kpts_cpu[i][1]].item<float>();
keypoints_no_nms.push_back(cv::KeyPoint(kpts[i][1].item<float>(), kpts[i][0].item<float>(), 8, -1, response)); keypoints_no_nms.push_back(cv::KeyPoint(kpts_cpu[i][1].item<float>(), kpts_cpu[i][0].item<float>(), 8, -1, response));
} }
} }
@@ -183,7 +187,7 @@ std::vector<cv::KeyPoint> SPDetector::detect(const cv::Mat &img, const cv::Mat &
for (size_t i = 0; i < keypoints_no_nms.size(); i++) { for (size_t i = 0; i < keypoints_no_nms.size(); i++) {
int x = keypoints_no_nms[i].pt.x; int x = keypoints_no_nms[i].pt.x;
int y = keypoints_no_nms[i].pt.y; int y = keypoints_no_nms[i].pt.y;
conf.at<float>(i, 0) = prob_[y][x].item<float>(); conf.at<float>(i, 0) = prob_cpu[y][x].item<float>();
} }
int border = 0; int border = 0;