added `max_num_detections` arg

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

@@ -93,9 +93,11 @@ def build_panoptic_maskrcnn(
       norm_epsilon=norm_activation_config.norm_epsilon,
       kernel_regularizer=l2_regularizer)
 
+  max_num_detections = model_config.detection_generator.max_num_detections
   panoptic_segmentation_generator_obj = \
       panoptic_segmentation_generator.PanopticSegmentationGenerator(
           output_size=postprocessing_config.output_size,
+          max_num_detections=max_num_detections,
           stuff_classes_offset=postprocessing_config.stuff_classes_offset,
           mask_binarize_threshold=postprocessing_config.mask_binarize_threshold,
           score_threshold=postprocessing_config.score_threshold,

official/vision/beta/projects/panoptic_maskrcnn/modeling/layers/panoptic_segmentation_generator.py

@@ -41,6 +41,7 @@ class PanopticSegmentationGenerator(tf.keras.layers.Layer):
   def __init__(
       self,
       output_size: List[int],
+      max_num_detections: int,
       stuff_classes_offset: int,
       mask_binarize_threshold: float = 0.5,
       score_threshold: float = 0.05,
@@ -70,6 +71,7 @@ class PanopticSegmentationGenerator(tf.keras.layers.Layer):
 
     """
     self._output_size = output_size
+    self._max_num_detections = max_num_detections
     self._stuff_classes_offset = stuff_classes_offset
     self._mask_binarize_threshold = mask_binarize_threshold
     self._score_threshold = score_threshold
@@ -79,6 +81,7 @@ class PanopticSegmentationGenerator(tf.keras.layers.Layer):
 
     self._config_dict = {
         'output_size': output_size,
+        'max_num_detections': max_num_detections,
         'stuff_classes_offset': stuff_classes_offset,
         'mask_binarize_threshold': mask_binarize_threshold,
         'score_threshold': score_threshold,
@@ -88,13 +91,8 @@ class PanopticSegmentationGenerator(tf.keras.layers.Layer):
     }
     super(PanopticSegmentationGenerator, self).__init__(**kwargs)
 
-  def _paste_masks(self, boxes, masks, num_detections):
-    pasted_masks = []
-
-    for idx in range(num_detections):
+  def _paste_mask(self, box, mask):
     pasted_mask = tf.ones(self._output_size + [1]) * self._void_class_label
-      mask = masks[idx]
-      box = boxes[idx]
 
     ymin = box[0]
     xmin = box[1]
@@ -116,12 +114,11 @@ class PanopticSegmentationGenerator(tf.keras.layers.Layer):
         end=[ymax, xmax],
         strides=[1, 1],
         value=resized_mask)
-      pasted_masks.append(pasted_mask)
 
-    return tf.stack(pasted_masks, axis=0)
+    return pasted_mask
 
   def _generate_panoptic_masks(self, boxes, scores, classes, detections_masks,
-                               segmentation_mask, num_detections):
+                               segmentation_mask):
     '''Generates panoptic masks for a single image
 
     This function implements the following steps to merge instance and semantic
@@ -145,11 +142,7 @@ class PanopticSegmentationGenerator(tf.keras.layers.Layer):
           for each object.
         segmentation_mask: A `tf.Tensor` of shape [height, width], representing
           the semantic segmentation output.
-        num_detections: A `int` representing the number of valid instance
-          detections.
     '''
-    masks = self._paste_masks(boxes, detections_masks, num_detections)
-
     # Offset stuff class predictions
     segmentation_mask = tf.where(
         tf.logical_or(
@@ -172,7 +165,7 @@ class PanopticSegmentationGenerator(tf.keras.layers.Layer):
     loop_end_idx = tf.where(sorted_scores > self._score_threshold)[-1, 0] + 1
     loop_end_idx = tf.minimum(
         tf.cast(loop_end_idx, dtype=tf.int32),
-        num_detections)
+        self._max_num_detections)
 
     # add things segmentation to panoptic masks
     for i in range(loop_end_idx):
@@ -180,11 +173,15 @@ class PanopticSegmentationGenerator(tf.keras.layers.Layer):
       # the overlaps are resolved based on confidence score
       instance_idx = sorted_indices[i]
 
+      pasted_mask = self._paste_mask(
+          box=boxes[instance_idx],
+          mask=detections_masks[instance_idx])
+
       class_id = tf.cast(classes[instance_idx], dtype=tf.float32)
 
       # convert sigmoid scores to binary values
       binary_mask = tf.greater(
-          masks[instance_idx], self._mask_binarize_threshold)
+          pasted_mask, self._mask_binarize_threshold)
 
       # filter empty instance masks
       if not tf.reduce_sum(tf.cast(binary_mask, tf.float32)) > 0:
@@ -228,7 +225,6 @@ class PanopticSegmentationGenerator(tf.keras.layers.Layer):
   def call(self, inputs):
     detections = inputs
 
-    num_detections = detections['num_detections']
     batched_scores = detections['detection_scores']
     batched_classes = detections['detection_classes']
     batched_boxes = tf.cast(detections['detection_boxes'], dtype=tf.int32)
@@ -257,8 +253,7 @@ class PanopticSegmentationGenerator(tf.keras.layers.Layer):
           scores=batched_scores[idx],
           classes=batched_classes[idx],
           detections_masks=batched_detections_masks[idx],
-          segmentation_mask=batched_segmentation_masks[idx],
-          num_detections=num_detections[idx])
+          segmentation_mask=batched_segmentation_masks[idx])
 
       category_mask.append(results['category_mask'])
       instance_mask.append(results['instance_mask'])

official/vision/beta/projects/panoptic_maskrcnn/modeling/layers/panoptic_segmentation_generator_test.py

@@ -17,6 +17,8 @@
 from absl.testing import parameterized
 import numpy as np
 import tensorflow as tf
+from tensorflow.python.distribute import combinations
+from tensorflow.python.distribute import strategy_combinations
 
 from official.vision.beta.projects.panoptic_maskrcnn.modeling.layers import panoptic_segmentation_generator
 
@@ -29,6 +31,7 @@ class PanopticSegmentationGeneratorTest(
   def test_serialize_deserialize(self):
     config = {
         'output_size': [640, 640],
+        'max_num_detections': 100,
         'stuff_classes_offset': 90,
         'mask_binarize_threshold': 0.5,
         'score_threshold': 0.005,
@@ -46,8 +49,13 @@ class PanopticSegmentationGeneratorTest(
 
     self.assertAllEqual(generator.get_config(), new_generator.get_config())
 
-
-  def test_outputs(self):
+  @combinations.generate(
+      combinations.combine(
+          strategy=[
+              strategy_combinations.cloud_tpu_strategy,
+              strategy_combinations.one_device_strategy_gpu,
+          ]))
+  def test_outputs(self, strategy):
 
     # 0 represents the void class label
     thing_class_ids = [0, 1, 2, 3, 4]
@@ -65,6 +73,7 @@ class PanopticSegmentationGeneratorTest(
 
     config = {
         'output_size': [640, 640],
+        'max_num_detections': 100,
         'stuff_classes_offset': 3,
         'mask_binarize_threshold': 0.5,
         'score_threshold': 0.005,
@@ -106,18 +115,26 @@ class PanopticSegmentationGeneratorTest(
         'segmentation_outputs': segmentation_mask_one_hot
         }
 
-    outputs = generator(inputs=inputs)
+    def _run(inputs):
+      return generator(inputs=inputs)
+
+    @tf.function
+    def _distributed_run(inputs):
+      outputs = strategy.run(_run, args=((inputs,)))
+      return strategy.gather(outputs, axis=0)
+
+    outputs = _distributed_run(inputs)
 
     self.assertIn('category_mask', outputs)
     self.assertIn('instance_mask', outputs)
 
     self.assertAllEqual(
-        outputs['category_mask'].get_shape().as_list(),
-        [1] + config['output_size'])
+        outputs['category_mask'][0].get_shape().as_list(),
+        config['output_size'])
 
     self.assertAllEqual(
-        outputs['instance_mask'].get_shape().as_list(),
-        [1] + config['output_size'])
+        outputs['instance_mask'][0].get_shape().as_list(),
+        config['output_size'])
 
     for category_id in np.unique(outputs['category_mask']):
       self.assertIn(category_id, all_class_ids)

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

@@ -100,7 +100,9 @@ class PanopticMaskRCNNModelTest(parameterized.TestCase, tf.test.TestCase):
     detection_generator_obj = detection_generator.DetectionGenerator()
     panoptic_segmentation_generator_obj = \
         panoptic_segmentation_generator.PanopticSegmentationGenerator(
-            output_size=[image_size, image_size], stuff_classes_offset=90)
+            output_size=[image_size, image_size],
+            max_num_detections=100,
+            stuff_classes_offset=90)
     mask_head = instance_heads.MaskHead(
         num_classes=num_classes, upsample_factor=2)
     mask_sampler_obj = mask_sampler.MaskSampler(
@@ -228,7 +230,9 @@ class PanopticMaskRCNNModelTest(parameterized.TestCase, tf.test.TestCase):
       detection_generator_obj = detection_generator.DetectionGenerator()
       panoptic_segmentation_generator_obj = \
           panoptic_segmentation_generator.PanopticSegmentationGenerator(
-              output_size=list(image_size), stuff_classes_offset=90)
+              output_size=list(image_size),
+              max_num_detections=100,
+              stuff_classes_offset=90)
       mask_head = instance_heads.MaskHead(
           num_classes=num_classes, upsample_factor=2)
       mask_sampler_obj = mask_sampler.MaskSampler(
@@ -338,7 +342,9 @@ class PanopticMaskRCNNModelTest(parameterized.TestCase, tf.test.TestCase):
     detection_generator_obj = detection_generator.DetectionGenerator()
     panoptic_segmentation_generator_obj = \
         panoptic_segmentation_generator.PanopticSegmentationGenerator(
-            output_size=[None, None], stuff_classes_offset=90)
+            output_size=[None, None],
+            max_num_detections=100,
+            stuff_classes_offset=90)
     segmentation_resnet_model_id = 101
     segmentation_output_stride = 16
     aspp_dilation_rates = [6, 12, 18]
@@ -416,7 +422,9 @@ class PanopticMaskRCNNModelTest(parameterized.TestCase, tf.test.TestCase):
     detection_generator_obj = detection_generator.DetectionGenerator()
     panoptic_segmentation_generator_obj = \
         panoptic_segmentation_generator.PanopticSegmentationGenerator(
-            output_size=[None, None], stuff_classes_offset=90)
+            output_size=[None, None],
+            max_num_detections=100,
+            stuff_classes_offset=90)
     segmentation_resnet_model_id = 101
     segmentation_output_stride = 16
     aspp_dilation_rates = [6, 12, 18]