rename output tensor keys

official/vision/beta/projects/panoptic_maskrcnn/modeling/heads/panoptic_deeplab_heads.py

@@ -381,7 +381,7 @@ class InstanceHead(PanopticDeeplabHead):
     """Creates the variables of the instance head."""
     super(InstanceHead, self).build(input_shape)
     self._instance_center_prediction_conv = tf.keras.layers.Conv2D(
-        name='instance_center_prediction',
+        name='instance_centers_heatmap',
         filters=1,
         kernel_size=self._config_dict['prediction_kernel_size'],
         padding='same',
@@ -391,7 +391,7 @@ class InstanceHead(PanopticDeeplabHead):
         bias_regularizer=self._config_dict['bias_regularizer'])
 
     self._instance_center_regression_conv = tf.keras.layers.Conv2D(
-        name='instance_center_regression',
+        name='instance_centers_offset',
         filters=2,
         kernel_size=self._config_dict['prediction_kernel_size'],
         padding='same',
@@ -409,10 +409,10 @@ class InstanceHead(PanopticDeeplabHead):
       training = tf.keras.backend.learning_phase()
 
     x = super(InstanceHead, self).call(inputs, training=training)
-    instance_center_prediction = self._instance_center_prediction_conv(x)
-    instance_center_regression = self._instance_center_regression_conv(x)
+    instance_centers_heatmap = self._instance_center_prediction_conv(x)
+    instance_centers_offset = self._instance_center_regression_conv(x)
     outputs = {
-        'instance_center_prediction': instance_center_prediction,
-        'instance_center_regression': instance_center_regression
+        'instance_centers_heatmap': instance_centers_heatmap,
+        'instance_centers_offset': instance_centers_offset
     }
     return outputs

official/vision/beta/projects/panoptic_maskrcnn/modeling/heads/panoptic_deeplab_heads_test.py

@@ -71,10 +71,10 @@ class PanopticDeeplabHeadsTest(parameterized.TestCase, tf.test.TestCase):
           semantic_outputs.numpy().shape,
           [2, h, w, num_classes])
       self.assertAllEqual(
-          instance_outputs['instance_center_prediction'].numpy().shape,
+          instance_outputs['instance_centers_heatmap'].numpy().shape,
           [2, h, w, 1])
       self.assertAllEqual(
-          instance_outputs['instance_center_regression'].numpy().shape,
+          instance_outputs['instance_centers_offset'].numpy().shape,
           [2, h, w, 2])
 
 

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

@@ -41,6 +41,11 @@ class PanopticDeeplabNetworkTest(parameterized.TestCase, tf.test.TestCase):
     batch_size = 2 if training else 1
     num_classes = 10
     inputs = np.random.rand(batch_size, input_size, input_size, 3)
+    
+    image_info = tf.convert_to_tensor(
+        [[[input_size, input_size], [input_size, input_size], [1, 1], [0, 0]]])
+    image_info = tf.tile(image_info, [batch_size, 1, 1])
+
     tf.keras.backend.set_image_data_format('channels_last')
     backbone = backbones.ResNet(model_id=50)
 
@@ -65,13 +70,15 @@ class PanopticDeeplabNetworkTest(parameterized.TestCase, tf.test.TestCase):
         low_level_num_filters=(64, 32))
 
     post_processor = panoptic_deeplab_merge.PostProcessor(
+        output_size=[input_size, input_size],
         center_score_threshold=0.1,
         thing_class_ids=[1, 2, 3, 4],
         label_divisor=[256],
         stuff_area_limit=4096,
         ignore_label=0,
         nms_kernel=41,
-        keep_k_centers=41)
+        keep_k_centers=41,
+        rescale_predictions=True)
 
     model = panoptic_deeplab_model.PanopticDeeplabModel(
         backbone=backbone,
@@ -81,12 +88,15 @@ class PanopticDeeplabNetworkTest(parameterized.TestCase, tf.test.TestCase):
         instance_head=instance_head,
         post_processor=post_processor)
 
-    outputs = model(inputs, training=training)
+    outputs = model(
+        inputs=inputs,
+        image_info=image_info,
+        training=training)
 
     if training:
       self.assertIn('segmentation_outputs', outputs)
-      self.assertIn('instance_center_prediction', outputs)
-      self.assertIn('instance_center_regression', outputs)
+      self.assertIn('instance_centers_heatmap', outputs)
+      self.assertIn('instance_centers_offset', outputs)
 
       self.assertAllEqual(
           [2, input_size // (2**low_level[-1]),
@@ -97,12 +107,12 @@ class PanopticDeeplabNetworkTest(parameterized.TestCase, tf.test.TestCase):
           [2, input_size // (2**low_level[-1]),
            input_size // (2**low_level[-1]),
            1],
-          outputs['instance_center_prediction'].numpy().shape)
+          outputs['instance_centers_heatmap'].numpy().shape)
       self.assertAllEqual(
           [2, input_size // (2**low_level[-1]),
            input_size // (2**low_level[-1]),
            2],
-          outputs['instance_center_regression'].numpy().shape)
+          outputs['instance_centers_offset'].numpy().shape)
 
     else:
       self.assertIn('panoptic_outputs', outputs)
@@ -110,6 +120,8 @@ class PanopticDeeplabNetworkTest(parameterized.TestCase, tf.test.TestCase):
       self.assertIn('instance_mask', outputs)
       self.assertIn('instance_centers', outputs)
       self.assertIn('instance_scores', outputs)
+      self.assertIn('segmentation_outputs', outputs)
+
 
   @combinations.generate(
       combinations.combine(
@@ -142,13 +154,15 @@ class PanopticDeeplabNetworkTest(parameterized.TestCase, tf.test.TestCase):
         low_level_num_filters=(64, 32))
 
     post_processor = panoptic_deeplab_merge.PostProcessor(
+        output_size=[640, 640],
         center_score_threshold=0.1,
         thing_class_ids=[1, 2, 3, 4],
         label_divisor=[256],
         stuff_area_limit=4096,
         ignore_label=0,
         nms_kernel=41,
-        keep_k_centers=41)
+        keep_k_centers=41,
+        rescale_predictions=True)
 
     model = panoptic_deeplab_model.PanopticDeeplabModel(
         backbone=backbone,