Internal change.

PiperOrigin-RevId: 365858684

official/vision/beta/configs/maskrcnn_test.py

+# Copyright 2021 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 maskrcnn."""
+# pylint: disable=unused-import
+from absl.testing import parameterized
+import tensorflow as tf
+
+from official.core import config_definitions as cfg
+from official.core import exp_factory
+from official.vision import beta
+from official.vision.beta.configs import maskrcnn as exp_cfg
+
+
+class MaskRCNNConfigTest(tf.test.TestCase, parameterized.TestCase):
+
+  @parameterized.parameters(
+      ('fasterrcnn_resnetfpn_coco',),
+      ('maskrcnn_resnetfpn_coco',),
+      ('maskrcnn_spinenet_coco',),
+  )
+  def test_maskrcnn_configs(self, config_name):
+    config = exp_factory.get_exp_config(config_name)
+    self.assertIsInstance(config, cfg.ExperimentConfig)
+    self.assertIsInstance(config.task, exp_cfg.MaskRCNNTask)
+    self.assertIsInstance(config.task.model, exp_cfg.MaskRCNN)
+    self.assertIsInstance(config.task.train_data, exp_cfg.DataConfig)
+    config.task.train_data.is_training = None
+    with self.assertRaisesRegex(KeyError, 'Found inconsistncy between key'):
+      config.validate()
+
+
+if __name__ == '__main__':
+  tf.test.main()

official/vision/beta/configs/retinanet_test.py

+# Copyright 2021 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 retinanet."""
+# pylint: disable=unused-import
+from absl.testing import parameterized
+import tensorflow as tf
+
+from official.core import config_definitions as cfg
+from official.core import exp_factory
+from official.vision import beta
+from official.vision.beta.configs import retinanet as exp_cfg
+
+
+class MaskRCNNConfigTest(tf.test.TestCase, parameterized.TestCase):
+
+  @parameterized.parameters(
+      ('retinanet_resnetfpn_coco',),
+      ('retinanet_spinenet_coco',),
+      ('retinanet_spinenet_mobile_coco',),
+  )
+  def test_maskrcnn_configs(self, config_name):
+    config = exp_factory.get_exp_config(config_name)
+    self.assertIsInstance(config, cfg.ExperimentConfig)
+    self.assertIsInstance(config.task, exp_cfg.RetinaNetTask)
+    self.assertIsInstance(config.task.model, exp_cfg.RetinaNet)
+    self.assertIsInstance(config.task.train_data, exp_cfg.DataConfig)
+    config.task.train_data.is_training = None
+    with self.assertRaisesRegex(KeyError, 'Found inconsistncy between key'):
+      config.validate()
+
+
+if __name__ == '__main__':
+  tf.test.main()

official/vision/beta/dataloaders/utils.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 """Data loader utils."""
+from typing import Dict
 
 # Import libraries
 import tensorflow as tf
@@ -20,8 +21,15 @@ import tensorflow as tf
 from official.vision.beta.ops import preprocess_ops
 
 
-def process_source_id(source_id):
-  """Processes source_id to the right format."""
+def process_source_id(source_id: tf.Tensor) -> tf.Tensor:
+  """Processes source_id to the right format.
+
+  Args:
+    source_id: A `tf.Tensor` that contains the source ID. It can be empty.
+
+  Returns:
+    A formatted source ID.
+  """
   if source_id.dtype == tf.string:
     source_id = tf.cast(tf.strings.to_number(source_id), tf.int64)
   with tf.control_dependencies([source_id]):
@@ -32,8 +40,20 @@ def process_source_id(source_id):
   return source_id
 
 
-def pad_groundtruths_to_fixed_size(groundtruths, size):
-  """Pads the first dimension of groundtruths labels to the fixed size."""
+def pad_groundtruths_to_fixed_size(groundtruths: Dict[str, tf.Tensor],
+                                   size: int) -> Dict[str, tf.Tensor]:
+  """Pads the first dimension of groundtruths labels to the fixed size.
+
+  Args:
+    groundtruths: A dictionary of {`str`: `tf.Tensor`} that contains groundtruth
+      annotations of `boxes`, `is_crowds`, `areas` and `classes`.
+    size: An `int` that specifies the expected size of the first dimension of
+      padded tensors.
+
+  Returns:
+    A dictionary of the same keys as input and padded tensors as values.
+
+  """
   groundtruths['boxes'] = preprocess_ops.clip_or_pad_to_fixed_size(
       groundtruths['boxes'], size, -1)
   groundtruths['is_crowds'] = preprocess_ops.clip_or_pad_to_fixed_size(

official/vision/beta/dataloaders/utils_test.py

+# Copyright 2021 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 dataloader utils functions."""
+
+# Import libraries
+
+from absl.testing import parameterized
+import tensorflow as tf
+
+from official.vision.beta.dataloaders import utils
+
+
+class UtilsTest(tf.test.TestCase, parameterized.TestCase):
+
+  def test_process_empty_source_id(self):
+    source_id = tf.constant([], dtype=tf.int64)
+    source_id = tf.strings.as_string(source_id)
+    self.assertEqual(-1, utils.process_source_id(source_id=source_id))
+
+  @parameterized.parameters(
+      ([128, 256], [128, 256]),
+      ([128, 32, 16], [128, 32, 16]),
+  )
+  def test_process_source_id(self, source_id, expected_result):
+    source_id = tf.constant(source_id, dtype=tf.int64)
+    source_id = tf.strings.as_string(source_id)
+    self.assertSequenceAlmostEqual(expected_result,
+                                   utils.process_source_id(source_id=source_id))
+
+  @parameterized.parameters(
+      ([[10, 20, 30, 40]], [[100]], [[0]], 10),
+      ([[0.1, 0.2, 0.5, 0.6]], [[0.5]], [[1]], 2),
+  )
+  def test_pad_groundtruths_to_fixed_size(self, boxes, area, classes, size):
+    groundtruths = {}
+    groundtruths['boxes'] = tf.constant(boxes)
+    groundtruths['is_crowds'] = tf.constant([[0]])
+    groundtruths['areas'] = tf.constant(area)
+    groundtruths['classes'] = tf.constant(classes)
+
+    actual_result = utils.pad_groundtruths_to_fixed_size(
+        groundtruths=groundtruths, size=size)
+
+    # Check that the first dimension is padded to the expected size.
+    for key in actual_result:
+      pad_shape = actual_result[key].shape[0]
+      self.assertEqual(size, pad_shape)
+
+
+if __name__ == '__main__':
+  tf.test.main()