Internal change

PiperOrigin-RevId: 275306010

official/vision/detection/configs/retinanet_config_test.py

+# Copyright 2019 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_config.py."""
+
+import tensorflow.compat.v2 as tf
+
+from official.vision.detection.configs import retinanet_config
+from official.modeling.hyperparams import params_dict
+
+
+class RetinanetConfigTest(tf.test.TestCase):
+
+  def test_restrictions_do_not_have_typos(self):
+    cfg = params_dict.ParamsDict(
+        retinanet_config.RETINANET_CFG, retinanet_config.RETINANET_RESTRICTIONS)
+    cfg.validate()
+
+
+if __name__ == '__main__':
+  tf.test.main()

official/vision/detection/dataloader/anchor_test.py

+# Copyright 2019 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 anchor.py."""
+
+from absl.testing import parameterized
+import numpy as np
+import tensorflow.compat.v2 as tf
+from official.vision.detection.dataloader import anchor
+
+
+class AnchorTest(parameterized.TestCase, tf.test.TestCase):
+
+  # The set of parameters are tailored for the MLPerf configuration, where
+  # the number of anchors is 495132, rpn_batch_size_per_im=256, and
+  # rpn_fg_fraction=0.5.
+  @parameterized.parameters(
+      (512, 25, 25, 25, 25, (512, 512)),
+      (512, 25, 25, 25, 25, (512, 640)),
+      (512, 25, 25, 25, 25, (640, 512)),
+      (495132, 100, 100, 100, 100, (512, 512)),
+      (495132, 200, 100, 128, 100, (512, 512)),
+      (495132, 100, 120, 100, 120, (512, 512)),
+      (495132, 100, 200, 100, 156, (512, 512)),
+      (495132, 200, 200, 128, 128, (512, 512)),
+  )
+  def testAnchorRpnSample(self, num_anchors, num_positives,
+                          num_negatives, expected_positives,
+                          expected_negatives, image_size):
+    num_anchors = num_anchors
+    num_positives = num_positives
+    num_negatives = num_negatives
+    match_results_np = np.empty([num_anchors])
+    match_results_np.fill(-2)
+    match_results_np[:num_positives] = 0
+    match_results_np[num_positives:num_positives + num_negatives] = -1
+    match_results = tf.convert_to_tensor(value=match_results_np, dtype=tf.int32)
+    input_anchor = anchor.Anchor(
+        min_level=3, max_level=7, num_scales=1,
+        aspect_ratios=[1.0, 2.0, 0.5], anchor_size=4,
+        image_size=image_size)
+    anchor_labeler = anchor.RpnAnchorLabeler(
+        input_anchor, match_threshold=0.7,
+        unmatched_threshold=0.3, rpn_batch_size_per_im=256,
+        rpn_fg_fraction=0.5)
+    rpn_sample_op = anchor_labeler._get_rpn_samples(match_results)
+    labels = [v.numpy() for v in rpn_sample_op]
+    self.assertLen(labels[0], num_anchors)
+    positives = np.sum(np.array(labels[0]) == 1)
+    negatives = np.sum(np.array(labels[0]) == 0)
+    self.assertEqual(positives, expected_positives)
+    self.assertEqual(negatives, expected_negatives)
+
+  @parameterized.parameters(
+      # Single scale anchor.
+      (5, 5, 1, [1.0], 2.0,
+       [[-16, -16, 48, 48], [-16, 16, 48, 80],
+        [16, -16, 80, 48], [16, 16, 80, 80]]),
+      # Multi scale anchor.
+      (5, 6, 1, [1.0], 2.0,
+       [[-16, -16, 48, 48], [-16, 16, 48, 80],
+        [16, -16, 80, 48], [16, 16, 80, 80], [-32, -32, 96, 96]]),
+      # # Multi aspect ratio anchor.
+      (6, 6, 1, [1.0, 4.0, 0.25], 2.0,
+       [[-32, -32, 96, 96], [-0, -96, 64, 160], [-96, -0, 160, 64]]),
+
+  )
+  def testAnchorGeneration(self, min_level, max_level, num_scales,
+                           aspect_ratios, anchor_size, expected_boxes):
+    image_size = [64, 64]
+    anchors = anchor.Anchor(min_level, max_level, num_scales, aspect_ratios,
+                            anchor_size, image_size)
+    boxes = anchors.boxes.numpy()
+    self.assertEqual(expected_boxes, boxes.tolist())
+
+  @parameterized.parameters(
+      # Single scale anchor.
+      (5, 5, 1, [1.0], 2.0,
+       [[-16, -16, 48, 48], [-16, 16, 48, 80],
+        [16, -16, 80, 48], [16, 16, 80, 80]]),
+      # Multi scale anchor.
+      (5, 6, 1, [1.0], 2.0,
+       [[-16, -16, 48, 48], [-16, 16, 48, 80],
+        [16, -16, 80, 48], [16, 16, 80, 80], [-32, -32, 96, 96]]),
+      # # Multi aspect ratio anchor.
+      (6, 6, 1, [1.0, 4.0, 0.25], 2.0,
+       [[-32, -32, 96, 96], [-0, -96, 64, 160], [-96, -0, 160, 64]]),
+
+  )
+  def testAnchorGenerationWithImageSizeAsTensor(self,
+                                                min_level,
+                                                max_level,
+                                                num_scales,
+                                                aspect_ratios,
+                                                anchor_size,
+                                                expected_boxes):
+    image_size = tf.constant([64, 64], tf.int32)
+    anchors = anchor.Anchor(min_level, max_level, num_scales, aspect_ratios,
+                            anchor_size, image_size)
+    boxes = anchors.boxes.numpy()
+    self.assertEqual(expected_boxes, boxes.tolist())
+
+  @parameterized.parameters(
+      (3, 6, 2, [1.0], 2.0),
+  )
+  def testLabelAnchors(self, min_level, max_level, num_scales,
+                       aspect_ratios, anchor_size):
+    input_size = [512, 512]
+    ground_truth_class_id = 2
+
+    # The matched anchors are the anchors used as ground truth and the anchors
+    # at the next octave scale on the same location.
+    expected_anchor_locations = [[0, 0, 0], [0, 0, 1]]
+    anchors = anchor.Anchor(min_level, max_level, num_scales, aspect_ratios,
+                            anchor_size, input_size)
+    anchor_labeler = anchor.AnchorLabeler(anchors)
+
+    # Uses the first anchors as ground truth. The ground truth should map to
+    # two anchors with two intermediate scales at the same location.
+    gt_boxes = anchors.boxes[0:1, :]
+    gt_classes = tf.constant([[ground_truth_class_id]], dtype=tf.float32)
+    cls_targets, box_targets, num_positives = anchor_labeler.label_anchors(
+        gt_boxes, gt_classes)
+
+    for k, v in cls_targets.items():
+      cls_targets[k] = v.numpy()
+    for k, v in box_targets.items():
+      box_targets[k] = v.numpy()
+    num_positives = num_positives.numpy()
+
+    anchor_locations = np.vstack(
+        np.where(cls_targets[min_level] > -1)).transpose()
+    self.assertAllClose(expected_anchor_locations, anchor_locations)
+    # Two anchor boxes on min_level got matched to the gt_boxes.
+    self.assertAllClose(num_positives, len(expected_anchor_locations))
+
+if __name__ == '__main__':
+  assert tf.version.VERSION.startswith('2.')
+  tf.test.main()

official/vision/detection/dataloader/tf_example_decoder_test.py

+# Copyright 2019 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 tf_example_decoder.py."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from absl import logging
+from absl.testing import parameterized
+import io
+import numpy as np
+from PIL import Image
+import tensorflow.compat.v2 as tf
+
+from official.vision.detection.dataloader import tf_example_decoder
+
+
+def _encode_image(image_array, fmt):
+  image = Image.fromarray(image_array)
+  with io.BytesIO() as output:
+    image.save(output, format=fmt)
+    return output.getvalue()
+
+
+class TfExampleDecoderTest(tf.test.TestCase, parameterized.TestCase):
+
+  @parameterized.parameters(
+      (100, 100, 0), (100, 100, 1), (100, 100, 2),
+  )
+  def test_result_shape(self, image_height, image_width, num_instances):
+    decoder = tf_example_decoder.TfExampleDecoder(include_mask=True)
+
+    image = _encode_image(
+        np.uint8(np.random.rand(image_height, image_width, 3) * 255),
+        fmt='JPEG')
+    if num_instances == 0:
+      xmins = []
+      xmaxs = []
+      ymins = []
+      ymaxs = []
+      labels = []
+      areas = []
+      is_crowds = []
+      masks = []
+    else:
+      xmins = list(np.random.rand(num_instances))
+      xmaxs = list(np.random.rand(num_instances))
+      ymins = list(np.random.rand(num_instances))
+      ymaxs = list(np.random.rand(num_instances))
+      labels = list(np.random.randint(100, size=num_instances))
+      areas = [(xmax - xmin) * (ymax - ymin) for xmin, xmax, ymin, ymax in
+               zip(xmins, xmaxs, ymins, ymaxs)]
+      is_crowds = [0] * num_instances
+      masks = []
+      for _ in range(num_instances):
+        mask = _encode_image(
+            np.uint8(np.random.rand(image_height, image_width) * 255),
+            fmt='PNG')
+        masks.append(mask)
+    serialized_example = tf.train.Example(
+        features=tf.train.Features(
+            feature={
+                'image/encoded': (
+                    tf.train.Feature(
+                        bytes_list=tf.train.BytesList(value=[image]))),
+                'image/source_id': (
+                    tf.train.Feature(
+                        bytes_list=tf.train.BytesList(value=['123']))),
+                'image/height': (
+                    tf.train.Feature(
+                        int64_list=tf.train.Int64List(value=[image_height]))),
+                'image/width': (
+                    tf.train.Feature(
+                        int64_list=tf.train.Int64List(value=[image_width]))),
+                'image/object/bbox/xmin': (
+                    tf.train.Feature(
+                        float_list=tf.train.FloatList(value=xmins))),
+                'image/object/bbox/xmax': (
+                    tf.train.Feature(
+                        float_list=tf.train.FloatList(value=xmaxs))),
+                'image/object/bbox/ymin': (
+                    tf.train.Feature(
+                        float_list=tf.train.FloatList(value=ymins))),
+                'image/object/bbox/ymax': (
+                    tf.train.Feature(
+                        float_list=tf.train.FloatList(value=ymaxs))),
+                'image/object/class/label': (
+                    tf.train.Feature(
+                        int64_list=tf.train.Int64List(value=labels))),
+                'image/object/is_crowd': (
+                    tf.train.Feature(
+                        int64_list=tf.train.Int64List(value=is_crowds))),
+                'image/object/area': (
+                    tf.train.Feature(
+                        float_list=tf.train.FloatList(value=areas))),
+                'image/object/mask': (
+                    tf.train.Feature(
+                        bytes_list=tf.train.BytesList(value=masks))),
+            })).SerializeToString()
+    decoded_tensors = decoder.decode(
+        tf.convert_to_tensor(value=serialized_example))
+
+    results = tf.nest.map_structure(lambda x: x.numpy(), decoded_tensors)
+
+    self.assertAllEqual(
+        (image_height, image_width, 3), results['image'].shape)
+    self.assertEqual('123', results['source_id'])
+    self.assertEqual(image_height, results['height'])
+    self.assertEqual(image_width, results['width'])
+    self.assertAllEqual(
+        (num_instances,), results['groundtruth_classes'].shape)
+    self.assertAllEqual(
+        (num_instances,), results['groundtruth_is_crowd'].shape)
+    self.assertAllEqual(
+        (num_instances,), results['groundtruth_area'].shape)
+    self.assertAllEqual(
+        (num_instances, 4), results['groundtruth_boxes'].shape)
+    self.assertAllEqual(
+        (num_instances, image_height, image_width),
+        results['groundtruth_instance_masks'].shape)
+    self.assertAllEqual(
+        (num_instances,), results['groundtruth_instance_masks_png'].shape)
+
+  def test_result_content(self):
+    decoder = tf_example_decoder.TfExampleDecoder(include_mask=True)
+
+    image_content = [[[0, 0, 0], [0, 0, 0], [0, 0, 0], [0, 0, 0]],
+                     [[0, 0, 0], [255, 255, 255], [255, 255, 255], [0, 0, 0]],
+                     [[0, 0, 0], [255, 255, 255], [255, 255, 255], [0, 0, 0]],
+                     [[0, 0, 0], [0, 0, 0], [0, 0, 0], [0, 0, 0]]]
+    image = _encode_image(np.uint8(image_content), fmt='PNG')
+    image_height = 4
+    image_width = 4
+    num_instances = 2
+    xmins = [0, 0.25]
+    xmaxs = [0.5, 1.0]
+    ymins = [0, 0]
+    ymaxs = [0.5, 1.0]
+    labels = [3, 1]
+    areas = [0.25, 0.75]
+    is_crowds = [1, 0]
+    mask_content = [[[255, 255, 0, 0],
+                     [255, 255, 0, 0],
+                     [0, 0, 0, 0],
+                     [0, 0, 0, 0]],
+                    [[0, 255, 255, 255],
+                     [0, 255, 255, 255],
+                     [0, 255, 255, 255],
+                     [0, 255, 255, 255]]]
+    masks = [_encode_image(np.uint8(m), fmt='PNG') for m in list(mask_content)]
+    serialized_example = tf.train.Example(
+        features=tf.train.Features(
+            feature={
+                'image/encoded': (
+                    tf.train.Feature(
+                        bytes_list=tf.train.BytesList(value=[image]))),
+                'image/source_id': (
+                    tf.train.Feature(
+                        bytes_list=tf.train.BytesList(value=['123']))),
+                'image/height': (
+                    tf.train.Feature(
+                        int64_list=tf.train.Int64List(value=[image_height]))),
+                'image/width': (
+                    tf.train.Feature(
+                        int64_list=tf.train.Int64List(value=[image_width]))),
+                'image/object/bbox/xmin': (
+                    tf.train.Feature(
+                        float_list=tf.train.FloatList(value=xmins))),
+                'image/object/bbox/xmax': (
+                    tf.train.Feature(
+                        float_list=tf.train.FloatList(value=xmaxs))),
+                'image/object/bbox/ymin': (
+                    tf.train.Feature(
+                        float_list=tf.train.FloatList(value=ymins))),
+                'image/object/bbox/ymax': (
+                    tf.train.Feature(
+                        float_list=tf.train.FloatList(value=ymaxs))),
+                'image/object/class/label': (
+                    tf.train.Feature(
+                        int64_list=tf.train.Int64List(value=labels))),
+                'image/object/is_crowd': (
+                    tf.train.Feature(
+                        int64_list=tf.train.Int64List(value=is_crowds))),
+                'image/object/area': (
+                    tf.train.Feature(
+                        float_list=tf.train.FloatList(value=areas))),
+                'image/object/mask': (
+                    tf.train.Feature(
+                        bytes_list=tf.train.BytesList(value=masks))),
+            })).SerializeToString()
+    decoded_tensors = decoder.decode(
+        tf.convert_to_tensor(value=serialized_example))
+
+    results = tf.nest.map_structure(lambda x: x.numpy(), decoded_tensors)
+
+    self.assertAllEqual(
+        (image_height, image_width, 3), results['image'].shape)
+    self.assertAllEqual(image_content, results['image'])
+    self.assertEqual('123', results['source_id'])
+    self.assertEqual(image_height, results['height'])
+    self.assertEqual(image_width, results['width'])
+    self.assertAllEqual(
+        (num_instances,), results['groundtruth_classes'].shape)
+    self.assertAllEqual(
+        (num_instances,), results['groundtruth_is_crowd'].shape)
+    self.assertAllEqual(
+        (num_instances,), results['groundtruth_area'].shape)
+    self.assertAllEqual(
+        (num_instances, 4), results['groundtruth_boxes'].shape)
+    self.assertAllEqual(
+        (num_instances, image_height, image_width),
+        results['groundtruth_instance_masks'].shape)
+    self.assertAllEqual(
+        (num_instances,), results['groundtruth_instance_masks_png'].shape)
+    self.assertAllEqual(
+        [3, 1], results['groundtruth_classes'])
+    self.assertAllEqual(
+        [True, False], results['groundtruth_is_crowd'])
+    self.assertNDArrayNear(
+        [0.25, 0.75], results['groundtruth_area'], 1e-4)
+    self.assertNDArrayNear(
+        [[0, 0, 0.5, 0.5], [0, 0.25, 1.0, 1.0]],
+        results['groundtruth_boxes'], 1e-4)
+    self.assertNDArrayNear(
+        mask_content, results['groundtruth_instance_masks'], 1e-4)
+    self.assertAllEqual(
+        masks, results['groundtruth_instance_masks_png'])
+
+  def test_handling_missing_fields(self):
+    decoder = tf_example_decoder.TfExampleDecoder(include_mask=True)
+
+    image_content = [[[0, 0, 0], [0, 0, 0], [0, 0, 0], [0, 0, 0]],
+                     [[0, 0, 0], [255, 255, 255], [255, 255, 255], [0, 0, 0]],
+                     [[0, 0, 0], [255, 255, 255], [255, 255, 255], [0, 0, 0]],
+                     [[0, 0, 0], [0, 0, 0], [0, 0, 0], [0, 0, 0]]]
+    image = _encode_image(np.uint8(image_content), fmt='PNG')
+    image_height = 4
+    image_width = 4
+    num_instances = 2
+    xmins = [0, 0.25]
+    xmaxs = [0.5, 1.0]
+    ymins = [0, 0]
+    ymaxs = [0.5, 1.0]
+    labels = [3, 1]
+    mask_content = [[[255, 255, 0, 0],
+                     [255, 255, 0, 0],
+                     [0, 0, 0, 0],
+                     [0, 0, 0, 0]],
+                    [[0, 255, 255, 255],
+                     [0, 255, 255, 255],
+                     [0, 255, 255, 255],
+                     [0, 255, 255, 255]]]
+    masks = [_encode_image(np.uint8(m), fmt='PNG') for m in list(mask_content)]
+    serialized_example = tf.train.Example(
+        features=tf.train.Features(
+            feature={
+                'image/encoded': (
+                    tf.train.Feature(
+                        bytes_list=tf.train.BytesList(value=[image]))),
+                'image/source_id': (
+                    tf.train.Feature(
+                        bytes_list=tf.train.BytesList(value=['123']))),
+                'image/height': (
+                    tf.train.Feature(
+                        int64_list=tf.train.Int64List(value=[image_height]))),
+                'image/width': (
+                    tf.train.Feature(
+                        int64_list=tf.train.Int64List(value=[image_width]))),
+                'image/object/bbox/xmin': (
+                    tf.train.Feature(
+                        float_list=tf.train.FloatList(value=xmins))),
+                'image/object/bbox/xmax': (
+                    tf.train.Feature(
+                        float_list=tf.train.FloatList(value=xmaxs))),
+                'image/object/bbox/ymin': (
+                    tf.train.Feature(
+                        float_list=tf.train.FloatList(value=ymins))),
+                'image/object/bbox/ymax': (
+                    tf.train.Feature(
+                        float_list=tf.train.FloatList(value=ymaxs))),
+                'image/object/class/label': (
+                    tf.train.Feature(
+                        int64_list=tf.train.Int64List(value=labels))),
+                'image/object/mask': (
+                    tf.train.Feature(
+                        bytes_list=tf.train.BytesList(value=masks))),
+            })).SerializeToString()
+    decoded_tensors = decoder.decode(
+        tf.convert_to_tensor(serialized_example))
+    results = tf.nest.map_structure(lambda x: x.numpy(), decoded_tensors)
+
+    self.assertAllEqual(
+        (image_height, image_width, 3), results['image'].shape)
+    self.assertAllEqual(image_content, results['image'])
+    self.assertEqual('123', results['source_id'])
+    self.assertEqual(image_height, results['height'])
+    self.assertEqual(image_width, results['width'])
+    self.assertAllEqual(
+        (num_instances,), results['groundtruth_classes'].shape)
+    self.assertAllEqual(
+        (num_instances,), results['groundtruth_is_crowd'].shape)
+    self.assertAllEqual(
+        (num_instances,), results['groundtruth_area'].shape)
+    self.assertAllEqual(
+        (num_instances, 4), results['groundtruth_boxes'].shape)
+    self.assertAllEqual(
+        (num_instances, image_height, image_width),
+        results['groundtruth_instance_masks'].shape)
+    self.assertAllEqual(
+        (num_instances,), results['groundtruth_instance_masks_png'].shape)
+    self.assertAllEqual(
+        [3, 1], results['groundtruth_classes'])
+    self.assertAllEqual(
+        [False, False], results['groundtruth_is_crowd'])
+    self.assertNDArrayNear(
+        [0.25, 0.75], results['groundtruth_area'], 1e-4)
+    self.assertNDArrayNear(
+        [[0, 0, 0.5, 0.5], [0, 0.25, 1.0, 1.0]],
+        results['groundtruth_boxes'], 1e-4)
+    self.assertNDArrayNear(
+        mask_content, results['groundtruth_instance_masks'], 1e-4)
+    self.assertAllEqual(
+        masks, results['groundtruth_instance_masks_png'])
+
+if __name__ == '__main__':
+  assert tf.version.VERSION.startswith('2.')
+  tf.test.main()

official/vision/detection/modeling/architecture/fpn_test.py

+# Copyright 2019 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 fpn.py."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from absl.testing import parameterized
+import tensorflow.compat.v2 as tf
+from official.vision.detection.modeling.architecture import fpn
+
+
+class FpnTest(parameterized.TestCase, tf.test.TestCase):
+
+  @parameterized.parameters(
+      (3, 7),
+      (3, 4),
+  )
+  def testFPNOutputShape(self, min_level, max_level):
+    backbone_min_level = 2
+    backbone_max_level = 5
+    fpn_feat_dims = 256
+    image_size = 256
+
+    inputs = {}
+    for level in range(backbone_min_level, backbone_max_level + 1):
+      inputs[level] = tf.zeros(
+          [1, image_size // 2**level, image_size // 2**level, fpn_feat_dims])
+
+    with tf.name_scope('min_level_%d_max_level_%d' % (min_level, max_level)):
+      fpn_fn = fpn.Fpn(min_level=min_level,
+                       max_level=max_level,
+                       fpn_feat_dims=fpn_feat_dims)
+      features = fpn_fn(inputs)
+
+      for level in range(min_level, max_level):
+        self.assertEqual(features[level].get_shape().as_list(),
+                         [1, image_size // 2**level,
+                          image_size // 2**level, fpn_feat_dims])
+      self.assertEqual(sorted(features.keys()), range(min_level, max_level + 1))
+
+
+if __name__ == '__main__':
+  assert tf.version.VERSION.startswith('2.')
+  tf.test.main()

official/vision/detection/modeling/architecture/resnet_test.py

+# Copyright 2019 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 resnet.py."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import tensorflow.compat.v2 as tf
+
+from absl.testing import parameterized
+from official.vision.detection.modeling.architecture import resnet
+
+
+class ResnetTest(parameterized.TestCase, tf.test.TestCase):
+
+  @parameterized.parameters(
+      (50, None, None, True),
+      (101, None, None, True),
+      (152, None, None, True),
+      (200, None, None, True),
+      (50, 0.9, 7, False),
+      (50, 0.9, 7, True),
+  )
+  def testResNetOutputShape(self, resnet_depth, dropblock_keep_prob,
+                            dropblock_size, is_training):
+
+    inputs = tf.zeros([1, 256, 256, 3])
+    resnet_fn = resnet.Resnet(
+        resnet_depth,
+        dropblock_keep_prob=dropblock_keep_prob,
+        dropblock_size=dropblock_size)
+    features = resnet_fn(inputs, is_training=is_training)
+    self.assertEqual(features[2].get_shape().as_list(), [1, 64, 64, 256])
+    self.assertEqual(features[3].get_shape().as_list(), [1, 32, 32, 512])
+    self.assertEqual(features[4].get_shape().as_list(), [1, 16, 16, 1024])
+    self.assertEqual(features[5].get_shape().as_list(), [1, 8, 8, 2048])
+
+if __name__ == '__main__':
+  assert tf.version.VERSION.startswith('2.')
+  tf.test.main()

official/vision/detection/modeling/base_model_test.py

+# Copyright 2019 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 base_model.py."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+
+from absl import logging
+import numpy as np
+import tensorflow.compat.v2 as tf
+
+from official.vision.detection.dataloader import mode_keys
+from official.vision.detection.modeling import base_model
+from official.modeling.hyperparams import params_dict
+
+
+class DummyModel(base_model.Model):
+
+  def build_model(self):
+    input_shape = [1]
+    input_layer = tf.keras.layers.Input(shape=input_shape)
+    outputs = self.model_outputs(inputs=input_layer, mode=None)
+
+    model = tf.keras.models.Model(
+        inputs=input_layer, outputs=outputs, name='dummy')
+    model.optimizer = self.build_optimizer()
+
+    return model
+
+  def build_loss_fn(self):
+    return tf.keras.losses.MeanSquaredError()
+
+  def build_outputs(self, features, mode):
+    return tf.keras.layers.Dense(1)(features)
+
+
+class BaseModelTest(tf.test.TestCase):
+
+  def setUp(self):
+    super(BaseModelTest, self).setUp()
+    self._model_dir = os.path.join(self.get_temp_dir(),
+                                   'model_dir')
+
+  def testBaseModelTrainAndEval(self):
+    params = params_dict.ParamsDict({
+        'batch_size': 1,
+        'model_dir': self._model_dir,
+        'train': {
+            'optimizer': {
+                'type': 'momentum',
+                'momentum': 0.9,
+            },
+            'learning_rate': {
+                'type': 'step',
+                'init_learning_rate': 0.2,
+                'warmup_learning_rate': 0.1,
+                'warmup_steps': 100,
+                'learning_rate_levels': [0.02, 0.002],
+                'learning_rate_steps': [200, 400],
+            },
+            'checkpoint': {
+                'path': '',
+                'prefix': '',
+                'skip_checkpoint_variables': True,
+            },
+            'iterations_per_loop': 1,
+            'frozen_variable_prefix': 'resnet50_conv2',
+        },
+        'enable_summary': False,
+        'architecture': {
+            'use_bfloat16': False,
+        },
+    })
+
+    def _input_fn(params):
+      features = tf.data.Dataset.from_tensor_slices([[1], [2], [3]])
+      labels = tf.data.Dataset.from_tensor_slices([[1], [2], [3]])
+      data = tf.data.Dataset.zip((features, labels)).repeat()
+      dataset = data.batch(params['batch_size'], drop_remainder=True)
+      return dataset
+
+    model_factory = DummyModel(params)
+
+    # Use local TPU for testing.
+    resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='')
+    tf.tpu.experimental.initialize_tpu_system(resolver)
+    strategy = tf.distribute.experimental.TPUStrategy(resolver)
+
+    with tf.device(''):
+      with strategy.scope():
+        model = model_factory.build_model()
+        metrics = [tf.keras.metrics.MeanSquaredError()]
+        loss = model_factory.build_loss_fn()
+        model.compile(optimizer=model.optimizer, loss=loss, metrics=metrics)
+        model.summary()
+
+      training_steps_per_epoch = 3
+      tensorboard_cb = tf.keras.callbacks.TensorBoard(log_dir=self._model_dir)
+      weights_file_path = os.path.join(self._model_dir,
+                                       'weights.{epoch:02d}-{val_loss:.2f}.tf')
+      checkpoint_cb = tf.keras.callbacks.ModelCheckpoint(weights_file_path)
+
+      training_callbacks = [checkpoint_cb, tensorboard_cb]
+
+      model.fit(
+          _input_fn({'batch_size': params.batch_size}),
+          epochs=2,
+          steps_per_epoch=training_steps_per_epoch,
+          callbacks=training_callbacks,
+          validation_data=_input_fn({'batch_size': params.batch_size}),
+          validation_steps=1,
+          validation_freq=1)
+      model.evaluate(_input_fn({'batch_size': params.batch_size}), steps=3)
+
+    out_files = tf.io.gfile.glob(os.path.join(self._model_dir, '*'))
+    logging.info('Model output files: %s', out_files)
+    self.assertNotEmpty(out_files)
+
+
+if __name__ == '__main__':
+  assert tf.version.VERSION.startswith('2.')
+  logging.set_verbosity(logging.INFO)
+  tf.test.main()

official/vision/detection/modeling/learning_rates_test.py

+# Copyright 2018 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 learning_rates.py."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from absl import logging
+import tensorflow.compat.v2 as tf
+
+from official.vision.detection.modeling import learning_rates
+from official.modeling.hyperparams import params_dict
+
+
+class StepLearningRateWithLinearWarmupTest(tf.test.TestCase):
+
+  def test_step_learning_rate_with_linear_warmup(self):
+    params = params_dict.ParamsDict({
+        'type': 'step',
+        'init_learning_rate': 0.2,
+        'warmup_learning_rate': 0.1,
+        'warmup_steps': 100,
+        'learning_rate_levels': [0.02, 0.002],
+        'learning_rate_steps': [200, 400],
+    })
+    learning_rate_fn = learning_rates.learning_rate_generator(params)
+    lr = learning_rate_fn(0).numpy()
+    self.assertAlmostEqual(0.1, lr)
+    lr = learning_rate_fn(50).numpy()
+    self.assertAlmostEqual(0.15, lr)
+    lr = learning_rate_fn(100).numpy()
+    self.assertAlmostEqual(0.2, lr)
+    lr = learning_rate_fn(150).numpy()
+    self.assertAlmostEqual(0.2, lr)
+    lr = learning_rate_fn(200).numpy()
+    self.assertAlmostEqual(0.02, lr)
+    lr = learning_rate_fn(300).numpy()
+    self.assertAlmostEqual(0.02, lr)
+    lr = learning_rate_fn(400).numpy()
+    self.assertAlmostEqual(0.002, lr)
+    lr = learning_rate_fn(500).numpy()
+    self.assertAlmostEqual(0.002, lr)
+    lr = learning_rate_fn(600).numpy()
+    self.assertAlmostEqual(0.002, lr)
+
+
+class CosinLearningRateWithLinearWarmupTest(tf.test.TestCase):
+
+  def test_cosine_learning_rate_with_linear_warmup(self):
+    params = params_dict.ParamsDict({
+        'type': 'cosine',
+        'init_learning_rate': 0.2,
+        'warmup_learning_rate': 0.1,
+        'warmup_steps': 100,
+        'total_steps': 1100,
+    })
+    learning_rate_fn = learning_rates.learning_rate_generator(params)
+    lr = learning_rate_fn(0).numpy()
+    self.assertAlmostEqual(0.1, lr)
+    lr = learning_rate_fn(50).numpy()
+    self.assertAlmostEqual(0.15, lr)
+    lr = learning_rate_fn(100).numpy()
+    self.assertAlmostEqual(0.2, lr)
+    lr = learning_rate_fn(350).numpy()
+    self.assertAlmostEqual(0.17071067811865476, lr)
+    lr = learning_rate_fn(600).numpy()
+    self.assertAlmostEqual(0.1, lr)
+    lr = learning_rate_fn(850).numpy()
+    self.assertAlmostEqual(0.029289321881345254, lr)
+    lr = learning_rate_fn(1100).numpy()
+    self.assertAlmostEqual(0.0, lr)
+
+
+if __name__ == '__main__':
+  assert tf.version.VERSION.startswith('2.')
+  logging.set_verbosity(logging.INFO)
+  tf.test.main()

official/vision/detection/modeling/postprocess_test.py

+# Copyright 2019 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 postprocess.py."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import numpy as np
+import tensorflow.compat.v2 as tf
+
+from absl.testing import parameterized
+from official.vision.detection.dataloader import anchor
+from official.vision.detection.modeling import postprocess
+from official.modeling.hyperparams import params_dict
+
+
+class GenerateDetectionsTest(parameterized.TestCase, tf.test.TestCase):
+
+  @parameterized.parameters(
+      (True),
+      (False),
+  )
+  def testDetectionsOutputShape(self, use_batched_nms):
+    min_level = 4
+    max_level = 6
+    num_scales = 2
+    max_total_size = 100
+    aspect_ratios = [1.0, 2.0,]
+    anchor_scale = 2.0
+    output_size = [64, 64]
+    num_classes = 4
+    # pre_nms_num_boxes = 5000
+    score_threshold = 0.01
+    batch_size = 1
+    postprocessor_params = params_dict.ParamsDict({
+        'use_batched_nms': use_batched_nms,
+        'max_total_size': max_total_size,
+        'nms_iou_threshold': 0.5,
+        'score_threshold': score_threshold,
+        'min_level': min_level,
+        'max_level': max_level,
+        'num_classes': num_classes,
+    })
+
+    input_anchor = anchor.Anchor(min_level, max_level, num_scales,
+                                 aspect_ratios, anchor_scale, output_size)
+    cls_outputs_all = (np.random.rand(84, num_classes) -
+                       0.5) * 3  # random 84x3 outputs.
+    box_outputs_all = np.random.rand(84, 4)  # random 84 boxes.
+    class_outputs = {
+        4:
+            tf.reshape(
+                tf.convert_to_tensor(
+                    value=cls_outputs_all[0:64], dtype=tf.float32),
+                [1, 8, 8, num_classes]),
+        5:
+            tf.reshape(
+                tf.convert_to_tensor(
+                    value=cls_outputs_all[64:80], dtype=tf.float32),
+                [1, 4, 4, num_classes]),
+        6:
+            tf.reshape(
+                tf.convert_to_tensor(
+                    value=cls_outputs_all[80:84], dtype=tf.float32),
+                [1, 2, 2, num_classes]),
+    }
+    box_outputs = {
+        4:
+            tf.reshape(
+                tf.convert_to_tensor(
+                    value=box_outputs_all[0:64], dtype=tf.float32),
+                [1, 8, 8, 4]),
+        5:
+            tf.reshape(
+                tf.convert_to_tensor(
+                    value=box_outputs_all[64:80], dtype=tf.float32),
+                [1, 4, 4, 4]),
+        6:
+            tf.reshape(
+                tf.convert_to_tensor(
+                    value=box_outputs_all[80:84], dtype=tf.float32),
+                [1, 2, 2, 4]),
+    }
+    image_info = tf.constant([[[1000, 1000], [100, 100], [0.1, 0.1], [0, 0]]],
+                             dtype=tf.float32)
+    predict_fn = postprocess.GenerateOneStageDetections(postprocessor_params)
+    boxes, scores, classes, valid_detections = predict_fn(
+        inputs=(box_outputs, class_outputs, input_anchor.multilevel_boxes,
+                image_info[:, 1:2, :]))
+    (boxes, scores, classes, valid_detections) = [
+        boxes.numpy(),
+        scores.numpy(),
+        classes.numpy(),
+        valid_detections.numpy()
+    ]
+    self.assertEqual(boxes.shape, (batch_size, max_total_size, 4))
+    self.assertEqual(scores.shape, (
+        batch_size,
+        max_total_size,
+    ))
+    self.assertEqual(classes.shape, (
+        batch_size,
+        max_total_size,
+    ))
+    self.assertEqual(valid_detections.shape, (batch_size,))
+
+
+if __name__ == '__main__':
+  assert tf.version.VERSION.startswith('2.')
+  tf.test.main()

official/vision/detection/utils/box_utils_test.py

+# Copyright 2019 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 box_utils.py."""
+
+import numpy as np
+import tensorflow.compat.v2 as tf
+from official.vision.detection.utils import box_utils
+
+
+def _transform_boxes_on_tpu_and_cpu(transform_fn, boxes, arg):
+  # Runs on TPU.
+  strategy = tf.distribute.experimental.TPUStrategy()
+  with strategy.scope():
+    normalized_op_tpu = transform_fn(boxes, arg)
+    normalized_boxes_tpu = normalized_op_tpu.numpy()
+
+  # Runs on CPU.
+  normalize_op = transform_fn(boxes, arg)
+  normalized_boxes_cpu = normalize_op.numpy()
+  return normalized_boxes_tpu, normalized_boxes_cpu
+
+
+class NormalizeBoxesTest(tf.test.TestCase):
+
+  def testNormalizeBoxes1DWithImageShapeAsList(self):
+    boxes = tf.constant([10, 30, 40, 90], tf.float32)
+    image_shape = [50, 100]
+
+    normalized_boxes_tpu, normalized_boxes_cpu = _transform_boxes_on_tpu_and_cpu(
+        box_utils.normalize_boxes, boxes, image_shape)
+
+    self.assertNDArrayNear(normalized_boxes_tpu, normalized_boxes_cpu, 1e-5)
+    self.assertNDArrayNear(normalized_boxes_tpu, [0.2, 0.3, 0.8, 0.9], 1e-5)
+
+  def testNormalizeBoxes1DWithImageShapeAsTensor(self):
+    boxes = tf.constant([10, 30, 40, 90], tf.float32)
+    image_shape = tf.constant([50, 100], tf.int32)
+
+    normalized_boxes_tpu, normalized_boxes_cpu = _transform_boxes_on_tpu_and_cpu(
+        box_utils.normalize_boxes, boxes, image_shape)
+
+    self.assertNDArrayNear(normalized_boxes_tpu, normalized_boxes_cpu, 1e-5)
+    self.assertNDArrayNear(normalized_boxes_tpu, [0.2, 0.3, 0.8, 0.9], 1e-5)
+
+  def testNormalizeBoxes2DWithImageShapeAsList(self):
+    boxes = tf.constant([[10, 30, 40, 90], [30, 10, 40, 50]], tf.float32)
+    image_shape = [50, 100]
+
+    normalized_boxes_tpu, normalized_boxes_cpu = _transform_boxes_on_tpu_and_cpu(
+        box_utils.normalize_boxes, boxes, image_shape)
+
+    self.assertNDArrayNear(normalized_boxes_tpu, normalized_boxes_cpu, 1e-5)
+    self.assertNDArrayNear(normalized_boxes_tpu,
+                           [[0.2, 0.3, 0.8, 0.9], [0.6, 0.1, 0.8, 0.5]], 1e-5)
+
+  def testNormalizeBoxes2DWithImageShapeAsVector(self):
+    boxes = tf.constant([[10, 30, 40, 90], [30, 10, 40, 50]], tf.float32)
+    image_shape = tf.constant([50, 100], dtype=tf.int32)
+
+    normalized_boxes_tpu, normalized_boxes_cpu = _transform_boxes_on_tpu_and_cpu(
+        box_utils.normalize_boxes, boxes, image_shape)
+
+    self.assertNDArrayNear(normalized_boxes_tpu, normalized_boxes_cpu, 1e-5)
+    self.assertNDArrayNear(normalized_boxes_tpu,
+                           [[0.2, 0.3, 0.8, 0.9], [0.6, 0.1, 0.8, 0.5]], 1e-5)
+
+  def testNormalizeBoxes2DWithImageShapeAsBroadcastableTensor(self):
+    boxes = tf.constant([[10, 30, 40, 90], [30, 10, 40, 50]], tf.float32)
+    image_shape = tf.constant([[50, 100]], tf.int32)
+
+    normalized_boxes_tpu, normalized_boxes_cpu = _transform_boxes_on_tpu_and_cpu(
+        box_utils.normalize_boxes, boxes, image_shape)
+
+    self.assertNDArrayNear(normalized_boxes_tpu, normalized_boxes_cpu, 1e-5)
+    self.assertNDArrayNear(normalized_boxes_tpu,
+                           [[0.2, 0.3, 0.8, 0.9], [0.6, 0.1, 0.8, 0.5]], 1e-5)
+
+  def testNormalizeBoxes2DWithImageShapeAsSameShapeTensor(self):
+    boxes = tf.constant([[10, 30, 40, 90], [30, 10, 40, 50]], tf.float32)
+    image_shape = tf.constant([[50, 100], [50, 100]], tf.int32)
+
+    normalized_boxes_tpu, normalized_boxes_cpu = _transform_boxes_on_tpu_and_cpu(
+        box_utils.normalize_boxes, boxes, image_shape)
+
+    self.assertNDArrayNear(normalized_boxes_tpu, normalized_boxes_cpu, 1e-5)
+    self.assertNDArrayNear(normalized_boxes_tpu,
+                           [[0.2, 0.3, 0.8, 0.9], [0.6, 0.1, 0.8, 0.5]], 1e-5)
+
+  def testNormalizeBoxes3DWithImageShapeAsList(self):
+    boxes = tf.constant([[[10, 30, 40, 90], [30, 10, 40, 50]],
+                         [[20, 40, 50, 80], [30, 50, 40, 90]]], tf.float32)
+    image_shape = [50, 100]
+
+    normalized_boxes_tpu, normalized_boxes_cpu = _transform_boxes_on_tpu_and_cpu(
+        box_utils.normalize_boxes, boxes, image_shape)
+
+    self.assertNDArrayNear(normalized_boxes_tpu, normalized_boxes_cpu, 1e-5)
+    self.assertNDArrayNear(normalized_boxes_tpu,
+                           [[[0.2, 0.3, 0.8, 0.9], [0.6, 0.1, 0.8, 0.5]],
+                            [[0.4, 0.4, 1.0, 0.8], [0.6, 0.5, 0.8, 0.9]]], 1e-5)
+
+  def testNormalizeBoxes3DWithImageShapeAsVector(self):
+    boxes = tf.constant([[[10, 30, 40, 90], [30, 10, 40, 50]],
+                         [[20, 40, 50, 80], [30, 50, 40, 90]]], tf.float32)
+    image_shape = tf.constant([50, 100], tf.int32)
+
+    normalized_boxes_tpu, normalized_boxes_cpu = _transform_boxes_on_tpu_and_cpu(
+        box_utils.normalize_boxes, boxes, image_shape)
+
+    self.assertNDArrayNear(normalized_boxes_tpu, normalized_boxes_cpu, 1e-5)
+    self.assertNDArrayNear(normalized_boxes_tpu,
+                           [[[0.2, 0.3, 0.8, 0.9], [0.6, 0.1, 0.8, 0.5]],
+                            [[0.4, 0.4, 1.0, 0.8], [0.6, 0.5, 0.8, 0.9]]], 1e-5)
+
+  def testNormalizeBoxes3DWithImageShapeAsBroadcastableTensor(self):
+    boxes = tf.constant([[[10, 30, 40, 90], [30, 10, 40, 50]],
+                         [[20, 40, 50, 80], [30, 50, 40, 90]]], tf.float32)
+    image_shape = tf.constant([[[50, 100]], [[500, 1000]]], tf.int32)
+
+    normalized_boxes_tpu, normalized_boxes_cpu = _transform_boxes_on_tpu_and_cpu(
+        box_utils.normalize_boxes, boxes, image_shape)
+
+    self.assertNDArrayNear(normalized_boxes_tpu, normalized_boxes_cpu, 1e-5)
+    self.assertNDArrayNear(
+        normalized_boxes_tpu,
+        [[[0.2, 0.3, 0.8, 0.9], [0.6, 0.1, 0.8, 0.5]],
+         [[0.04, 0.04, 0.1, 0.08], [0.06, 0.05, 0.08, 0.09]]], 1e-5)
+
+  def testNormalizeBoxes3DWithImageShapeAsSameShapeTensor(self):
+    boxes = tf.constant([[[10, 30, 40, 90], [30, 10, 40, 50]],
+                         [[20, 40, 50, 80], [30, 50, 40, 90]]], tf.float32)
+    image_shape = tf.constant(
+        [[[50, 100], [50, 100]], [[500, 1000], [500, 1000]]], tf.int32)
+
+    normalized_boxes_tpu, normalized_boxes_cpu = _transform_boxes_on_tpu_and_cpu(
+        box_utils.normalize_boxes, boxes, image_shape)
+
+    self.assertNDArrayNear(normalized_boxes_tpu, normalized_boxes_cpu, 1e-5)
+    self.assertNDArrayNear(
+        normalized_boxes_tpu,
+        [[[0.2, 0.3, 0.8, 0.9], [0.6, 0.1, 0.8, 0.5]],
+         [[0.04, 0.04, 0.1, 0.08], [0.06, 0.05, 0.08, 0.09]]], 1e-5)
+
+
+class DenormalizeBoxesTest(tf.test.TestCase):
+
+  def testDenormalizeBoxes1DWithImageShapeAsList(self):
+    boxes = tf.constant([0.2, 0.3, 0.8, 0.9], tf.float32)
+    image_shape = [50, 100]
+    normalized_boxes_tpu, normalized_boxes_cpu = _transform_boxes_on_tpu_and_cpu(
+        box_utils.denormalize_boxes, boxes, image_shape)
+
+    self.assertNDArrayNear(normalized_boxes_tpu, normalized_boxes_cpu, 1e-5)
+    self.assertNDArrayNear(normalized_boxes_tpu, [10, 30, 40, 90], 1e-5)
+
+  def testDenormalizeBoxes1DWithImageShapeAsTensor(self):
+    boxes = tf.constant([0.2, 0.3, 0.8, 0.9], tf.float32)
+    image_shape = tf.constant([50, 100], tf.int32)
+    normalized_boxes_tpu, normalized_boxes_cpu = _transform_boxes_on_tpu_and_cpu(
+        box_utils.denormalize_boxes, boxes, image_shape)
+
+    self.assertNDArrayNear(normalized_boxes_tpu, normalized_boxes_cpu, 1e-5)
+    self.assertNDArrayNear(normalized_boxes_tpu, [10, 30, 40, 90], 1e-5)
+
+  def testDenormalizeBoxes2DWithImageShapeAsList(self):
+    boxes = tf.constant([[0.2, 0.3, 0.8, 0.9], [0.6, 0.1, 0.8, 0.5]],
+                        tf.float32)
+    image_shape = [50, 100]
+    normalized_boxes_tpu, normalized_boxes_cpu = _transform_boxes_on_tpu_and_cpu(
+        box_utils.denormalize_boxes, boxes, image_shape)
+
+    self.assertNDArrayNear(normalized_boxes_tpu, normalized_boxes_cpu, 1e-5)
+    self.assertNDArrayNear(normalized_boxes_tpu,
+                           [[10, 30, 40, 90], [30, 10, 40, 50]], 1e-5)
+
+  def testDenormalizeBoxes2DWithImageShapeAsVector(self):
+    boxes = tf.constant([[0.2, 0.3, 0.8, 0.9], [0.6, 0.1, 0.8, 0.5]],
+                        tf.float32)
+    image_shape = tf.constant([50, 100], dtype=tf.int32)
+    normalized_boxes_tpu, normalized_boxes_cpu = _transform_boxes_on_tpu_and_cpu(
+        box_utils.denormalize_boxes, boxes, image_shape)
+
+    self.assertNDArrayNear(normalized_boxes_tpu, normalized_boxes_cpu, 1e-5)
+    self.assertNDArrayNear(normalized_boxes_tpu,
+                           [[10, 30, 40, 90], [30, 10, 40, 50]], 1e-5)
+
+  def testDenormalizeBoxes2DWithImageShapeAsBroadcastableTensor(self):
+    boxes = tf.constant([[0.2, 0.3, 0.8, 0.9], [0.6, 0.1, 0.8, 0.5]],
+                        tf.float32)
+    image_shape = tf.constant([[50, 100]], tf.int32)
+    normalized_boxes_tpu, normalized_boxes_cpu = _transform_boxes_on_tpu_and_cpu(
+        box_utils.denormalize_boxes, boxes, image_shape)
+
+    self.assertNDArrayNear(normalized_boxes_tpu, normalized_boxes_cpu, 1e-5)
+    self.assertNDArrayNear(normalized_boxes_tpu,
+                           [[10, 30, 40, 90], [30, 10, 40, 50]], 1e-5)
+
+  def testDenormalizeBoxes2DWithImageShapeAsSameShapeTensor(self):
+    boxes = tf.constant([[0.2, 0.3, 0.8, 0.9], [0.6, 0.1, 0.8, 0.5]],
+                        tf.float32)
+    image_shape = tf.constant([[50, 100], [50, 100]], tf.int32)
+    normalized_boxes_tpu, normalized_boxes_cpu = _transform_boxes_on_tpu_and_cpu(
+        box_utils.denormalize_boxes, boxes, image_shape)
+
+    self.assertNDArrayNear(normalized_boxes_tpu, normalized_boxes_cpu, 1e-5)
+    self.assertNDArrayNear(normalized_boxes_tpu,
+                           [[10, 30, 40, 90], [30, 10, 40, 50]], 1e-5)
+
+  def testDenormalizeBoxes3DWithImageShapeAsList(self):
+    boxes = tf.constant([[[0.2, 0.3, 0.8, 0.9], [0.6, 0.1, 0.8, 0.5]],
+                         [[0.4, 0.4, 1.0, 0.8], [0.6, 0.5, 0.8, 0.9]]],
+                        tf.float32)
+    image_shape = [50, 100]
+    normalized_boxes_tpu, normalized_boxes_cpu = _transform_boxes_on_tpu_and_cpu(
+        box_utils.denormalize_boxes, boxes, image_shape)
+
+    self.assertNDArrayNear(normalized_boxes_tpu, normalized_boxes_cpu, 1e-5)
+    self.assertNDArrayNear(normalized_boxes_tpu,
+                           [[[10, 30, 40, 90], [30, 10, 40, 50]],
+                            [[20, 40, 50, 80], [30, 50, 40, 90]]], 1e-5)
+
+  def testDenormalizeBoxes3DWithImageShapeAsVector(self):
+    boxes = tf.constant([[[0.2, 0.3, 0.8, 0.9], [0.6, 0.1, 0.8, 0.5]],
+                         [[0.4, 0.4, 1.0, 0.8], [0.6, 0.5, 0.8, 0.9]]],
+                        tf.float32)
+    image_shape = tf.constant([50, 100], tf.int32)
+    normalized_boxes_tpu, normalized_boxes_cpu = _transform_boxes_on_tpu_and_cpu(
+        box_utils.denormalize_boxes, boxes, image_shape)
+
+    self.assertNDArrayNear(normalized_boxes_tpu, normalized_boxes_cpu, 1e-5)
+    self.assertNDArrayNear(normalized_boxes_tpu,
+                           [[[10, 30, 40, 90], [30, 10, 40, 50]],
+                            [[20, 40, 50, 80], [30, 50, 40, 90]]], 1e-5)
+
+  def testDenormalizeBoxes3DWithImageShapeAsBroadcastableTensor(self):
+    boxes = tf.constant([[[0.2, 0.3, 0.8, 0.9], [0.6, 0.1, 0.8, 0.5]],
+                         [[0.04, 0.04, 0.1, 0.08], [0.06, 0.05, 0.08, 0.09]]],
+                        tf.float32)
+    image_shape = tf.constant([[[50, 100]], [[500, 1000]]], tf.int32)
+    normalized_boxes_tpu, normalized_boxes_cpu = _transform_boxes_on_tpu_and_cpu(
+        box_utils.denormalize_boxes, boxes, image_shape)
+
+    self.assertNDArrayNear(normalized_boxes_tpu, normalized_boxes_cpu, 1e-5)
+    self.assertNDArrayNear(normalized_boxes_tpu,
+                           [[[10, 30, 40, 90], [30, 10, 40, 50]],
+                            [[20, 40, 50, 80], [30, 50, 40, 90]]], 1e-5)
+
+  def testDenormalizeBoxes3DWithImageShapeAsSameShapeTensor(self):
+    boxes = tf.constant([[[0.2, 0.3, 0.8, 0.9], [0.6, 0.1, 0.8, 0.5]],
+                         [[0.04, 0.04, 0.1, 0.08], [0.06, 0.05, 0.08, 0.09]]],
+                        tf.float32)
+    image_shape = tf.constant(
+        [[[50, 100], [50, 100]], [[500, 1000], [500, 1000]]], tf.int32)
+    normalized_boxes_tpu, normalized_boxes_cpu = _transform_boxes_on_tpu_and_cpu(
+        box_utils.denormalize_boxes, boxes, image_shape)
+
+    self.assertNDArrayNear(normalized_boxes_tpu, normalized_boxes_cpu, 1e-5)
+    self.assertNDArrayNear(normalized_boxes_tpu,
+                           [[[10, 30, 40, 90], [30, 10, 40, 50]],
+                            [[20, 40, 50, 80], [30, 50, 40, 90]]], 1e-5)
+
+
+class ClipBoxesTest(tf.test.TestCase):
+
+  def testClipBoxesImageShapeAsList(self):
+    boxes_data = [[0, 0, 1, 1], [0, 0.1, 1, 1.1], [0, 0.3, 1, 1.3],
+                  [0, 0.5, 1, 1.5], [0, 0.7, 1, 1.7], [0, 1.9, 1, 1.9]]
+    image_shape = [3, 3]
+    boxes = tf.constant(boxes_data)
+
+    clipped_boxes_tpu, clipped_boxes_cpu = _transform_boxes_on_tpu_and_cpu(
+        box_utils.clip_boxes, boxes, image_shape)
+
+    self.assertAllClose(clipped_boxes_tpu, clipped_boxes_cpu)
+    self.assertAllClose(clipped_boxes_tpu,
+                        [[0, 0, 1, 1], [0, 0.1, 1, 1.1], [0, 0.3, 1, 1.3],
+                         [0, 0.5, 1, 1.5], [0, 0.7, 1, 1.7], [0, 1.9, 1, 1.9]])
+
+  def testClipBoxesImageShapeAsVector(self):
+    boxes_data = [[0, 0, 1, 1], [0, 0.1, 1, 1.1], [0, 0.3, 1, 1.3],
+                  [0, 0.5, 1, 1.5], [0, 0.7, 1, 1.7], [0, 1.9, 1, 1.9]]
+    boxes = tf.constant(boxes_data)
+    image_shape = np.array([3, 3], dtype=np.float32)
+    clipped_boxes_tpu, clipped_boxes_cpu = _transform_boxes_on_tpu_and_cpu(
+        box_utils.clip_boxes, boxes, image_shape)
+
+    self.assertAllClose(clipped_boxes_tpu, clipped_boxes_cpu)
+    self.assertAllClose(clipped_boxes_tpu,
+                        [[0, 0, 1, 1], [0, 0.1, 1, 1.1], [0, 0.3, 1, 1.3],
+                         [0, 0.5, 1, 1.5], [0, 0.7, 1, 1.7], [0, 1.9, 1, 1.9]])
+
+  def testClipBoxesImageShapeAsTensor(self):
+    boxes_data = [[0, 0, 1, 1], [0, 0.1, 1, 1.1], [0, 0.3, 1, 1.3],
+                  [0, 0.5, 1, 1.5], [0, 0.7, 1, 1.7], [0, 1.9, 1, 1.9]]
+    boxes = tf.constant(boxes_data)
+    image_shape = tf.constant([[3, 3], [3, 3], [3, 3], [3, 3], [3, 3], [3, 3]],
+                              dtype=tf.float32)
+    clipped_boxes_tpu, clipped_boxes_cpu = _transform_boxes_on_tpu_and_cpu(
+        box_utils.clip_boxes, boxes, image_shape)
+
+    self.assertAllClose(clipped_boxes_tpu, clipped_boxes_cpu)
+    self.assertAllClose(clipped_boxes_tpu,
+                        [[0, 0, 1, 1], [0, 0.1, 1, 1.1], [0, 0.3, 1, 1.3],
+                         [0, 0.5, 1, 1.5], [0, 0.7, 1, 1.7], [0, 1.9, 1, 1.9]])
+
+
+class EncodeDecodeBoxesTest(tf.test.TestCase):
+
+  def test_encode_decode_boxes(self):
+    boxes_np = np.array([[[1.0, 2.0, 3.0, 4.0], [2.0, 3.0, 4.0, 5.0]],
+                         [[4.0, 5.0, 6.0, 7.0], [5.0, 6.0, 7.0, 8.0]]])
+    boxes = tf.constant(boxes_np, dtype=tf.float32)
+    anchors = tf.constant([[[1.5, 2.5, 3.5, 4.5], [2.5, 3.5, 4.5, 5.5]],
+                           [[1.5, 2.5, 3.5, 4.5], [2.5, 3.5, 4.5, 5.5]]],
+                          dtype=tf.float32)
+    weights = [1.0, 1.0, 1.0, 1.0]
+
+    def test_fn(boxes, anchors):
+      encoded_boxes = box_utils.encode_boxes(boxes, anchors, weights)
+      decoded_boxes = box_utils.decode_boxes(encoded_boxes, anchors, weights)
+      return decoded_boxes
+
+    decoded_boxes_tpu, decoded_boxes_cpu = _transform_boxes_on_tpu_and_cpu(
+        test_fn, boxes, anchors)
+
+    self.assertNDArrayNear(decoded_boxes_tpu, decoded_boxes_cpu, 1e-5)
+    self.assertNDArrayNear(decoded_boxes_tpu, boxes_np, 1e-5)
+
+
+if __name__ == '__main__':
+  assert tf.version.VERSION.startswith('2.')
+  tf.test.main()

official/vision/detection/utils/input_utils_test.py

+# Copyright 2019 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 input_utils.py."""
+
+import numpy as np
+import tensorflow.compat.v2 as tf
+from absl.testing import parameterized
+from official.vision.detection.utils import input_utils
+
+
+class InputUtilsTest(parameterized.TestCase, tf.test.TestCase):
+
+  @parameterized.parameters(
+      ([1], 10),
+      ([1, 2], 10),
+      ([1, 2, 3,], 10),
+  )
+  def testPadToFixedSize(self, input_shape, output_size):
+    # Copies input shape to padding shape.
+    padding_shape = input_shape[:]
+    padding_shape[0] = output_size - input_shape[0]
+    expected_outputs = np.concatenate(
+        [np.ones(input_shape), np.zeros(padding_shape)], axis=0)
+
+    data = tf.ones(input_shape)
+    output_data = input_utils.pad_to_fixed_size(
+        data, output_size, constant_values=0)
+    output_data = output_data.numpy()
+    self.assertAllClose(output_size, output_data.shape[0])
+    self.assertAllClose(expected_outputs, output_data)
+
+  @parameterized.parameters(
+      (100, 200, 100, 200, 32, 1.0, 1.0, 100, 200, 128, 224),
+      (100, 256, 128, 256, 32, 1.0, 1.0, 100, 256, 128, 256),
+      (200, 512, 200, 128, 32, 0.25, 0.25, 50, 128, 224, 128),
+  )
+  def testResizeAndCropImageRectangluarCase(self,
+                                            input_height,
+                                            input_width,
+                                            desired_height,
+                                            desired_width,
+                                            stride,
+                                            scale_y,
+                                            scale_x,
+                                            scaled_height,
+                                            scaled_width,
+                                            output_height,
+                                            output_width):
+    image = tf.convert_to_tensor(
+        value=np.random.rand(input_height, input_width, 3))
+
+    desired_size = (desired_height, desired_width)
+    resized_image, image_info = input_utils.resize_and_crop_image(
+        image,
+        desired_size=desired_size,
+        padded_size=input_utils.compute_padded_size(desired_size, stride))
+    resized_image_shape = tf.shape(input=resized_image)
+    resized_shape_np = resized_image_shape.numpy()
+    image_info_np = image_info.numpy()
+
+    self.assertAllEqual([output_height, output_width, 3], resized_shape_np)
+    self.assertNDArrayNear(
+        [[input_height, input_width], [scaled_height, scaled_width],
+         [scale_y, scale_x], [0.0, 0.0]], image_info_np, 1e-5)
+
+  @parameterized.parameters(
+      (100, 200, 220, 220, 32, 1.1, 1.1, 110, 220, 224, 224),
+      (512, 512, 1024, 1024, 32, 2.0, 2.0, 1024, 1024, 1024, 1024),
+  )
+  def testResizeAndCropImageSquareCase(self,
+                                       input_height,
+                                       input_width,
+                                       desired_height,
+                                       desired_width,
+                                       stride,
+                                       scale_y,
+                                       scale_x,
+                                       scaled_height,
+                                       scaled_width,
+                                       output_height,
+                                       output_width):
+    image = tf.convert_to_tensor(
+        value=np.random.rand(input_height, input_width, 3))
+
+    desired_size = (desired_height, desired_width)
+    resized_image, image_info = input_utils.resize_and_crop_image(
+        image,
+        desired_size=desired_size,
+        padded_size=input_utils.compute_padded_size(desired_size, stride))
+    resized_image_shape = tf.shape(input=resized_image)
+
+    resized_shape_np, image_info_np = ([
+        resized_image_shape.numpy(),
+        image_info.numpy()
+    ])
+    self.assertAllEqual([output_height, output_width, 3], resized_shape_np)
+    self.assertNDArrayNear(
+        [[input_height, input_width], [scaled_height, scaled_width],
+         [scale_y, scale_x], [0.0, 0.0]], image_info_np, 1e-5)
+
+
+if __name__ == '__main__':
+  assert tf.version.VERSION.startswith('2.')
+  tf.test.main()

official/vision/detection/utils/spatial_transform_test.py

+# Copyright 2019 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 spatial_transform.py."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import sys
+
+from absl import logging
+import numpy as np
+import tensorflow.compat.v2 as tf
+
+from official.vision.detection.utils import spatial_transform
+
+
+class MultiLevelCropAndResizeTest(tf.test.TestCase):
+
+  def test_multilevel_crop_and_resize_square(self):
+    """Example test case.
+
+    Input =
+    [
+      [0, 1, 2, 3],
+      [4, 5, 6, 7],
+      [8, 9, 10, 11],
+      [12, 13, 14, 15],
+    ]
+    output_size = 2x2
+    box =
+    [
+      [[0, 0, 2, 2]]
+    ]
+    Gathered data =
+    [
+      [0, 1, 1, 2],
+      [4, 5, 5, 6],
+      [4, 5, 5, 6],
+      [8, 9, 9, 10],
+    ]
+    Interpolation kernel =
+    [
+      [1, 1, 1, 1],
+      [1, 1, 1, 1],
+      [1, 1, 1, 1],
+      [1, 1, 1, 1],
+    ]
+    Output =
+    [
+      [2.5, 3.5],
+      [6.5, 7.5]
+    ]
+    """
+    input_size = 4
+    min_level = 0
+    max_level = 0
+    batch_size = 1
+    output_size = 2
+    num_filters = 1
+    features = {}
+    for level in range(min_level, max_level + 1):
+      feat_size = int(input_size / 2**level)
+
+      features[level] = tf.range(
+          batch_size * feat_size * feat_size * num_filters, dtype=tf.float32)
+      features[level] = tf.reshape(
+          features[level], [batch_size, feat_size, feat_size, num_filters])
+    boxes = tf.constant([
+        [[0, 0, 2, 2]],
+    ], dtype=tf.float32)
+    tf_roi_features = spatial_transform.multilevel_crop_and_resize(
+        features, boxes, output_size)
+    roi_features = tf_roi_features.numpy()
+    self.assertAllClose(
+        roi_features,
+        np.array([[2.5, 3.5],
+                  [6.5,
+                   7.5]]).reshape([batch_size, 1, output_size, output_size, 1]))
+
+  def test_multilevel_crop_and_resize_rectangle(self):
+    """Example test case.
+
+    Input =
+    [
+      [0, 1, 2, 3],
+      [4, 5, 6, 7],
+      [8, 9, 10, 11],
+      [12, 13, 14, 15],
+    ]
+    output_size = 2x2
+    box =
+    [
+      [[0, 0, 2, 3]]
+    ]
+    Box vertices =
+    [
+      [[0.5, 0.75], [0.5, 2.25]],
+      [[1.5, 0.75], [1.5, 2.25]],
+    ]
+    Gathered data =
+    [
+      [0, 1, 2, 3],
+      [4, 5, 6, 7],
+      [4, 5, 6, 7],
+      [8, 9, 10, 11],
+    ]
+    Interpolation kernel =
+    [
+      [0.5 1.5 1.5 0.5],
+      [0.5 1.5 1.5 0.5],
+      [0.5 1.5 1.5 0.5],
+      [0.5 1.5 1.5 0.5],
+    ]
+    Output =
+    [
+      [2.75, 4.25],
+      [6.75, 8.25]
+    ]
+    """
+    input_size = 4
+    min_level = 0
+    max_level = 0
+    batch_size = 1
+    output_size = 2
+    num_filters = 1
+    features = {}
+    for level in range(min_level, max_level + 1):
+      feat_size = int(input_size / 2**level)
+
+      features[level] = tf.range(
+          batch_size * feat_size * feat_size * num_filters, dtype=tf.float32)
+      features[level] = tf.reshape(
+          features[level], [batch_size, feat_size, feat_size, num_filters])
+    boxes = tf.constant([
+        [[0, 0, 2, 3]],
+    ], dtype=tf.float32)
+    tf_roi_features = spatial_transform.multilevel_crop_and_resize(
+        features, boxes, output_size)
+    roi_features = tf_roi_features.numpy()
+    self.assertAllClose(
+        roi_features,
+        np.array([[2.75, 4.25],
+                  [6.75,
+                   8.25]]).reshape([batch_size, 1, output_size, output_size,
+                                    1]))
+
+  def test_multilevel_crop_and_resize_two_boxes(self):
+    """Test two boxes."""
+    input_size = 4
+    min_level = 0
+    max_level = 0
+    batch_size = 1
+    output_size = 2
+    num_filters = 1
+    features = {}
+    for level in range(min_level, max_level + 1):
+      feat_size = int(input_size / 2**level)
+
+      features[level] = tf.range(
+          batch_size * feat_size * feat_size * num_filters, dtype=tf.float32)
+      features[level] = tf.reshape(
+          features[level], [batch_size, feat_size, feat_size, num_filters])
+    boxes = tf.constant([
+        [[0, 0, 2, 2], [0, 0, 2, 3]],
+    ], dtype=tf.float32)
+    tf_roi_features = spatial_transform.multilevel_crop_and_resize(
+        features, boxes, output_size)
+    roi_features = tf_roi_features.numpy()
+    self.assertAllClose(
+        roi_features,
+        np.array([[[2.5, 3.5], [6.5, 7.5]], [[2.75, 4.25], [6.75, 8.25]]
+                 ]).reshape([batch_size, 2, output_size, output_size, 1]))
+
+  def test_multilevel_crop_and_resize_feature_level_assignment(self):
+    """Test feature level assignment."""
+    input_size = 640
+    min_level = 2
+    max_level = 5
+    batch_size = 1
+    output_size = 2
+    num_filters = 1
+    features = {}
+    for level in range(min_level, max_level + 1):
+      feat_size = int(input_size / 2**level)
+
+      features[level] = float(level) * tf.ones(
+          [batch_size, feat_size, feat_size, num_filters], dtype=tf.float32)
+    boxes = tf.constant(
+        [
+            [
+                [0, 0, 111, 111],  # Level 2.
+                [0, 0, 113, 113],  # Level 3.
+                [0, 0, 223, 223],  # Level 3.
+                [0, 0, 225, 225],  # Level 4.
+                [0, 0, 449, 449]
+            ],  # Level 5.
+        ],
+        dtype=tf.float32)
+    tf_roi_features = spatial_transform.multilevel_crop_and_resize(
+        features, boxes, output_size)
+    roi_features = tf_roi_features.numpy()
+    self.assertAllClose(roi_features[0, 0], 2 * np.ones((2, 2, 1)))
+    self.assertAllClose(roi_features[0, 1], 3 * np.ones((2, 2, 1)))
+    self.assertAllClose(roi_features[0, 2], 3 * np.ones((2, 2, 1)))
+    self.assertAllClose(roi_features[0, 3], 4 * np.ones((2, 2, 1)))
+    self.assertAllClose(roi_features[0, 4], 5 * np.ones((2, 2, 1)))
+
+  def test_multilevel_crop_and_resize_large_input(self):
+    """Test 512 boxes on TPU."""
+    input_size = 1408
+    min_level = 2
+    max_level = 6
+    batch_size = 2
+    num_boxes = 512
+    num_filters = 256
+    output_size = 7
+    strategy = tf.distribute.experimental.TPUStrategy()
+    with strategy.scope():
+      features = {}
+      for level in range(min_level, max_level + 1):
+        feat_size = int(input_size / 2**level)
+        features[level] = tf.constant(
+            np.reshape(
+                np.arange(
+                    batch_size * feat_size * feat_size * num_filters,
+                    dtype=np.float32),
+                [batch_size, feat_size, feat_size, num_filters]),
+            dtype=tf.bfloat16)
+      boxes = np.array([
+          [[0, 0, 256, 256]]*num_boxes,
+      ], dtype=np.float32)
+      boxes = np.tile(boxes, [batch_size, 1, 1])
+      tf_boxes = tf.constant(boxes, dtype=tf.float32)
+
+      tf_roi_features = spatial_transform.multilevel_crop_and_resize(
+          features, tf_boxes)
+      roi_features = tf_roi_features.numpy()
+      self.assertEqual(
+          roi_features.shape,
+          (batch_size, num_boxes, output_size, output_size, num_filters))
+
+  def test_crop_mask_in_target_box(self):
+    batch_size = 1
+    num_masks = 2
+    height = 2
+    width = 2
+    output_size = 2
+
+    masks = tf.ones([batch_size, num_masks, height, width])
+    boxes = tf.constant([[0., 0., 1., 1.], [0., 0., 1., 1.]])
+    target_boxes = tf.constant([[0., 0., 1., 1.], [-1., -1., 1., 1.]])
+    expected_outputs = np.array([[[[1., 1.], [1., 1.]], [[0., 0.], [0., 1.]]]])
+    boxes = tf.reshape(boxes, [batch_size, num_masks, 4])
+    target_boxes = tf.reshape(target_boxes, [batch_size, num_masks, 4])
+
+    tf_cropped_masks = spatial_transform.crop_mask_in_target_box(
+        masks, boxes, target_boxes, output_size)
+    cropped_masks = tf_cropped_masks.numpy()
+    self.assertAllEqual(cropped_masks, expected_outputs)
+
+
+if __name__ == '__main__':
+  logging.set_verbosity(logging.INFO)
+  assert tf.version.VERSION.startswith('2.')
+  tf.test.main()