Merge pull request #8943 from kmindspark:detr-push-2

PiperOrigin-RevId: 324623313

research/object_detection/matchers/hungarian_matcher.py

+# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""Hungarian bipartite matcher implementation."""
+
+import numpy as np
+from scipy.optimize import linear_sum_assignment
+
+import tensorflow.compat.v1 as tf
+from object_detection.core import matcher
+
+
+class HungarianBipartiteMatcher(matcher.Matcher):
+  """Wraps a Hungarian bipartite matcher into TensorFlow."""
+
+  def _match(self, similarity_matrix, valid_rows):
+    """Optimally bipartite matches a collection rows and columns.
+
+    Args:
+      similarity_matrix: Float tensor of shape [N, M] with pairwise similarity
+        where higher values mean more similar.
+      valid_rows: A boolean tensor of shape [N] indicating the rows that are
+        valid.
+
+    Returns:
+      match_results: int32 tensor of shape [M] with match_results[i]=-1
+        meaning that column i is not matched and otherwise that it is matched to
+        row match_results[i].
+    """
+    valid_row_sim_matrix = tf.gather(similarity_matrix,
+                                     tf.squeeze(tf.where(valid_rows), axis=-1))
+    distance_matrix = -1 * valid_row_sim_matrix
+
+    def numpy_wrapper(inputs):
+      def numpy_matching(input_matrix):
+        row_indices, col_indices = linear_sum_assignment(input_matrix)
+        match_results = np.full(input_matrix.shape[1], -1)
+        match_results[col_indices] = row_indices
+        return match_results.astype(np.int32)
+
+      return tf.numpy_function(numpy_matching, inputs, Tout=[tf.int32])
+
+    matching_result = tf.autograph.experimental.do_not_convert(
+        numpy_wrapper)([distance_matrix])
+
+    return tf.reshape(matching_result, [-1])

research/object_detection/matchers/hungarian_matcher_tf2_test.py

+# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""Tests for object_detection.core.bipartite_matcher."""
+import unittest
+import numpy as np
+import tensorflow.compat.v1 as tf
+
+from object_detection.utils import test_case
+from object_detection.utils import tf_version
+
+if tf_version.is_tf2():
+  from object_detection.matchers import hungarian_matcher  # pylint: disable=g-import-not-at-top
+
+
+@unittest.skipIf(tf_version.is_tf1(), 'Skipping TF2.X only test.')
+class HungarianBipartiteMatcherTest(test_case.TestCase):
+
+  def test_get_expected_matches_when_all_rows_are_valid(self):
+    similarity_matrix = np.array([[0.50, 0.1, 0.8], [0.15, 0.2, 0.3]],
+                                 dtype=np.float32)
+    valid_rows = np.ones([2], dtype=np.bool)
+    expected_match_results = [-1, 1, 0]
+
+    matcher = hungarian_matcher.HungarianBipartiteMatcher()
+    match_results_out = matcher.match(similarity_matrix, valid_rows=valid_rows)
+
+    self.assertAllEqual(match_results_out._match_results.numpy(),
+                        expected_match_results)
+
+  def test_get_expected_matches_with_all_rows_be_default(self):
+    similarity_matrix = np.array([[0.50, 0.1, 0.8], [0.15, 0.2, 0.3]],
+                                 dtype=np.float32)
+    expected_match_results = [-1, 1, 0]
+
+    matcher = hungarian_matcher.HungarianBipartiteMatcher()
+    match_results_out = matcher.match(similarity_matrix)
+
+    self.assertAllEqual(match_results_out._match_results.numpy(),
+                        expected_match_results)
+
+  def test_get_no_matches_with_zero_valid_rows(self):
+    similarity_matrix = np.array([[0.50, 0.1, 0.8], [0.15, 0.2, 0.3]],
+                                 dtype=np.float32)
+    valid_rows = np.zeros([2], dtype=np.bool)
+    expected_match_results = [-1, -1, -1]
+
+    matcher = hungarian_matcher.HungarianBipartiteMatcher()
+    match_results_out = matcher.match(similarity_matrix, valid_rows=valid_rows)
+
+    self.assertAllEqual(match_results_out._match_results.numpy(),
+                        expected_match_results)
+
+  def test_get_expected_matches_with_only_one_valid_row(self):
+    similarity_matrix = np.array([[0.50, 0.1, 0.8], [0.15, 0.2, 0.3]],
+                                 dtype=np.float32)
+    valid_rows = np.array([True, False], dtype=np.bool)
+    expected_match_results = [-1, -1, 0]
+
+    matcher = hungarian_matcher.HungarianBipartiteMatcher()
+    match_results_out = matcher.match(similarity_matrix, valid_rows=valid_rows)
+
+    self.assertAllEqual(match_results_out._match_results.numpy(),
+                        expected_match_results)
+
+  def test_get_expected_matches_with_only_one_valid_row_at_bottom(self):
+    similarity_matrix = np.array([[0.15, 0.2, 0.3], [0.50, 0.1, 0.8]],
+                                 dtype=np.float32)
+    valid_rows = np.array([False, True], dtype=np.bool)
+    expected_match_results = [-1, -1, 0]
+
+    matcher = hungarian_matcher.HungarianBipartiteMatcher()
+    match_results_out = matcher.match(similarity_matrix, valid_rows=valid_rows)
+
+    self.assertAllEqual(match_results_out._match_results.numpy(),
+                        expected_match_results)
+
+  def test_get_expected_matches_with_two_valid_rows(self):
+    similarity_matrix = np.array([[0.15, 0.2, 0.3], [0.50, 0.1, 0.8],
+                                  [0.84, 0.32, 0.2]],
+                                 dtype=np.float32)
+    valid_rows = np.array([True, False, True], dtype=np.bool)
+    expected_match_results = [1, -1, 0]
+
+    matcher = hungarian_matcher.HungarianBipartiteMatcher()
+    match_results_out = matcher.match(similarity_matrix, valid_rows=valid_rows)
+
+    self.assertAllEqual(match_results_out._match_results.numpy(),
+                        expected_match_results)
+
+
+if __name__ == '__main__':
+  tf.test.main()