Added a test utility method to compare protos directly

mediapipe/tasks/python/test/test_utils.py

@@ -14,7 +14,12 @@
 """Test util for MediaPipe Tasks."""
 
 import os
+import difflib
+import six
 
+from google.protobuf import descriptor
+from google.protobuf import descriptor_pool
+from google.protobuf import text_format
 from absl import flags
 
 from mediapipe.python._framework_bindings import image as image_module
@@ -53,3 +58,110 @@ def create_calibration_file(file_dir: str,
   with open(calibration_file, mode="w") as file:
     file.write(content)
   return calibration_file
+
+
+def assertProtoEqual(self, a, b, check_initialized=True,
+                     normalize_numbers=False, msg=None):
+  """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.
+  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.Merge(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.
+  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, 6) for x in values]
+    elif desc.type == descriptor.FieldDescriptor.TYPE_DOUBLE:
+      normalized_values = [round(float(x), 7) 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

@@ -51,8 +51,8 @@ _MAX_RESULTS = 3
 
 
 # TODO: Port assertProtoEquals
-def _assert_proto_equals(expected, actual):  # pylint: disable=unused-argument
+def _assert_proto_equals(self, expected, actual):  # pylint: disable=unused-argument
-  pass
+  test_utils.assertProtoEqual(self, expected, actual)
 
 
 def _generate_empty_results(timestamp_ms: int) -> _ClassificationResult:
@@ -186,7 +186,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(),
+    _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 +214,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(),
+      _assert_proto_equals(self, image_result.to_pb2(),
                            expected_classification_result.to_pb2())
 
   def test_classify_succeeds_with_region_of_interest(self):
@@ -232,7 +232,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(),
+      _assert_proto_equals(self, image_result.to_pb2(),
                            _generate_soccer_ball_results(0).to_pb2())
 
   def test_score_threshold_option(self):
@@ -401,7 +401,7 @@ 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(),
+        _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 +420,8 @@ 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,
+        _assert_proto_equals(self, classification_result.to_pb2(),
-                         _generate_soccer_ball_results(timestamp))
+                             _generate_soccer_ball_results(timestamp).to_pb2())
 
   def test_calling_classify_in_live_stream_mode(self):
     options = _ImageClassifierOptions(
@@ -463,7 +463,7 @@ class ImageClassifierTest(parameterized.TestCase):
 
     def check_result(result: _ClassificationResult, output_image: _Image,
                      timestamp_ms: int):
-      _assert_proto_equals(result.to_pb2(),
+      _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 +493,7 @@ class ImageClassifierTest(parameterized.TestCase):
 
     def check_result(result: _ClassificationResult, output_image: _Image,
                      timestamp_ms: int):
-      _assert_proto_equals(result.to_pb2(),
+      _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)