Merge pull request #2410 from kartikdutt18/Add-NMS

Add Fast NMS.

HISTORY.md

@@ -17,6 +17,10 @@
 
   * Add Load and Save of Sparse Matrix (#2344).
 
+  * Add Intersection over Union (IoU) metric for bounding boxes (#2402).
+
+  * Add Non Maximal Supression (NMS) metric for bounding boxes (#2410).
+
   * Fix `no_intercept` and probability computation for linear SVM bindings
     (#2419).
 

src/mlpack/core/metrics/CMakeLists.txt

@@ -9,6 +9,8 @@ set(SOURCES
   lmetric_impl.hpp
   mahalanobis_distance.hpp
   mahalanobis_distance_impl.hpp
+  non_maximal_supression.hpp
+  non_maximal_supression_impl.hpp
 )
 
 # add directory name to sources

src/mlpack/core/metrics/iou_metric.hpp

@@ -1,5 +1,5 @@
 /**
- * @file iou_metric.hpp
+ * @file core/metrics/iou_metric.hpp
  * @author Kartik Dutt
  *
  * Definition of Intersection Over Union metric. It is defined as intersection

src/mlpack/core/metrics/iou_metric_impl.hpp

@@ -1,5 +1,5 @@
 /**
- * @file iou_metric_impl.hpp
+ * @file core/metrics/iou_metric_impl.hpp
  * @author Kartik Dutt
  *
  * Implementation of Intersection Over Union metric.
@@ -70,7 +70,7 @@ void IoU<UseCoordinates>::serialize(
     Archive& ar,
     const unsigned int /* version */)
 {
-  ar & BOOST_SERIALIZATION_NVP(useCoordinates);
+  // Nothing to do here.
 }
 
 } // namespace metric

src/mlpack/core/metrics/non_maximal_supression.hpp

@@ -0,0 +1,86 @@
+/**
+ * @file core/metrics/non_maximal_supression.hpp
+ * @author Kartik Dutt
+ *
+ * Definition of Non Maximal Supression metric.
+ *
+ * mlpack is free software; you may redistribute it and/or modify it under the
+ * terms of the 3-clause BSD license.  You should have received a copy of the
+ * 3-clause BSD license along with mlpack.  If not, see
+ * http://www.opensource.org/licenses/BSD-3-Clause for more information.
+ */
+#ifndef MLPACK_CORE_METRICS_NMS_HPP
+#define MLPACK_CORE_METRICS_NMS_HPP
+
+namespace mlpack {
+namespace metric {
+
+/**
+ * Definition of Non Maximal Supression.
+ *
+ * Performs non-maximal suppression (NMS) on the boxes according to their
+ * Intersection-over-Union (IoU). NMS iteratively removes lower scoring boxes
+ * which have an IoU greater than threshold with another high scoring box.
+ *
+ * For bounding box representation there are two common representation
+ * either as coordinates i.e. each value in vector represents a
+ * coordinate in the format x0, y0, x1, y1 where x0, y0 represent the
+ * lower left coordinate and x1, y1, represent upper right coordinate.
+ * 
+ * Second representation follows the following representation : x0, y0, h, w.
+ * Where x0 and y0 are bottom left bounding box coordinates and h, w are
+ * height and width of the bounding box.
+ *
+ * @tparam UseCoordinates Toggles between the two representation of bounding box.
+ *                        If true, each value in vector represents a coordinate 
+ *                        in the formate x0, y0, x1, y1. Else the bounding box is
+ *                        represented as x0, y0, h, w.
+ */
+template<bool UseCoordinates = false>
+class NMS
+{
+ public:
+  //! Default constructor required to satisfy the Metric policy.
+  NMS() { /* Nothing to do here. */ }
+
+  /**
+   * Performs non-maximal suppression.
+   *
+   * @param boundingBoxes Column major representation of bounding boxes
+   *                      i.e. Each column corresponds to a different bounding
+   *                      box. Each bounding box should contain 4 points only
+   *                      either {x1, y1, x2, y2} or {x1, y1, h, w} depending
+   *                      on UseCoordinates parameter.
+   * @param confidenceScores Vector containing confidence score corresponding
+   *                         to each bounding box.
+   * @param selectedIndices Output of Non Maximal Suppression (NMS) is stored
+   *                        here. It contains a list of indices corresponding
+   *                        to bounding boxes in input parameter, sorted
+   *                        in descending order of the confidence scores.
+   * @param threshold Threshold used to discard all overlapping bounding boxes
+   *                  that have IoU greater than the threshold.
+   */
+  template<
+      typename BoundingBoxesType,
+      typename ConfidenceScoreType,
+      typename OutputType
+  >
+  static void Evaluate(const BoundingBoxesType& boundingBoxes,
+                       const ConfidenceScoreType& confidenceScores,
+                       OutputType& selectedIndices,
+                       const double threshold = 0.5);
+
+  static const bool useCoordinates = UseCoordinates;
+
+  //! Serialize the metric.
+  template <typename Archive>
+  void serialize(Archive &ar, const unsigned int /* version */);
+}; // Class NMS.
+
+} // namespace metric
+} // namespace mlpack
+
+// Include implementation.
+#include "non_maximal_supression_impl.hpp"
+
+#endif

src/mlpack/core/metrics/non_maximal_supression_impl.hpp

@@ -0,0 +1,138 @@
+/**
+ * @file core/metrics/nms_metric_impl.hpp
+ * @author Kartik Dutt
+ *
+ * Implementation of Non Maximal Supression metric.
+ *
+ * mlpack is free software; you may redistribute it and/or modify it under the
+ * terms of the 3-clause BSD license.  You should have received a copy of the
+ * 3-clause BSD license along with mlpack.  If not, see
+ * http://www.opensource.org/licenses/BSD-3-Clause for more information.
+ */
+#ifndef MLPACK_CORE_METRICS_NMS_IMPL_HPP
+#define MLPACK_CORE_METRICS_NMS_IMPL_HPP
+
+// In case it hasn't been included.
+#include "non_maximal_supression.hpp"
+
+namespace mlpack {
+namespace metric {
+
+template<bool UseCoordinates>
+template<
+    typename BoundingBoxesType,
+    typename ConfidenceScoreType,
+    typename OutputType
+>
+void NMS<UseCoordinates>::Evaluate(
+    const BoundingBoxesType& boundingBoxes,
+    const ConfidenceScoreType& confidenceScores,
+    OutputType& selectedIndices,
+    const double threshold)
+{
+  Log::Assert(boundingBoxes.n_rows == 4, "Bounding boxes must \
+      contain only 4 rows determining coordinates of bounding \
+      box either in {x1, y1, x2, y2} or {x1, y1, h, w} format.\
+      Refer to the documentation for more information.");
+
+  Log::Assert(confidenceScores.n_cols != boundingBoxes.n_cols, "Each \
+      bounding box must correspond to atleast and only 1 bounding box. \
+      Found " + std::to_string(confidenceScores.n_cols) + " confidence \
+      scores for " + std::to_string(boundingBoxes.n_cols) + " bounding boxes.");
+
+  // Clear selected bounding boxes.
+  selectedIndices.clear();
+
+  // Obtain Sorted indices for bounding boxes according to
+  // their confidence scores.
+  arma::ucolvec sortedIndices = arma::sort_index(confidenceScores);
+
+  // Pre-Compute area of each bounding box.
+  arma::mat area;
+  if (UseCoordinates)
+  {
+    area = (boundingBoxes.row(2) - boundingBoxes.row(0)) %
+        (boundingBoxes.row(3) - boundingBoxes.row(1));
+  }
+  else
+  {
+    area = boundingBoxes.row(2) % boundingBoxes.row(3);
+  }
+
+  while (sortedIndices.n_elem > 0)
+  {
+    size_t selectedIndex = sortedIndices(sortedIndices.n_elem - 1);
+
+    // Choose the box with the largest probability.
+    selectedIndices.insert_rows(0, arma::uvec(1).fill(selectedIndex));
+
+    // Check if there are other bounding boxes to compare with.
+    if (sortedIndices.n_elem == 1)
+    {
+      break;
+    }
+
+    // Remove the last index.
+    sortedIndices = sortedIndices(arma::span(0, sortedIndices.n_rows - 2),
+        arma::span());
+
+    // Get x and y coordinates for remaining bounding boxes.
+    BoundingBoxesType x2 = boundingBoxes.submat(arma::uvec(1).fill(2),
+        sortedIndices);
+
+    BoundingBoxesType x1 = boundingBoxes.submat(arma::uvec(1).fill(0),
+        sortedIndices);;
+
+    BoundingBoxesType y2 = boundingBoxes.submat(arma::uvec(1).fill(3),
+        sortedIndices);
+
+    BoundingBoxesType y1 = boundingBoxes.submat(arma::uvec(1).fill(1),
+        sortedIndices);
+
+    size_t selectedX2 = boundingBoxes(2, selectedIndex);
+    size_t selectedY2 = boundingBoxes(3, selectedIndex);
+    size_t selectedX1 = boundingBoxes(0, selectedIndex);
+    size_t selectedY1 = boundingBoxes(1, selectedIndex);
+
+    if (!UseCoordinates)
+    {
+      // Change height - width representation to coordinate represention.
+      x2 = x2 + x1;
+      y2 = y2 + y1;
+      selectedX2 = selectedX2 + selectedX1;
+      selectedY2 = selectedY2 + selectedY1;
+    }
+
+    // Calculate points of intersection between the bounding box with
+    // highest confidence score and remaining bounding boxes.
+    x2 = arma::clamp(x2, DBL_MIN, selectedX2);
+    y2 = arma::clamp(y2, DBL_MIN, selectedY2);
+    x1 = arma::clamp(x1, selectedX1, DBL_MAX);
+    y1 = arma::clamp(y1, selectedY1, DBL_MAX);
+
+    BoundingBoxesType intersectionArea = arma::clamp(x2 - x1, 0.0, DBL_MAX) %
+          arma::clamp(y2 - y1, 0.0, DBL_MAX);
+
+    // Calculate IoU of remaining boxes with the last bounding box with
+    // the highest confidence score.
+    BoundingBoxesType calculateIoU = intersectionArea /
+        (area(sortedIndices).t() - intersectionArea + area(selectedIndex));
+
+    sortedIndices = sortedIndices(arma::find(calculateIoU <= threshold));
+  }
+
+  selectedIndices = arma::flipud(selectedIndices);
+}
+
+template<bool UseCoordinates>
+template<typename Archive>
+void NMS<UseCoordinates>::serialize(
+    Archive& ar,
+    const unsigned int /* version */)
+{
+  // Nothing to do here.
+}
+
+} // namespace metric
+} // namespace mlpack
+#endif

src/mlpack/tests/metric_test.cpp

@@ -12,6 +12,7 @@
 #include <mlpack/core/metrics/lmetric.hpp>
 #include <boost/test/unit_test.hpp>
 #include <mlpack/core/metrics/iou_metric.hpp>
+#include <mlpack/core/metrics/non_maximal_supression.hpp>
 #include "test_tools.hpp"
 
 using namespace std;
@@ -133,4 +134,171 @@ BOOST_AUTO_TEST_CASE(IoUMetricTest)
   BOOST_REQUIRE_CLOSE(IoU<>::Evaluate(bbox1, bbox2), 0.7309670, 1e-4);
 }
 
+BOOST_AUTO_TEST_CASE(NMSMetricTest)
+{
+  arma::mat bbox, selectedBoundingBox, desiredBoundingBox;
+  arma::vec bbox1(4), bbox2(4), bbox3(4);
+  arma::uvec selectedIndices, desiredIndices;
+
+  // Set values of each bounding box.
+  // Use coordinate system to represent bounding boxes.
+  // Bounding boxes represent {x0, y0, x1, y1}.
+  bbox1 << 0.5 << 0.5 << 41.0 << 31.0;
+  bbox2 << 1.0 << 1.0 << 42.0 << 22.0;
+  bbox3 << 10.0 << 13.0 << 90.0 << 100.0;
+
+  // Fill bounding box.
+  bbox.insert_cols(0, bbox3);
+  bbox.insert_cols(0, bbox2);
+  bbox.insert_cols(0, bbox1);
+
+  // Fill confidence scores for each bounding box.
+  arma::vec confidenceScores(3);
+  confidenceScores << 0.7 << 0.6 << 0.4;
+
+  // Selected bounding box using torchvision.ops.nms().
+  desiredBoundingBox.insert_cols(0, bbox3);
+  desiredBoundingBox.insert_cols(0, bbox1);
+
+  // Selected indices of bounding boxes using
+  // torchvision.ops.nms().
+  desiredIndices = arma::ucolvec(2);
+  desiredIndices << 0 << 2;
+
+  // Evaluate the bounding box.
+  NMS<true>::Evaluate(bbox, confidenceScores,
+      selectedIndices);
+
+  selectedBoundingBox = bbox.cols(selectedIndices);
+
+  BOOST_REQUIRE_EQUAL(selectedBoundingBox.n_cols, 2);
+  BOOST_REQUIRE_EQUAL(selectedBoundingBox.n_rows, 4);
+  CheckMatrices(desiredBoundingBox, selectedBoundingBox);
+
+  for (size_t i = 0; i < desiredIndices.n_elem; i++)
+  {
+    BOOST_REQUIRE_EQUAL(desiredIndices[i], selectedIndices[i]);
+  }
+
+  // Clean up.
+  bbox.clear();
+  desiredBoundingBox.clear();
+  selectedBoundingBox.clear();
+
+  // Fill new bounding boxes.
+  bbox.insert_cols(0, bbox1);
+  bbox.insert_cols(0, bbox2);
+  bbox.insert_cols(0, bbox1);
+  confidenceScores << 1.0 << 0.6 << 0.9;
+
+  // Output calculated using using torchvision.ops.nms().
+  desiredBoundingBox.insert_cols(0, bbox2);
+  desiredBoundingBox.insert_cols(0, bbox1);
+
+  NMS<true>::Evaluate(bbox, confidenceScores,
+      selectedIndices, 0.9);
+
+  selectedBoundingBox = bbox.cols(selectedIndices);
+
+  BOOST_REQUIRE_EQUAL(selectedBoundingBox.n_cols, 2);
+  BOOST_REQUIRE_EQUAL(selectedBoundingBox.n_rows, 4);
+  CheckMatrices(desiredBoundingBox, selectedBoundingBox);
+
+  // Clean up.
+  bbox.clear();
+  desiredBoundingBox.clear();
+  selectedBoundingBox.clear();
+
+  // Use coordinate system to represent bounding boxes.
+  // Bounding boxes represent {x0, y0, x1, y1}.
+  bbox1 << 39 << 63 << 203 << 112;
+  bbox2 << 31 << 69 << 201 << 125;
+  bbox3 << 54 << 66 << 198 << 114;
+
+  // Fill bounding box.
+  bbox.insert_cols(0, bbox3);
+  bbox.insert_cols(0, bbox2);
+  bbox.insert_cols(0, bbox1);
+
+  // Fill confidence scores of bounding boxes.
+  confidenceScores << 1.0 << 0.6 << 0.9;
+
+  // Selected bounding box using torchvision.ops.nms().
+  desiredBoundingBox.insert_cols(0, bbox2);
+  desiredBoundingBox.insert_cols(0, bbox1);
+
+  NMS<true>::Evaluate(bbox, confidenceScores,
+      selectedIndices, 0.7);
+
+  selectedBoundingBox = bbox.cols(selectedIndices);
+
+  BOOST_REQUIRE_EQUAL(selectedBoundingBox.n_cols, 2);
+  BOOST_REQUIRE_EQUAL(selectedBoundingBox.n_rows, 4);
+  CheckMatrices(desiredBoundingBox, selectedBoundingBox);
+
+  // Clean up.
+  bbox.clear();
+  desiredBoundingBox.clear();
+  selectedBoundingBox.clear();
+
+  // Set values of each bounding box.
+  // Use coordinate system to represent bounding boxes.
+  // Bounding boxes represent {x0, y0, h, w}.
+  bbox1 << 0.0 << 0.0 << 41.0 << 31.0;
+  bbox2 << 1.0 << 1.0 << 41.0 << 21.0;
+  bbox3 << 10.0 << 13.0 << 80.0 << 87.0;
+
+  // Fill bounding box.
+  bbox.insert_cols(0, bbox3);
+  bbox.insert_cols(0, bbox2);
+  bbox.insert_cols(0, bbox1);
+
+  // Fill confidence scores for each bounding box.
+  confidenceScores << 0.7 << 0.6 << 0.4;
+
+  // Selected bounding box using torchvision.ops.nms().
+  desiredBoundingBox.insert_cols(0, bbox3);
+  desiredBoundingBox.insert_cols(0, bbox1);
+
+  // Evaluate the bounding box.
+  NMS<>::Evaluate(bbox, confidenceScores,
+    selectedIndices);
+
+  selectedBoundingBox = bbox.cols(selectedIndices);
+  BOOST_REQUIRE_EQUAL(selectedBoundingBox.n_cols, 2);
+  BOOST_REQUIRE_EQUAL(selectedBoundingBox.n_rows, 4);
+  CheckMatrices(desiredBoundingBox, selectedBoundingBox);
+
+  // Clean up.
+  bbox.clear();
+  desiredBoundingBox.clear();
+  selectedBoundingBox.clear();
+
+  // Use coordinate system to represent bounding boxes.
+  // Bounding boxes represent {x0, y0, h, w}.
+  bbox1 << 39 << 63 << 164 << 49;
+  bbox2 << 31 << 69 << 170 << 56;
+  bbox3 << 54 << 66 << 144 << 48;
+
+  // Fill bounding box.
+  bbox.insert_cols(0, bbox3);
+  bbox.insert_cols(0, bbox2);
+  bbox.insert_cols(0, bbox1);
+
+  // Fill confidence scores of bounding boxes.
+  confidenceScores << 1.0 << 0.6 << 0.4;
+
+  // Selected bounding box using torchvision.ops.nms().
+  desiredBoundingBox.insert_cols(0, bbox2);
+  desiredBoundingBox.insert_cols(0, bbox1);
+
+  NMS<false>::Evaluate(bbox, confidenceScores,
+    selectedIndices, 0.7);
+
+  selectedBoundingBox = bbox.cols(selectedIndices);
+  BOOST_REQUIRE_EQUAL(selectedBoundingBox.n_cols, 2);
+  BOOST_REQUIRE_EQUAL(selectedBoundingBox.n_rows, 4);
+  CheckMatrices(desiredBoundingBox, selectedBoundingBox);
+}
+
 BOOST_AUTO_TEST_SUITE_END();