Internal change

PiperOrigin-RevId: 381516130

Author: fyangf

official/vision/beta/projects/volumetric_models/README.md

@@ -0,0 +1,31 @@
+# Volumetric Models
+
+**DISCLAIMER**: This implementation is still under development. No support will
+be provided during the development phase.
+
+This folder contains implementation of volumetric models, i.e., UNet 3D model,
+for 3D semantic segmentation.
+
+## Modeling
+
+Following the style of TF-Vision, a UNet 3D model is implemented as a backbone
+and a decoder.
+
+## Backbone
+
+The backbone is the left U-shape of the complete UNet model. It takes batch of
+images as input, and outputs a dictionary in a form of `{level: features}`.
+`features` in the output is a tensor of feature maps.
+
+## Decoder
+
+The decoder is the right U-shape of the complete UNet model. It takes the output
+dictionary from the backbone and connects the feature maps from each level to
+the decoder's decoding branches. The final output is the raw segmentation
+predictions.
+
+An additional head is attached to the output of the decoder to optionally
+perform more operations and then generate the prediction map of logits.
+
+The `factory.py` file builds and connects the backbone, decoder and head
+together to form the complete UNet model.

official/vision/beta/projects/volumetric_models/configs/backbones.py

@@ -0,0 +1,48 @@
+# 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.
+
+# Lint as: python3
+"""Backbones configurations."""
+from typing import Optional, Sequence
+
+import dataclasses
+
+from official.modeling import hyperparams
+
+
+@dataclasses.dataclass
+class UNet3D(hyperparams.Config):
+  """UNet3D config."""
+  model_id: int = 4
+  pool_size: Sequence[int] = (2, 2, 2)
+  kernel_size: Sequence[int] = (3, 3, 3)
+  base_filters: int = 32
+  use_batch_normalization: bool = True
+
+
+@dataclasses.dataclass
+class Backbone(hyperparams.OneOfConfig):
+  """Configuration for backbones.
+
+  Attributes:
+    type: 'str', type of backbone be used, one the of fields below.
+    resnet: resnet backbone config.
+    dilated_resnet: dilated resnet backbone for semantic segmentation config.
+    revnet: revnet backbone config.
+    efficientnet: efficientnet backbone config.
+    spinenet: spinenet backbone config.
+    mobilenet: mobilenet backbone config.
+  """
+  type: Optional[str] = None
+  unet_3d: UNet3D = UNet3D()

official/vision/beta/projects/volumetric_models/configs/decoders.py

@@ -0,0 +1,43 @@
+# 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.
+
+# Lint as: python3
+"""Decoders configurations."""
+from typing import Optional, Sequence
+
+import dataclasses
+
+from official.modeling import hyperparams
+
+
+@dataclasses.dataclass
+class UNet3DDecoder(hyperparams.Config):
+  """UNet3D decoder config."""
+  model_id: int = 4
+  pool_size: Sequence[int] = (2, 2, 2)
+  kernel_size: Sequence[int] = (3, 3, 3)
+  use_batch_normalization: bool = True
+  use_deconvolution: bool = True
+
+
+@dataclasses.dataclass
+class Decoder(hyperparams.OneOfConfig):
+  """Configuration for decoders.
+
+  Attributes:
+    type: 'str', type of decoder be used, on the of fields below.
+    fpn: fpn config.
+  """
+  type: Optional[str] = None
+  unet_3d_decoder: UNet3DDecoder = UNet3DDecoder()

official/vision/beta/projects/volumetric_models/configs/semantic_segmentation_3d.py

@@ -0,0 +1,163 @@
+# 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.
+
+# Lint as: python3
+"""Semantic segmentation configuration definition."""
+from typing import List, Optional, Union
+
+import dataclasses
+
+from official.core import exp_factory
+from official.modeling import hyperparams
+from official.modeling import optimization
+from official.modeling.hyperparams import config_definitions as cfg
+from official.vision.beta.configs import common
+from official.vision.beta.projects.volumetric_models.configs import backbones
+from official.vision.beta.projects.volumetric_models.configs import decoders
+
+
+@dataclasses.dataclass
+class DataConfig(cfg.DataConfig):
+  """Input config for training."""
+  output_size: List[int] = dataclasses.field(default_factory=list)
+  input_size: List[int] = dataclasses.field(default_factory=list)
+  num_classes: int = 0
+  num_channels: int = 1
+  input_path: str = ''
+  global_batch_size: int = 0
+  is_training: bool = True
+  dtype: str = 'float32'
+  label_dtype: str = 'float32'
+  image_field_key: str = 'image/encoded'
+  label_field_key: str = 'image/class/label'
+  shuffle_buffer_size: int = 1000
+  cycle_length: int = 10
+  drop_remainder: bool = False
+  file_type: str = 'tfrecord'
+
+
+@dataclasses.dataclass
+class SegmentationHead3D(hyperparams.Config):
+  """Segmentation head config."""
+  num_classes: int = 0
+  level: int = 1
+  num_convs: int = 0
+  num_filters: int = 256
+  upsample_factor: int = 1
+  output_logits: bool = True
+
+
+@dataclasses.dataclass
+class SemanticSegmentationModel3D(hyperparams.Config):
+  """Semantic segmentation model config."""
+  num_classes: int = 0
+  num_channels: int = 1
+  input_size: List[int] = dataclasses.field(default_factory=list)
+  min_level: int = 3
+  max_level: int = 6
+  head: SegmentationHead3D = SegmentationHead3D()
+  backbone: backbones.Backbone = backbones.Backbone(
+      type='unet_3d', unet_3d=backbones.UNet3D())
+  decoder: decoders.Decoder = decoders.Decoder(
+      type='unet_3d_decoder', unet_3d_decoder=decoders.UNet3DDecoder())
+  norm_activation: common.NormActivation = common.NormActivation()
+
+
+@dataclasses.dataclass
+class Losses(hyperparams.Config):
+  # Supported `loss_type` are `adaptive` and `generalized`.
+  loss_type: str = 'adaptive'
+  l2_weight_decay: float = 0.0
+
+
+@dataclasses.dataclass
+class Evaluation(hyperparams.Config):
+  report_per_class_metric: bool = False  # Whether to report per-class metrics.
+
+
+@dataclasses.dataclass
+class SemanticSegmentation3DTask(cfg.TaskConfig):
+  """The model config."""
+  model: SemanticSegmentationModel3D = SemanticSegmentationModel3D()
+  train_data: DataConfig = DataConfig(is_training=True)
+  validation_data: DataConfig = DataConfig(is_training=False)
+  losses: Losses = Losses()
+  evaluation: Evaluation = Evaluation()
+  train_input_partition_dims: List[int] = dataclasses.field(
+      default_factory=list)
+  eval_input_partition_dims: List[int] = dataclasses.field(default_factory=list)
+  init_checkpoint: Optional[str] = None
+  init_checkpoint_modules: Union[
+      str, List[str]] = 'all'  # all, backbone, and/or decoder
+
+
+@exp_factory.register_config_factory('seg_unet3d_test')
+def seg_unet3d_test() -> cfg.ExperimentConfig:
+  """Image segmentation on a dummy dataset with 3D UNet for testing purpose."""
+  train_batch_size = 2
+  eval_batch_size = 2
+  steps_per_epoch = 10
+  config = cfg.ExperimentConfig(
+      task=SemanticSegmentation3DTask(
+          model=SemanticSegmentationModel3D(
+              num_classes=2,
+              input_size=[32, 32, 32],
+              num_channels=2,
+              backbone=backbones.Backbone(
+                  type='unet_3d', unet_3d=backbones.UNet3D(model_id=2)),
+              decoder=decoders.Decoder(
+                  type='unet_3d_decoder',
+                  unet_3d_decoder=decoders.UNet3DDecoder(model_id=2)),
+              head=SegmentationHead3D(num_convs=0, num_classes=2),
+              norm_activation=common.NormActivation(
+                  activation='relu', use_sync_bn=False)),
+          train_data=DataConfig(
+              input_path='train.tfrecord',
+              num_classes=2,
+              input_size=[32, 32, 32],
+              num_channels=2,
+              is_training=True,
+              global_batch_size=train_batch_size),
+          validation_data=DataConfig(
+              input_path='val.tfrecord',
+              num_classes=2,
+              input_size=[32, 32, 32],
+              num_channels=2,
+              is_training=False,
+              global_batch_size=eval_batch_size),
+          losses=Losses(loss_type='adaptive')),
+      trainer=cfg.TrainerConfig(
+          steps_per_loop=steps_per_epoch,
+          summary_interval=steps_per_epoch,
+          checkpoint_interval=steps_per_epoch,
+          train_steps=10,
+          validation_steps=10,
+          validation_interval=steps_per_epoch,
+          optimizer_config=optimization.OptimizationConfig({
+              'optimizer': {
+                  'type': 'sgd',
+              },
+              'learning_rate': {
+                  'type': 'constant',
+                  'constant': {
+                      'learning_rate': 0.000001
+                  }
+              }
+          })),
+      restrictions=[
+          'task.train_data.is_training != None',
+          'task.validation_data.is_training != None'
+      ])
+
+  return config

official/vision/beta/projects/volumetric_models/configs/semantic_segmentation_3d_test.py

@@ -0,0 +1,44 @@
+# 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.
+
+# Lint as: python3
+"""Tests for semantic_segmentation."""
+
+# pylint: disable=unused-import
+from absl.testing import parameterized
+import tensorflow as tf
+
+from official.core import exp_factory
+from official.modeling.hyperparams import config_definitions as cfg
+from official.vision.beta.projects.volumetric_models.configs import semantic_segmentation_3d as exp_cfg
+
+
+class ImageSegmentationConfigTest(tf.test.TestCase, parameterized.TestCase):
+
+  @parameterized.parameters(
+      ('seg_unet3d_test',),)
+  def test_semantic_segmentation_configs(self, config_name):
+    config = exp_factory.get_exp_config(config_name)
+    self.assertIsInstance(config, cfg.ExperimentConfig)
+    self.assertIsInstance(config.task, exp_cfg.SemanticSegmentation3DTask)
+    self.assertIsInstance(config.task.model,
+                          exp_cfg.SemanticSegmentationModel3D)
+    self.assertIsInstance(config.task.train_data, exp_cfg.DataConfig)
+    config.task.train_data.is_training = None
+    with self.assertRaises(KeyError):
+      config.validate()
+
+
+if __name__ == '__main__':
+  tf.test.main()