| 50 | } |
| 51 | |
| 52 | static void NonMaxSuppressionSingleClasssImpl(const float* boxesPtr, const float* scores, int numBoxes, int maxDetections, |
| 53 | float iouThreshold, float scoreThreshold, std::vector<int32_t>* selected) { |
| 54 | MNN_ASSERT(iouThreshold >= 0.0f && iouThreshold <= 1.0f); |
| 55 | |
| 56 | const int outputNum = std::min(maxDetections, numBoxes); |
| 57 | std::vector<float> scoresData(numBoxes); |
| 58 | std::copy_n(scores, numBoxes, scoresData.begin()); |
| 59 | |
| 60 | struct Candidate { |
| 61 | int boxIndex; |
| 62 | float score; |
| 63 | }; |
| 64 | |
| 65 | auto cmp = [](const Candidate bsI, const Candidate bsJ) { return bsI.score < bsJ.score; }; |
| 66 | |
| 67 | std::priority_queue<Candidate, std::deque<Candidate>, decltype(cmp)> candidatePriorityQueue(cmp); |
| 68 | |
| 69 | for (int i = 0; i < scoresData.size(); ++i) { |
| 70 | if (scoresData[i] > scoreThreshold) { |
| 71 | candidatePriorityQueue.emplace(Candidate({i, scoresData[i]})); |
| 72 | } |
| 73 | } |
| 74 | |
| 75 | // std::vector<float> selectedScores; |
| 76 | Candidate nextCandidate; |
| 77 | float iou, originalScore; |
| 78 | |
| 79 | while (selected->size() < outputNum && !candidatePriorityQueue.empty()) { |
| 80 | nextCandidate = candidatePriorityQueue.top(); |
| 81 | originalScore = nextCandidate.score; |
| 82 | candidatePriorityQueue.pop(); |
| 83 | |
| 84 | // Overlapping boxes are likely to have similar scores, |
| 85 | // therefore we iterate through the previously selected boxes backwards |
| 86 | // in order to see if `next_candidate` should be suppressed. |
| 87 | bool shouldSelect = true; |
| 88 | for (int j = (int)selected->size() - 1; j >= 0; --j) { |
| 89 | iou = IOU(boxesPtr, nextCandidate.boxIndex, selected->at(j)); |
| 90 | if (iou == 0.0) { |
| 91 | continue; |
| 92 | } |
| 93 | if (iou > iouThreshold) { |
| 94 | shouldSelect = false; |
| 95 | } |
| 96 | } |
| 97 | |
| 98 | if (shouldSelect) { |
| 99 | selected->push_back(nextCandidate.boxIndex); |
| 100 | // selectedScores.push_back(nextCandidate.score); |
| 101 | } |
| 102 | } |
| 103 | } |
| 104 | |
| 105 | std::vector<Express::VARP> NMSModule::onForward(const std::vector<Express::VARP>& inputs) { |
| 106 | const int maxDetections = inputs[2]->readMap<int>()[0]; |