Merge pull request #3800 from kinaryml:python-test-proto-equals

PiperOrigin-RevId: 485340924

mediapipe/tasks/python/test/BUILD

@@ -27,5 +27,8 @@ py_library(
         "//mediapipe/model_maker/python/vision/gesture_recognizer:__pkg__",
         "//mediapipe/tasks:internal",
     ],
-    deps = ["//mediapipe/python:_framework_bindings"],
+    deps = [
+        "//mediapipe/python:_framework_bindings",
+        "@com_google_protobuf//:protobuf_python",
+    ],
 )

mediapipe/tasks/python/test/test_utils.py

@@ -13,9 +13,15 @@
 # limitations under the License.
 """Test util for MediaPipe Tasks."""
 
+import difflib
 import os
 
 from absl import flags
+import six
+
+from google.protobuf import descriptor
+from google.protobuf import descriptor_pool
+from google.protobuf import text_format
 
 from mediapipe.python._framework_bindings import image as image_module
 from mediapipe.python._framework_bindings import image_frame as image_frame_module
@@ -53,3 +59,126 @@ def create_calibration_file(file_dir: str,
   with open(calibration_file, mode="w") as file:
     file.write(content)
   return calibration_file
+
+
+def assert_proto_equals(self,
+                        a,
+                        b,
+                        check_initialized=True,
+                        normalize_numbers=True,
+                        msg=None):
+  """assert_proto_equals() is useful for unit tests.
+
+  It produces much more helpful output than assertEqual() for proto2 messages.
+  Fails with a useful error if a and b aren't equal. Comparison of repeated
+  fields matches the semantics of unittest.TestCase.assertEqual(), ie order and
+  extra duplicates fields matter.
+
+  This is a fork of https://github.com/tensorflow/tensorflow/blob/
+  master/tensorflow/python/util/protobuf/compare.py#L73. We use slightly
+  different rounding cutoffs to support Mac usage.
+
+  Args:
+    self: absltest.testing.parameterized.TestCase
+    a: proto2 PB instance, or text string representing one.
+    b: proto2 PB instance -- message.Message or subclass thereof.
+    check_initialized: boolean, whether to fail if either a or b isn't
+      initialized.
+    normalize_numbers: boolean, whether to normalize types and precision of
+      numbers before comparison.
+    msg: if specified, is used as the error message on failure.
+  """
+  pool = descriptor_pool.Default()
+  if isinstance(a, six.string_types):
+    a = text_format.Parse(a, b.__class__(), descriptor_pool=pool)
+
+  for pb in a, b:
+    if check_initialized:
+      errors = pb.FindInitializationErrors()
+      if errors:
+        self.fail("Initialization errors: %s\n%s" % (errors, pb))
+    if normalize_numbers:
+      _normalize_number_fields(pb)
+
+  a_str = text_format.MessageToString(a, descriptor_pool=pool)
+  b_str = text_format.MessageToString(b, descriptor_pool=pool)
+
+  # Some Python versions would perform regular diff instead of multi-line
+  # diff if string is longer than 2**16. We substitute this behavior
+  # with a call to unified_diff instead to have easier-to-read diffs.
+  # For context, see: https://bugs.python.org/issue11763.
+  if len(a_str) < 2**16 and len(b_str) < 2**16:
+    self.assertMultiLineEqual(a_str, b_str, msg=msg)
+  else:
+    diff = "".join(
+        difflib.unified_diff(a_str.splitlines(True), b_str.splitlines(True)))
+    if diff:
+      self.fail("%s :\n%s" % (msg, diff))
+
+
+def _normalize_number_fields(pb):
+  """Normalizes types and precisions of number fields in a protocol buffer.
+
+  Due to subtleties in the python protocol buffer implementation, it is possible
+  for values to have different types and precision depending on whether they
+  were set and retrieved directly or deserialized from a protobuf. This function
+  normalizes integer values to ints and longs based on width, 32-bit floats to
+  five digits of precision to account for python always storing them as 64-bit,
+  and ensures doubles are floating point for when they're set to integers.
+  Modifies pb in place. Recurses into nested objects. https://github.com/tensorf
+  low/tensorflow/blob/master/tensorflow/python/util/protobuf/compare.py#L118
+
+  Args:
+    pb: proto2 message.
+
+  Returns:
+    the given pb, modified in place.
+  """
+  for desc, values in pb.ListFields():
+    is_repeated = True
+    if desc.label != descriptor.FieldDescriptor.LABEL_REPEATED:
+      is_repeated = False
+      values = [values]
+
+    normalized_values = None
+
+    # We force 32-bit values to int and 64-bit values to long to make
+    # alternate implementations where the distinction is more significant
+    # (e.g. the C++ implementation) simpler.
+    if desc.type in (descriptor.FieldDescriptor.TYPE_INT64,
+                     descriptor.FieldDescriptor.TYPE_UINT64,
+                     descriptor.FieldDescriptor.TYPE_SINT64):
+      normalized_values = [int(x) for x in values]
+    elif desc.type in (descriptor.FieldDescriptor.TYPE_INT32,
+                       descriptor.FieldDescriptor.TYPE_UINT32,
+                       descriptor.FieldDescriptor.TYPE_SINT32,
+                       descriptor.FieldDescriptor.TYPE_ENUM):
+      normalized_values = [int(x) for x in values]
+    elif desc.type == descriptor.FieldDescriptor.TYPE_FLOAT:
+      normalized_values = [round(x, 5) for x in values]
+    elif desc.type == descriptor.FieldDescriptor.TYPE_DOUBLE:
+      normalized_values = [round(float(x), 6) for x in values]
+
+    if normalized_values is not None:
+      if is_repeated:
+        pb.ClearField(desc.name)
+        getattr(pb, desc.name).extend(normalized_values)
+      else:
+        setattr(pb, desc.name, normalized_values[0])
+
+    if (desc.type == descriptor.FieldDescriptor.TYPE_MESSAGE or
+        desc.type == descriptor.FieldDescriptor.TYPE_GROUP):
+      if (desc.type == descriptor.FieldDescriptor.TYPE_MESSAGE and
+          desc.message_type.has_options and
+          desc.message_type.GetOptions().map_entry):
+        # This is a map, only recurse if the values have a message type.
+        if (desc.message_type.fields_by_number[2].type ==
+            descriptor.FieldDescriptor.TYPE_MESSAGE):
+          for v in six.itervalues(values):
+            _normalize_number_fields(v)
+      else:
+        for v in values:
+          # recursive step
+          _normalize_number_fields(v)
+
+  return pb

mediapipe/tasks/python/test/vision/image_classifier_test.py

@@ -50,11 +50,6 @@ _SCORE_THRESHOLD = 0.5
 _MAX_RESULTS = 3
 
 
-# TODO: Port assertProtoEquals
-def _assert_proto_equals(expected, actual):  # pylint: disable=unused-argument
-  pass
-
-
 def _generate_empty_results(timestamp_ms: int) -> _ClassificationResult:
   return _ClassificationResult(classifications=[
       _Classifications(
@@ -74,22 +69,22 @@ def _generate_burger_results(timestamp_ms: int) -> _ClassificationResult:
                   categories=[
                       _Category(
                           index=934,
-                          score=0.7939587831497192,
+                          score=0.793959,
                           display_name='',
                           category_name='cheeseburger'),
                       _Category(
                           index=932,
-                          score=0.02739289402961731,
+                          score=0.0273929,
                           display_name='',
                           category_name='bagel'),
                       _Category(
                           index=925,
-                          score=0.01934075355529785,
+                          score=0.0193408,
                           display_name='',
                           category_name='guacamole'),
                       _Category(
                           index=963,
-                          score=0.006327860057353973,
+                          score=0.00632786,
                           display_name='',
                           category_name='meat loaf')
                   ],
@@ -108,7 +103,7 @@ def _generate_soccer_ball_results(timestamp_ms: int) -> _ClassificationResult:
                   categories=[
                       _Category(
                           index=806,
-                          score=0.9965274930000305,
+                          score=0.996527,
                           display_name='',
                           category_name='soccer ball')
                   ],
@@ -186,7 +181,7 @@ class ImageClassifierTest(parameterized.TestCase):
     # Performs image classification on the input.
     image_result = classifier.classify(self.test_image)
     # Comparing results.
-    _assert_proto_equals(image_result.to_pb2(),
+    test_utils.assert_proto_equals(self, image_result.to_pb2(),
                                    expected_classification_result.to_pb2())
     # Closes the classifier explicitly when the classifier is not used in
     # a context.
@@ -214,7 +209,7 @@ class ImageClassifierTest(parameterized.TestCase):
       # Performs image classification on the input.
       image_result = classifier.classify(self.test_image)
       # Comparing results.
-      _assert_proto_equals(image_result.to_pb2(),
+      test_utils.assert_proto_equals(self, image_result.to_pb2(),
                                      expected_classification_result.to_pb2())
 
   def test_classify_succeeds_with_region_of_interest(self):
@@ -232,7 +227,7 @@ class ImageClassifierTest(parameterized.TestCase):
       # Performs image classification on the input.
       image_result = classifier.classify(test_image, roi)
       # Comparing results.
-      _assert_proto_equals(image_result.to_pb2(),
+      test_utils.assert_proto_equals(self, image_result.to_pb2(),
                                      _generate_soccer_ball_results(0).to_pb2())
 
   def test_score_threshold_option(self):
@@ -401,7 +396,8 @@ class ImageClassifierTest(parameterized.TestCase):
       for timestamp in range(0, 300, 30):
         classification_result = classifier.classify_for_video(
             self.test_image, timestamp)
-        _assert_proto_equals(classification_result.to_pb2(),
+        test_utils.assert_proto_equals(
+            self, classification_result.to_pb2(),
             _generate_burger_results(timestamp).to_pb2())
 
   def test_classify_for_video_succeeds_with_region_of_interest(self):
@@ -420,8 +416,9 @@ class ImageClassifierTest(parameterized.TestCase):
       for timestamp in range(0, 300, 30):
         classification_result = classifier.classify_for_video(
             test_image, timestamp, roi)
-        self.assertEqual(classification_result,
+        test_utils.assert_proto_equals(
-                         _generate_soccer_ball_results(timestamp))
+            self, classification_result.to_pb2(),
+            _generate_soccer_ball_results(timestamp).to_pb2())
 
   def test_calling_classify_in_live_stream_mode(self):
     options = _ImageClassifierOptions(
@@ -463,7 +460,7 @@ class ImageClassifierTest(parameterized.TestCase):
 
     def check_result(result: _ClassificationResult, output_image: _Image,
                      timestamp_ms: int):
-      _assert_proto_equals(result.to_pb2(),
+      test_utils.assert_proto_equals(self, result.to_pb2(),
                                      expected_result_fn(timestamp_ms).to_pb2())
       self.assertTrue(
           np.array_equal(output_image.numpy_view(),
@@ -493,7 +490,8 @@ class ImageClassifierTest(parameterized.TestCase):
 
     def check_result(result: _ClassificationResult, output_image: _Image,
                      timestamp_ms: int):
-      _assert_proto_equals(result.to_pb2(),
+      test_utils.assert_proto_equals(
+          self, result.to_pb2(),
           _generate_soccer_ball_results(timestamp_ms).to_pb2())
       self.assertEqual(output_image.width, test_image.width)
       self.assertEqual(output_image.height, test_image.height)