added tests for `PanopticSegmentationGenerator`

official/vision/beta/projects/panoptic_maskrcnn/modeling/layers/panoptic_segmentation_generator_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 panoptic_segmentation_generator.py."""
+
+from absl.testing import parameterized
+import numpy as np
+import tensorflow as tf
+
+from official.vision.beta.projects.panoptic_maskrcnn.modeling.layers import panoptic_segmentation_generator
+
+PANOPTIC_SEGMENTATION_GENERATOR = \
+    panoptic_segmentation_generator.PanopticSegmentationGenerator
+
+class PanopticSegmentationGeneratorTest(
+    parameterized.TestCase, tf.test.TestCase):
+
+  def test_serialize_deserialize(self):
+    config = {
+        'output_size': [640, 640],
+        'stuff_classes_offset': 90,
+        'mask_binarize_threshold': 0.5,
+        'score_threshold': 0.005,
+        'things_class_label': 1,
+        'void_class_label': 0,
+        'void_instance_id': -1
+    }
+    generator = PANOPTIC_SEGMENTATION_GENERATOR(**config)
+
+    expected_config = dict(config)
+    self.assertEqual(generator.get_config(), expected_config)
+
+    new_generator = PANOPTIC_SEGMENTATION_GENERATOR.from_config(
+        generator.get_config())
+
+    self.assertAllEqual(generator.get_config(), new_generator.get_config())
+
+
+  def test_outputs(self):
+
+    # 0 represents the void class label
+    thing_class_ids = [0, 1, 2, 3, 4]
+    stuff_class_ids = [0, 5, 6, 7, 8, 9, 10]
+    all_class_ids = set(thing_class_ids + stuff_class_ids)
+
+    num_thing_classes = len(thing_class_ids)
+    num_stuff_classes = len(stuff_class_ids)
+    num_classes_for_segmentation = num_stuff_classes + 1
+
+    # all thing classes are mapped to class_id=1, stuff class ids are offset
+    # such that the stuff class_ids start from 2, this means the semantic
+    # segmentation head will have ground truths with class_ids belonging to 
+    # [0, 1, 2, 3, 4, 5, 6, 7]
+
+    config = {
+        'output_size': [640, 640],
+        'stuff_classes_offset': 3,
+        'mask_binarize_threshold': 0.5,
+        'score_threshold': 0.005,
+        'things_class_label': 1,
+        'void_class_label': 0,
+        'void_instance_id': -1
+    }
+    generator = PANOPTIC_SEGMENTATION_GENERATOR(**config)
+
+    crop_height = 112
+    crop_width = 112
+
+    boxes = tf.constant([[
+        [167, 398, 342, 619],
+        [192, 171, 363, 449],
+        [211, 1, 382, 74]
+    ]])
+
+    num_detections = boxes.get_shape().as_list()[1]
+    scores = tf.random.uniform([1, num_detections], 0, 1)
+    classes = tf.random.uniform(
+        [1, num_detections],
+        1, num_thing_classes, dtype=tf.int32)
+    masks = tf.random.normal(
+        [1, num_detections, crop_height, crop_width])
+
+    segmentation_mask = tf.random.uniform(
+        [1, *config['output_size']],
+        0, num_classes_for_segmentation, dtype=tf.int32)
+    segmentation_mask_one_hot = tf.one_hot(
+        segmentation_mask, depth=num_stuff_classes + 1)
+
+    inputs = {
+        'detection_boxes': boxes,
+        'detection_scores': scores,
+        'detection_classes': classes,
+        'detection_masks': masks,
+        'num_detections': tf.constant([num_detections]),
+        'segmentation_outputs': segmentation_mask_one_hot
+        }
+
+    outputs = generator(inputs=inputs)
+
+    self.assertIn('category_mask', outputs)
+    self.assertIn('instance_mask', outputs)
+
+    self.assertAllEqual(
+        outputs['category_mask'].get_shape().as_list(),
+        [1] + config['output_size'])
+
+    self.assertAllEqual(
+        outputs['instance_mask'].get_shape().as_list(),
+        [1] + config['output_size'])
+
+    for category_id in np.unique(outputs['category_mask']):
+      self.assertIn(category_id, all_class_ids)
+
+if __name__ == '__main__':
+  tf.test.main()

official/vision/beta/projects/panoptic_maskrcnn/modeling/panoptic_maskrcnn_model_test.py

@@ -15,7 +15,6 @@
 """Tests for panoptic_maskrcnn_model.py."""
 
 import os
-
 from absl.testing import parameterized
 import numpy as np
 import tensorflow as tf
@@ -35,7 +34,7 @@ from official.vision.beta.modeling.layers import roi_generator
 from official.vision.beta.modeling.layers import roi_sampler
 from official.vision.beta.ops import anchor
 from official.vision.beta.projects.panoptic_maskrcnn.modeling import panoptic_maskrcnn_model
-
+from official.vision.beta.projects.panoptic_maskrcnn.modeling.layers import panoptic_segmentation_generator
 
 class PanopticMaskRCNNModelTest(parameterized.TestCase, tf.test.TestCase):
 
@@ -99,6 +98,9 @@ class PanopticMaskRCNNModelTest(parameterized.TestCase, tf.test.TestCase):
     roi_sampler_obj = roi_sampler.ROISampler()
     roi_aligner_obj = roi_aligner.MultilevelROIAligner()
     detection_generator_obj = detection_generator.DetectionGenerator()
+    panoptic_segmentation_generator_obj = \
+        panoptic_segmentation_generator.PanopticSegmentationGenerator(
+            output_size=[image_size, image_size], stuff_classes_offset=90)
     mask_head = instance_heads.MaskHead(
         num_classes=num_classes, upsample_factor=2)
     mask_sampler_obj = mask_sampler.MaskSampler(
@@ -131,6 +133,7 @@ class PanopticMaskRCNNModelTest(parameterized.TestCase, tf.test.TestCase):
         roi_sampler_obj,
         roi_aligner_obj,
         detection_generator_obj,
+        panoptic_segmentation_generator_obj,
         mask_head,
         mask_sampler_obj,
         mask_roi_aligner_obj,
@@ -223,6 +226,9 @@ class PanopticMaskRCNNModelTest(parameterized.TestCase, tf.test.TestCase):
       roi_sampler_cascade.append(roi_sampler_obj)
       roi_aligner_obj = roi_aligner.MultilevelROIAligner()
       detection_generator_obj = detection_generator.DetectionGenerator()
+      panoptic_segmentation_generator_obj = \
+          panoptic_segmentation_generator.PanopticSegmentationGenerator(
+              output_size=list(image_size), stuff_classes_offset=90)
       mask_head = instance_heads.MaskHead(
           num_classes=num_classes, upsample_factor=2)
       mask_sampler_obj = mask_sampler.MaskSampler(
@@ -255,6 +261,7 @@ class PanopticMaskRCNNModelTest(parameterized.TestCase, tf.test.TestCase):
           roi_sampler_obj,
           roi_aligner_obj,
           detection_generator_obj,
+          panoptic_segmentation_generator_obj,
           mask_head,
           mask_sampler_obj,
           mask_roi_aligner_obj,
@@ -300,9 +307,19 @@ class PanopticMaskRCNNModelTest(parameterized.TestCase, tf.test.TestCase):
       self.assertIn('num_detections', results)
       self.assertIn('detection_masks', results)
       self.assertIn('segmentation_outputs', results)
+      self.assertIn('panoptic_outputs', results)
+      self.assertIn('category_mask', results['panoptic_outputs'])
+      self.assertIn('instance_mask', results['panoptic_outputs'])
+
       self.assertAllEqual(
           [2, image_size[0] // (2**level), image_size[1] // (2**level), 2],
           results['segmentation_outputs'].numpy().shape)
+      self.assertAllEqual(
+          [2, image_size[0], image_size[1]],
+          results['panoptic_outputs']['category_mask'].numpy().shape)
+      self.assertAllEqual(
+          [2, image_size[0], image_size[1]],
+          results['panoptic_outputs']['instance_mask'].numpy().shape)
 
   @combinations.generate(
       combinations.combine(
@@ -319,6 +336,9 @@ class PanopticMaskRCNNModelTest(parameterized.TestCase, tf.test.TestCase):
     roi_sampler_obj = roi_sampler.ROISampler()
     roi_aligner_obj = roi_aligner.MultilevelROIAligner()
     detection_generator_obj = detection_generator.DetectionGenerator()
+    panoptic_segmentation_generator_obj = \
+        panoptic_segmentation_generator.PanopticSegmentationGenerator(
+            output_size=[None, None], stuff_classes_offset=90)
     segmentation_resnet_model_id = 101
     segmentation_output_stride = 16
     aspp_dilation_rates = [6, 12, 18]
@@ -356,6 +376,7 @@ class PanopticMaskRCNNModelTest(parameterized.TestCase, tf.test.TestCase):
         roi_sampler_obj,
         roi_aligner_obj,
         detection_generator_obj,
+        panoptic_segmentation_generator_obj,
         mask_head,
         mask_sampler_obj,
         mask_roi_aligner_obj,
@@ -393,6 +414,9 @@ class PanopticMaskRCNNModelTest(parameterized.TestCase, tf.test.TestCase):
     roi_sampler_obj = roi_sampler.ROISampler()
     roi_aligner_obj = roi_aligner.MultilevelROIAligner()
     detection_generator_obj = detection_generator.DetectionGenerator()
+    panoptic_segmentation_generator_obj = \
+        panoptic_segmentation_generator.PanopticSegmentationGenerator(
+            output_size=[None, None], stuff_classes_offset=90)
     segmentation_resnet_model_id = 101
     segmentation_output_stride = 16
     aspp_dilation_rates = [6, 12, 18]
@@ -430,6 +454,7 @@ class PanopticMaskRCNNModelTest(parameterized.TestCase, tf.test.TestCase):
         roi_sampler_obj,
         roi_aligner_obj,
         detection_generator_obj,
+        panoptic_segmentation_generator_obj,
         mask_head,
         mask_sampler_obj,
         mask_roi_aligner_obj,