panoptic deeplab modeling

official/vision/beta/projects/panoptic_maskrcnn/configs/panoptic_deeplab.py

@@ -0,0 +1,78 @@
+# 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.
+
+"""Panoptic Deeplab configuration definition."""
+
+import dataclasses
+from typing import List, Tuple, Union
+
+from official.modeling import hyperparams
+from official.vision.beta.configs import common
+from official.vision.beta.configs import backbones
+from official.vision.beta.configs import decoders
+
+_COCO_INPUT_PATH_BASE = 'coco/tfrecords'
+_COCO_TRAIN_EXAMPLES = 118287
+_COCO_VAL_EXAMPLES = 5000
+
+
+@dataclasses.dataclass
+class PanopticDeeplabHead(hyperparams.Config):
+  """Panoptic Deeplab head config."""
+  level: int = 3
+  num_convs: int = 2
+  num_filters: int = 256
+  kernel_size: int = 5
+  use_depthwise_convolution: bool = False
+  upsample_factor: int = 1
+  low_level: Union[List[int], Tuple[int]] = (3, 2)
+  low_level_num_filters: Union[List[int], Tuple[int]] = (64, 32)
+
+
+@dataclasses.dataclass
+class SemanticHead(PanopticDeeplabHead):
+  """Semantic head config."""
+  prediction_kernel_size: int = 1
+
+@dataclasses.dataclass
+class InstanceHead(PanopticDeeplabHead):
+  """Instance head config."""
+  prediction_kernel_size: int = 1
+
+@dataclasses.dataclass
+class PanopticDeeplabPostProcessor(hyperparams.Config):
+  """Panoptic Deeplab PostProcessing config."""
+  center_score_threshold: float = 0.1
+  thing_class_ids: List[int] = dataclasses.field(default_factory=list)
+  label_divisor: int = 256 * 256 * 256
+  stuff_area_limit: int = 4096
+  ignore_label: int = 0
+  nms_kernel: int = 41
+  keep_k_centers: int = 400
+
+@dataclasses.dataclass
+class PanopticDeeplab(hyperparams.Config):
+  """Panoptic Deeplab model config."""
+  num_classes: int = 0
+  input_size: List[int] = dataclasses.field(default_factory=list)
+  min_level: int = 3
+  max_level: int = 6
+  norm_activation: common.NormActivation = common.NormActivation()
+  backbone: backbones.Backbone = backbones.Backbone(
+      type='resnet', resnet=backbones.ResNet())
+  decoder: decoders.Decoder = decoders.Decoder(type='aspp')
+  semantic_head: SemanticHead = SemanticHead()
+  instance_head: InstanceHead = InstanceHead()
+  shared_decoder: bool = False
+  post_processor: PanopticDeeplabPostProcessor = PanopticDeeplabPostProcessor()

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

@@ -16,10 +16,15 @@
 
 import tensorflow as tf
 
+
+from official.vision.beta.projects.panoptic_maskrcnn.configs import panoptic_deeplab as panoptic_deeplab_cfg
 from official.vision.beta.projects.deepmac_maskrcnn.tasks import deep_mask_head_rcnn
 from official.vision.beta.projects.panoptic_maskrcnn.configs import panoptic_maskrcnn as panoptic_maskrcnn_cfg
+from official.vision.beta.projects.panoptic_maskrcnn.modeling import panoptic_deeplab_model
+from official.vision.beta.projects.panoptic_maskrcnn.modeling.heads import panoptic_deeplab_heads
 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
+from official.vision.beta.projects.panoptic_maskrcnn.modeling.layers import panoptic_deeplab_merge
 from official.vision.modeling import backbones
 from official.vision.modeling.decoders import factory as decoder_factory
 from official.vision.modeling.heads import segmentation_heads
@@ -142,3 +147,97 @@ def build_panoptic_maskrcnn(
       aspect_ratios=model_config.anchor.aspect_ratios,
       anchor_size=model_config.anchor.anchor_size)
   return model
+
+
+def build_panoptic_deeplab(
+        input_specs: tf.keras.layers.InputSpec,
+        model_config: panoptic_deeplab_cfg.PanopticDeeplab,
+        l2_regularizer: tf.keras.regularizers.Regularizer = None) -> tf.keras.Model:  # pytype: disable=annotation-type-mismatch  # typed-keras
+  """Builds Panoptic Deeplab model.
+
+
+  Args:
+    input_specs: `tf.keras.layers.InputSpec` specs of the input tensor.
+    model_config: Config instance for the panoptic maskrcnn model.
+    l2_regularizer: Optional `tf.keras.regularizers.Regularizer`, if specified,
+      the model is built with the provided regularization layer.
+  Returns:
+    tf.keras.Model for the panoptic segmentation model.
+  """
+  norm_activation_config = model_config.norm_activation
+  backbone = backbones.factory.build_backbone(
+      input_specs=input_specs,
+      backbone_config=model_config.backbone,
+      norm_activation_config=norm_activation_config,
+      l2_regularizer=l2_regularizer)
+
+  semantic_decoder = decoder_factory.build_decoder(
+      input_specs=backbone.output_specs,
+      model_config=model_config,
+      l2_regularizer=l2_regularizer)
+
+  if model_config.shared_decoder:
+    instance_decoder = None
+  else:
+    # TODO(srihari-humbarwadi): decouple semantic and 
+    # instance decoder types
+    instance_decoder = decoder_factory.build_decoder(
+        input_specs=backbone.output_specs,
+        model_config=model_config,
+        l2_regularizer=l2_regularizer)
+
+  semantic_head_config = model_config.semantic_head
+  instance_head_config = model_config.instance_head
+
+  semantic_head = panoptic_deeplab_heads.SemanticHead(
+      num_classes=model_config.num_classes,
+      level=semantic_head_config.level,
+      num_convs=semantic_head_config.num_convs,
+      kernel_size=semantic_head_config.kernel_size,
+      prediction_kernel_size=semantic_head_config.prediction_kernel_size,
+      num_filters=semantic_head_config.num_filters,
+      use_depthwise_convolution=semantic_head_config.use_depthwise_convolution,
+      upsample_factor=semantic_head_config.upsample_factor,
+      low_level=semantic_head_config.low_level,
+      low_level_num_filters=semantic_head_config.low_level_num_filters,
+      activation=norm_activation_config.activation,
+      use_sync_bn=norm_activation_config.use_sync_bn,
+      norm_momentum=norm_activation_config.norm_momentum,
+      norm_epsilon=norm_activation_config.norm_epsilon,
+      kernel_regularizer=l2_regularizer)
+
+  instance_head = panoptic_deeplab_heads.InstanceHead(
+      level=instance_head_config.level,
+      num_convs=instance_head_config.num_convs,
+      kernel_size=instance_head_config.kernel_size,
+      prediction_kernel_size=instance_head_config.prediction_kernel_size,
+      num_filters=instance_head_config.num_filters,
+      use_depthwise_convolution=instance_head_config.use_depthwise_convolution,
+      upsample_factor=instance_head_config.upsample_factor,
+      low_level=instance_head_config.low_level,
+      low_level_num_filters=instance_head_config.low_level_num_filters,
+      activation=norm_activation_config.activation,
+      use_sync_bn=norm_activation_config.use_sync_bn,
+      norm_momentum=norm_activation_config.norm_momentum,
+      norm_epsilon=norm_activation_config.norm_epsilon,
+      kernel_regularizer=l2_regularizer)
+
+  post_processing_config = model_config.post_processor
+  post_processor = panoptic_deeplab_merge.PostProcessor(
+      center_score_threshold=post_processing_config.center_score_threshold,
+      thing_class_ids=post_processing_config.thing_class_ids,
+      label_divisor=post_processing_config.label_divisor,
+      stuff_area_limit=post_processing_config.stuff_area_limit,
+      ignore_label=post_processing_config.ignore_label,
+      nms_kernel=post_processing_config.nms_kernel,
+      keep_k_centers=post_processing_config.keep_k_centers)
+
+  model = panoptic_deeplab_model.PanopticDeeplabModel(
+      backbone=backbone, 
+      semantic_decoder=semantic_decoder,
+      instance_decoder=instance_decoder,
+      semantic_head=semantic_head,
+      instance_head=instance_head, 
+      post_processor=post_processor)
+
+  return model

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

@@ -17,8 +17,10 @@
 from absl.testing import parameterized
 import numpy as np
 import tensorflow as tf
+from tensorflow.python.distribute import combinations
 
 from official.vision.beta.projects.panoptic_maskrcnn.configs import panoptic_maskrcnn as panoptic_maskrcnn_cfg
+from official.vision.beta.projects.panoptic_maskrcnn.configs import panoptic_deeplab as panoptic_deeplab_cfg
 from official.vision.beta.projects.panoptic_maskrcnn.modeling import factory
 from official.vision.configs import backbones
 from official.vision.configs import decoders
@@ -62,5 +64,47 @@ def test_builder(self, backbone_type, input_size, segmentation_backbone_type,
         model_config=model_config,
         l2_regularizer=l2_regularizer)
 
+class PanopticDeeplabBuilderTest(parameterized.TestCase, tf.test.TestCase):
+
+  @combinations.generate(
+      combinations.combine(
+          input_size=[(640, 640), (512, 512)],
+          backbone_type=['resnet', 'dilated_resnet'],
+          decoder_type=['aspp', 'fpn'],
+          level=[2, 3, 4],
+          low_level=[(4, 3), (3, 2)],
+          shared_decoder=[True, False]))
+  def test_builder(self, input_size, backbone_type, level, 
+                   low_level, decoder_type, shared_decoder):
+    num_classes = 10
+    input_specs = tf.keras.layers.InputSpec(
+        shape=[None, input_size[0], input_size[1], 3])
+
+    model_config = panoptic_deeplab_cfg.PanopticDeeplab(
+        num_classes=num_classes,
+        input_size=input_size,
+        backbone=backbones.Backbone(type=backbone_type),
+        decoder=decoders.Decoder(type=decoder_type),
+        semantic_head=panoptic_deeplab_cfg.SemanticHead(
+            level=level, 
+            num_convs=1, 
+            kernel_size=5, 
+            prediction_kernel_size=1, 
+            low_level=low_level),
+        instance_head=panoptic_deeplab_cfg.InstanceHead(
+            level=level,
+            num_convs=1,
+            kernel_size=5,
+            prediction_kernel_size=1,
+            low_level=low_level),
+        shared_decoder=shared_decoder)
+
+    l2_regularizer = tf.keras.regularizers.l2(5e-5)
+    _ = factory.build_panoptic_deeplab(
+        input_specs=input_specs,
+        model_config=model_config,
+        l2_regularizer=l2_regularizer)
+
+
 if __name__ == '__main__':
   tf.test.main()