Function NonMaxSuppression

samples/ObjectDetection/src/NonMaxSuppression.cpp:48–90

}

std::vector<int> NonMaxSuppression(DetectedObjects& inputDetections, float iouThresh)
{
    // Sort indicies of detections by highest score to lowest.
    std::vector<unsigned int> sortedIndicies = GenerateRangeK(inputDetections.size());
    std::sort(sortedIndicies.begin(), sortedIndicies.end(),
        [&inputDetections](int idx1, int idx2)
        {
            return inputDetections[idx1].GetScore() > inputDetections[idx2].GetScore();
        });

    std::vector<bool> visited(inputDetections.size(), false);
    std::vector<int> outputIndiciesAfterNMS;

    for (int i=0; i < inputDetections.size(); ++i)
    {
        // Each new unvisited detect should be kept.
        if (!visited[sortedIndicies[i]])
        {
            outputIndiciesAfterNMS.emplace_back(sortedIndicies[i]);
            visited[sortedIndicies[i]] = true;
        }

        // Look for detections to suppress.
        for (int j=i+1; j<inputDetections.size(); ++j)
        {
            // Skip if already kept or suppressed.
            if (!visited[sortedIndicies[j]])
            {
                // Detects must have the same label to be suppressed.
                if (inputDetections[sortedIndicies[j]].GetLabel() == inputDetections[sortedIndicies[i]].GetLabel())
                {
                    auto iou = IntersectionOverUnion(inputDetections[sortedIndicies[i]],
                                                    inputDetections[sortedIndicies[j]]);
                    if (iou > iouThresh)
                    {
                        visited[sortedIndicies[j]] = true;
                    }
                }
            }
        }
    }
    return outputIndiciesAfterNMS;
}

} // namespace od