Internal change

PiperOrigin-RevId: 413981781

official/legacy/detection/README.md

+# Object Detection Models on TensorFlow 2
+
+**WARNING**: This repository will be deprecated and replaced by the solid
+implementations inside vision/beta/.
+
+## Prerequsite
+To get started, download the code from TensorFlow models GitHub repository or
+use the pre-installed Google Cloud VM.
+
+```bash
+git clone https://github.com/tensorflow/models.git
+```
+
+Next, make sure to use TensorFlow 2.1+ on Google Cloud. Also here are
+a few package you need to install to get started:
+
+```bash
+sudo apt-get install -y python-tk && \
+pip3 install -r ~/models/official/requirements.txt
+```
+
+## Train RetinaNet on TPU
+
+### Train a vanilla ResNet-50 based RetinaNet.
+
+```bash
+TPU_NAME="<your GCP TPU name>"
+MODEL_DIR="<path to the directory to store model files>"
+RESNET_CHECKPOINT="<path to the pre-trained Resnet-50 checkpoint>"
+TRAIN_FILE_PATTERN="<path to the TFRecord training data>"
+EVAL_FILE_PATTERN="<path to the TFRecord validation data>"
+VAL_JSON_FILE="<path to the validation annotation JSON file>"
+python3 ~/models/official/vision/detection/main.py \
+  --strategy_type=tpu \
+  --tpu="${TPU_NAME?}" \
+  --model_dir="${MODEL_DIR?}" \
+  --mode=train \
+  --params_override="{ type: retinanet, train: { checkpoint: { path: ${RESNET_CHECKPOINT?}, prefix: resnet50/ }, train_file_pattern: ${TRAIN_FILE_PATTERN?} }, eval: { val_json_file: ${VAL_JSON_FILE?}, eval_file_pattern: ${EVAL_FILE_PATTERN?} } }"
+```
+
+The pre-trained ResNet-50 checkpoint can be downloaded [here](https://storage.cloud.google.com/cloud-tpu-checkpoints/model-garden-vision/detection/resnet50-2018-02-07.tar.gz).
+
+Note: The ResNet implementation under
+[detection/](https://github.com/tensorflow/models/tree/master/official/vision/detection)
+is currently different from the one under
+[classification/](https://github.com/tensorflow/models/tree/master/official/vision/image_classification),
+so the checkpoints are not compatible.
+We will unify the implementation soon.
+
+
+### Train a SpineNet-49 based RetinaNet.
+
+```bash
+TPU_NAME="<your GCP TPU name>"
+MODEL_DIR="<path to the directory to store model files>"
+TRAIN_FILE_PATTERN="<path to the TFRecord training data>"
+EVAL_FILE_PATTERN="<path to the TFRecord validation data>"
+VAL_JSON_FILE="<path to the validation annotation JSON file>"
+python3 ~/models/official/vision/detection/main.py \
+  --strategy_type=tpu \
+  --tpu="${TPU_NAME?}" \
+  --model_dir="${MODEL_DIR?}" \
+  --mode=train \
+  --params_override="{ type: retinanet, architecture: {backbone: spinenet, multilevel_features: identity}, spinenet: {model_id: 49}, train_file_pattern: ${TRAIN_FILE_PATTERN?} }, eval: { val_json_file: ${VAL_JSON_FILE?}, eval_file_pattern: ${EVAL_FILE_PATTERN?} } }"
+```
+
+
+### Train a custom RetinaNet using the config file.
+
+First, create a YAML config file, e.g. *my_retinanet.yaml*. This file specifies
+the parameters to be overridden, which should at least include the following
+fields.
+
+```YAML
+# my_retinanet.yaml
+type: 'retinanet'
+train:
+  train_file_pattern: <path to the TFRecord training data>
+eval:
+  eval_file_pattern: <path to the TFRecord validation data>
+  val_json_file: <path to the validation annotation JSON file>
+```
+
+Once the YAML config file is created, you can launch the training using the
+following command.
+
+```bash
+TPU_NAME="<your GCP TPU name>"
+MODEL_DIR="<path to the directory to store model files>"
+python3 ~/models/official/vision/detection/main.py \
+  --strategy_type=tpu \
+  --tpu="${TPU_NAME?}" \
+  --model_dir="${MODEL_DIR?}" \
+  --mode=train \
+  --config_file="my_retinanet.yaml"
+```
+
+## Train RetinaNet on GPU
+
+Training on GPU is similar to that on TPU. The major change is the strategy
+type (use "[mirrored](https://www.tensorflow.org/api_docs/python/tf/distribute/MirroredStrategy)" for multiple GPU and
+"[one_device](https://www.tensorflow.org/api_docs/python/tf/distribute/OneDeviceStrategy)" for single GPU).
+
+Multi-GPUs example (assuming there are 8GPU connected to the host):
+
+```bash
+MODEL_DIR="<path to the directory to store model files>"
+python3 ~/models/official/vision/detection/main.py \
+  --strategy_type=mirrored \
+  --num_gpus=8 \
+  --model_dir="${MODEL_DIR?}" \
+  --mode=train \
+  --config_file="my_retinanet.yaml"
+```
+
+```bash
+MODEL_DIR="<path to the directory to store model files>"
+python3 ~/models/official/vision/detection/main.py \
+  --strategy_type=one_device \
+  --num_gpus=1 \
+  --model_dir="${MODEL_DIR?}" \
+  --mode=train \
+  --config_file="my_retinanet.yaml"
+```
+
+An example with inline configuration (YAML or JSON format):
+
+```
+python3 ~/models/official/vision/detection/main.py \
+  --model_dir=<model folder> \
+  --strategy_type=one_device \
+  --num_gpus=1 \
+  --mode=train \
+  --params_override="eval:
+ eval_file_pattern: <Eval TFRecord file pattern>
+ batch_size: 8
+ val_json_file: <COCO format groundtruth JSON file>
+predict:
+ predict_batch_size: 8
+architecture:
+ use_bfloat16: False
+train:
+ total_steps: 1
+ batch_size: 8
+ train_file_pattern: <Eval TFRecord file pattern>
+use_tpu: False
+"
+```
+
+---
+
+## Train Mask R-CNN on TPU
+
+### Train a vanilla ResNet-50 based Mask R-CNN.
+
+```bash
+TPU_NAME="<your GCP TPU name>"
+MODEL_DIR="<path to the directory to store model files>"
+RESNET_CHECKPOINT="<path to the pre-trained Resnet-50 checkpoint>"
+TRAIN_FILE_PATTERN="<path to the TFRecord training data>"
+EVAL_FILE_PATTERN="<path to the TFRecord validation data>"
+VAL_JSON_FILE="<path to the validation annotation JSON file>"
+python3 ~/models/official/vision/detection/main.py \
+  --strategy_type=tpu \
+  --tpu=${TPU_NAME} \
+  --model_dir=${MODEL_DIR} \
+  --mode=train \
+  --model=mask_rcnn \
+  --params_override="{train: { checkpoint: { path: ${RESNET_CHECKPOINT}, prefix: resnet50/ }, train_file_pattern: ${TRAIN_FILE_PATTERN} }, eval: { val_json_file: ${VAL_JSON_FILE}, eval_file_pattern: ${EVAL_FILE_PATTERN} } }"
+```
+
+The pre-trained ResNet-50 checkpoint can be downloaded [here](https://storage.cloud.google.com/cloud-tpu-checkpoints/model-garden-vision/detection/resnet50-2018-02-07.tar.gz).
+
+Note: The ResNet implementation under
+[detection/](https://github.com/tensorflow/models/tree/master/official/vision/detection)
+is currently different from the one under
+[classification/](https://github.com/tensorflow/models/tree/master/official/vision/image_classification),
+so the checkpoints are not compatible.
+We will unify the implementation soon.
+
+
+### Train a SpineNet-49 based Mask R-CNN.
+
+```bash
+TPU_NAME="<your GCP TPU name>"
+MODEL_DIR="<path to the directory to store model files>"
+TRAIN_FILE_PATTERN="<path to the TFRecord training data>"
+EVAL_FILE_PATTERN="<path to the TFRecord validation data>"
+VAL_JSON_FILE="<path to the validation annotation JSON file>"
+python3 ~/models/official/vision/detection/main.py \
+  --strategy_type=tpu \
+  --tpu="${TPU_NAME?}" \
+  --model_dir="${MODEL_DIR?}" \
+  --mode=train \
+  --model=mask_rcnn \
+  --params_override="{architecture: {backbone: spinenet, multilevel_features: identity}, spinenet: {model_id: 49}, train_file_pattern: ${TRAIN_FILE_PATTERN?} }, eval: { val_json_file: ${VAL_JSON_FILE?}, eval_file_pattern: ${EVAL_FILE_PATTERN?} } }"
+```
+
+
+### Train a custom Mask R-CNN using the config file.
+
+First, create a YAML config file, e.g. *my_maskrcnn.yaml*.
+This file specifies the parameters to be overridden,
+which should at least include the following fields.
+
+```YAML
+# my_maskrcnn.yaml
+train:
+  train_file_pattern: <path to the TFRecord training data>
+eval:
+  eval_file_pattern: <path to the TFRecord validation data>
+  val_json_file: <path to the validation annotation JSON file>
+```
+
+Once the YAML config file is created, you can launch the training using the
+following command.
+
+```bash
+TPU_NAME="<your GCP TPU name>"
+MODEL_DIR="<path to the directory to store model files>"
+python3 ~/models/official/vision/detection/main.py \
+  --strategy_type=tpu \
+  --tpu=${TPU_NAME} \
+  --model_dir=${MODEL_DIR} \
+  --mode=train \
+  --model=mask_rcnn \
+  --config_file="my_maskrcnn.yaml"
+```
+
+## Train Mask R-CNN on GPU
+
+Training on GPU is similar to that on TPU. The major change is the strategy type
+(use
+"[mirrored](https://www.tensorflow.org/api_docs/python/tf/distribute/MirroredStrategy)"
+for multiple GPU and
+"[one_device](https://www.tensorflow.org/api_docs/python/tf/distribute/OneDeviceStrategy)"
+for single GPU).
+
+Multi-GPUs example (assuming there are 8GPU connected to the host):
+
+```bash
+MODEL_DIR="<path to the directory to store model files>"
+python3 ~/models/official/vision/detection/main.py \
+  --strategy_type=mirrored \
+  --num_gpus=8 \
+  --model_dir=${MODEL_DIR} \
+  --mode=train \
+  --model=mask_rcnn \
+  --config_file="my_maskrcnn.yaml"
+```
+
+```bash
+MODEL_DIR="<path to the directory to store model files>"
+python3 ~/models/official/vision/detection/main.py \
+  --strategy_type=one_device \
+  --num_gpus=1 \
+  --model_dir=${MODEL_DIR} \
+  --mode=train \
+  --model=mask_rcnn \
+  --config_file="my_maskrcnn.yaml"
+```
+
+An example with inline configuration (YAML or JSON format):
+
+```
+python3 ~/models/official/vision/detection/main.py \
+  --model_dir=<model folder> \
+  --strategy_type=one_device \
+  --num_gpus=1 \
+  --mode=train \
+  --model=mask_rcnn \
+  --params_override="eval:
+ eval_file_pattern: <Eval TFRecord file pattern>
+ batch_size: 8
+ val_json_file: <COCO format groundtruth JSON file>
+predict:
+ predict_batch_size: 8
+architecture:
+ use_bfloat16: False
+train:
+ total_steps: 1000
+ batch_size: 8
+ train_file_pattern: <Eval TFRecord file pattern>
+use_tpu: False
+"
+```
+
+## Train ShapeMask on TPU
+
+### Train a ResNet-50 based ShapeMask.
+
+```bash
+TPU_NAME="<your GCP TPU name>"
+MODEL_DIR="<path to the directory to store model files>"
+RESNET_CHECKPOINT="<path to the pre-trained Resnet-50 checkpoint>"
+TRAIN_FILE_PATTERN="<path to the TFRecord training data>"
+EVAL_FILE_PATTERN="<path to the TFRecord validation data>"
+VAL_JSON_FILE="<path to the validation annotation JSON file>"
+SHAPE_PRIOR_PATH="<path to shape priors>"
+python3 ~/models/official/vision/detection/main.py \
+  --strategy_type=tpu \
+  --tpu=${TPU_NAME} \
+  --model_dir=${MODEL_DIR} \
+  --mode=train \
+  --model=shapemask \
+  --params_override="{train: { checkpoint: { path: ${RESNET_CHECKPOINT}, prefix: resnet50/ }, train_file_pattern: ${TRAIN_FILE_PATTERN} }, eval: { val_json_file: ${VAL_JSON_FILE}, eval_file_pattern: ${EVAL_FILE_PATTERN} } shapemask_head: {use_category_for_mask: true, shape_prior_path: ${SHAPE_PRIOR_PATH}} }"
+```
+
+The pre-trained ResNet-50 checkpoint can be downloaded [here](https://storage.cloud.google.com/cloud-tpu-checkpoints/model-garden-vision/detection/resnet50-2018-02-07.tar.gz).
+
+The shape priors can be downloaded [here]
+(https://storage.googleapis.com/cloud-tpu-checkpoints/shapemask/kmeans_class_priors_91x20x32x32.npy)
+
+
+### Train a custom ShapeMask using the config file.
+
+First, create a YAML config file, e.g. *my_shapemask.yaml*.
+This file specifies the parameters to be overridden:
+
+```YAML
+# my_shapemask.yaml
+train:
+  train_file_pattern: <path to the TFRecord training data>
+  total_steps: <total steps to train>
+  batch_size: <training batch size>
+eval:
+  eval_file_pattern: <path to the TFRecord validation data>
+  val_json_file: <path to the validation annotation JSON file>
+  batch_size: <evaluation batch size>
+shapemask_head:
+  shape_prior_path: <path to shape priors>
+```
+
+Once the YAML config file is created, you can launch the training using the
+following command.
+
+```bash
+TPU_NAME="<your GCP TPU name>"
+MODEL_DIR="<path to the directory to store model files>"
+python3 ~/models/official/vision/detection/main.py \
+  --strategy_type=tpu \
+  --tpu=${TPU_NAME} \
+  --model_dir=${MODEL_DIR} \
+  --mode=train \
+  --model=shapemask \
+  --config_file="my_shapemask.yaml"
+```
+
+## Train ShapeMask on GPU
+
+Training on GPU is similar to that on TPU. The major change is the strategy type
+(use
+"[mirrored](https://www.tensorflow.org/api_docs/python/tf/distribute/MirroredStrategy)"
+for multiple GPU and
+"[one_device](https://www.tensorflow.org/api_docs/python/tf/distribute/OneDeviceStrategy)"
+for single GPU).
+
+Multi-GPUs example (assuming there are 8GPU connected to the host):
+
+```bash
+MODEL_DIR="<path to the directory to store model files>"
+python3 ~/models/official/vision/detection/main.py \
+  --strategy_type=mirrored \
+  --num_gpus=8 \
+  --model_dir=${MODEL_DIR} \
+  --mode=train \
+  --model=shapemask \
+  --config_file="my_shapemask.yaml"
+```
+
+A single GPU example
+
+```bash
+MODEL_DIR="<path to the directory to store model files>"
+python3 ~/models/official/vision/detection/main.py \
+  --strategy_type=one_device \
+  --num_gpus=1 \
+  --model_dir=${MODEL_DIR} \
+  --mode=train \
+  --model=shapemask \
+  --config_file="my_shapemask.yaml"
+```
+
+
+An example with inline configuration (YAML or JSON format):
+
+```
+python3 ~/models/official/vision/detection/main.py \
+  --model_dir=<model folder> \
+  --strategy_type=one_device \
+  --num_gpus=1 \
+  --mode=train \
+  --model=shapemask \
+  --params_override="eval:
+ eval_file_pattern: <Eval TFRecord file pattern>
+ batch_size: 8
+ val_json_file: <COCO format groundtruth JSON file>
+train:
+ total_steps: 1000
+ batch_size: 8
+ train_file_pattern: <Eval TFRecord file pattern>
+use_tpu: False
+"
+```
+
+
+### Run the evaluation (after training)
+
+```
+python3 /usr/share/models/official/vision/detection/main.py \
+   --strategy_type=tpu \
+   --tpu=${TPU_NAME} \
+   --model_dir=${MODEL_DIR} \
+   --mode=eval \
+   --model=shapemask \
+   --params_override="{eval: { val_json_file: ${VAL_JSON_FILE}, eval_file_pattern: ${EVAL_FILE_PATTERN}, eval_samples: 5000 } }"
+```
+
+`MODEL_DIR` needs to point to the trained path of ShapeMask model.
+Change `strategy_type=mirrored` and `num_gpus=1` to run on a GPU.
+
+Note: The JSON groundtruth file is useful for [COCO dataset](http://cocodataset.org/#home) and can be
+downloaded from the [COCO website](http://cocodataset.org/#download). For custom dataset, it is unncessary because the groundtruth can be included in the TFRecord files.
+
+## References
+
+1.  [Focal Loss for Dense Object Detection](https://arxiv.org/abs/1708.02002).
+    Tsung-Yi Lin, Priya Goyal, Ross Girshick, Kaiming He, and Piotr Dollár. IEEE
+    International Conference on Computer Vision (ICCV), 2017.

official/legacy/detection/__init__.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.
+

official/legacy/detection/configs/__init__.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.
+

official/legacy/detection/configs/base_config.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.
+
+"""Base config template."""
+
+
+BACKBONES = [
+    'resnet',
+    'spinenet',
+]
+
+MULTILEVEL_FEATURES = [
+    'fpn',
+    'identity',
+]
+
+# pylint: disable=line-too-long
+# For ResNet, this freezes the variables of the first conv1 and conv2_x
+# layers [1], which leads to higher training speed and slightly better testing
+# accuracy. The intuition is that the low-level architecture (e.g., ResNet-50)
+# is able to capture low-level features such as edges; therefore, it does not
+# need to be fine-tuned for the detection task.
+# Note that we need to trailing `/` to avoid the incorrect match.
+# [1]: https://github.com/facebookresearch/Detectron/blob/master/detectron/core/config.py#L198
+RESNET_FROZEN_VAR_PREFIX = r'(resnet\d+)\/(conv2d(|_([1-9]|10))|batch_normalization(|_([1-9]|10)))\/'
+REGULARIZATION_VAR_REGEX = r'.*(kernel|weight):0$'
+
+BASE_CFG = {
+    'model_dir': '',
+    'use_tpu': True,
+    'strategy_type': 'tpu',
+    'isolate_session_state': False,
+    'train': {
+        'iterations_per_loop': 100,
+        'batch_size': 64,
+        'total_steps': 22500,
+        'num_cores_per_replica': None,
+        'input_partition_dims': None,
+        'optimizer': {
+            'type': 'momentum',
+            'momentum': 0.9,
+            'nesterov': True,  # `False` is better for TPU v3-128.
+        },
+        'learning_rate': {
+            'type': 'step',
+            'warmup_learning_rate': 0.0067,
+            'warmup_steps': 500,
+            'init_learning_rate': 0.08,
+            'learning_rate_levels': [0.008, 0.0008],
+            'learning_rate_steps': [15000, 20000],
+        },
+        'checkpoint': {
+            'path': '',
+            'prefix': '',
+        },
+        # One can use 'RESNET_FROZEN_VAR_PREFIX' to speed up ResNet training
+        # when loading from the checkpoint.
+        'frozen_variable_prefix': '',
+        'train_file_pattern': '',
+        'train_dataset_type': 'tfrecord',
+        # TODO(b/142174042): Support transpose_input option.
+        'transpose_input': False,
+        'regularization_variable_regex': REGULARIZATION_VAR_REGEX,
+        'l2_weight_decay': 0.0001,
+        'gradient_clip_norm': 0.0,
+        'input_sharding': False,
+    },
+    'eval': {
+        'input_sharding': True,
+        'batch_size': 8,
+        'eval_samples': 5000,
+        'min_eval_interval': 180,
+        'eval_timeout': None,
+        'num_steps_per_eval': 1000,
+        'type': 'box',
+        'use_json_file': True,
+        'val_json_file': '',
+        'eval_file_pattern': '',
+        'eval_dataset_type': 'tfrecord',
+        # When visualizing images, set evaluation batch size to 40 to avoid
+        # potential OOM.
+        'num_images_to_visualize': 0,
+    },
+    'predict': {
+        'batch_size': 8,
+    },
+    'architecture': {
+        'backbone': 'resnet',
+        'min_level': 3,
+        'max_level': 7,
+        'multilevel_features': 'fpn',
+        'use_bfloat16': True,
+        # Note that `num_classes` is the total number of classes including
+        # one background classes whose index is 0.
+        'num_classes': 91,
+    },
+    'anchor': {
+        'num_scales': 3,
+        'aspect_ratios': [1.0, 2.0, 0.5],
+        'anchor_size': 4.0,
+    },
+    'norm_activation': {
+        'activation': 'relu',
+        'batch_norm_momentum': 0.997,
+        'batch_norm_epsilon': 1e-4,
+        'batch_norm_trainable': True,
+        'use_sync_bn': False,
+    },
+    'resnet': {
+        'resnet_depth': 50,
+    },
+    'spinenet': {
+        'model_id': '49',
+    },
+    'fpn': {
+        'fpn_feat_dims': 256,
+        'use_separable_conv': False,
+        'use_batch_norm': True,
+    },
+    'postprocess': {
+        'use_batched_nms': False,
+        'max_total_size': 100,
+        'nms_iou_threshold': 0.5,
+        'score_threshold': 0.05,
+        'pre_nms_num_boxes': 5000,
+    },
+    'enable_summary': False,
+}
+# pylint: enable=line-too-long

official/legacy/detection/configs/factory.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.
+
+"""Factory to provide model configs."""
+
+from official.legacy.detection.configs import maskrcnn_config
+from official.legacy.detection.configs import olnmask_config
+from official.legacy.detection.configs import retinanet_config
+from official.legacy.detection.configs import shapemask_config
+from official.modeling.hyperparams import params_dict
+
+
+def config_generator(model):
+  """Model function generator."""
+  if model == 'retinanet':
+    default_config = retinanet_config.RETINANET_CFG
+    restrictions = retinanet_config.RETINANET_RESTRICTIONS
+  elif model == 'mask_rcnn':
+    default_config = maskrcnn_config.MASKRCNN_CFG
+    restrictions = maskrcnn_config.MASKRCNN_RESTRICTIONS
+  elif model == 'olnmask':
+    default_config = olnmask_config.OLNMASK_CFG
+    restrictions = olnmask_config.OLNMASK_RESTRICTIONS
+  elif model == 'shapemask':
+    default_config = shapemask_config.SHAPEMASK_CFG
+    restrictions = shapemask_config.SHAPEMASK_RESTRICTIONS
+  else:
+    raise ValueError('Model %s is not supported.' % model)
+
+  return params_dict.ParamsDict(default_config, restrictions)

official/legacy/detection/configs/maskrcnn_config.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.
+
+"""Config template to train Mask R-CNN."""
+
+from official.legacy.detection.configs import base_config
+from official.modeling.hyperparams import params_dict
+
+
+# pylint: disable=line-too-long
+MASKRCNN_CFG = params_dict.ParamsDict(base_config.BASE_CFG)
+MASKRCNN_CFG.override({
+    'type': 'mask_rcnn',
+    'eval': {
+        'type': 'box_and_mask',
+        'num_images_to_visualize': 0,
+    },
+    'architecture': {
+        'parser': 'maskrcnn_parser',
+        'min_level': 2,
+        'max_level': 6,
+        'include_mask': True,
+        'mask_target_size': 28,
+    },
+    'maskrcnn_parser': {
+        'output_size': [1024, 1024],
+        'num_channels': 3,
+        'rpn_match_threshold': 0.7,
+        'rpn_unmatched_threshold': 0.3,
+        'rpn_batch_size_per_im': 256,
+        'rpn_fg_fraction': 0.5,
+        'aug_rand_hflip': True,
+        'aug_scale_min': 1.0,
+        'aug_scale_max': 1.0,
+        'skip_crowd_during_training': True,
+        'max_num_instances': 100,
+        'mask_crop_size': 112,
+    },
+    'anchor': {
+        'num_scales': 1,
+        'anchor_size': 8,
+    },
+    'rpn_head': {
+        'num_convs': 2,
+        'num_filters': 256,
+        'use_separable_conv': False,
+        'use_batch_norm': False,
+    },
+    'frcnn_head': {
+        'num_convs': 0,
+        'num_filters': 256,
+        'use_separable_conv': False,
+        'num_fcs': 2,
+        'fc_dims': 1024,
+        'use_batch_norm': False,
+    },
+    'mrcnn_head': {
+        'num_convs': 4,
+        'num_filters': 256,
+        'use_separable_conv': False,
+        'use_batch_norm': False,
+    },
+    'rpn_score_loss': {
+        'rpn_batch_size_per_im': 256,
+    },
+    'rpn_box_loss': {
+        'huber_loss_delta': 1.0 / 9.0,
+    },
+    'frcnn_box_loss': {
+        'huber_loss_delta': 1.0,
+    },
+    'roi_proposal': {
+        'rpn_pre_nms_top_k': 2000,
+        'rpn_post_nms_top_k': 1000,
+        'rpn_nms_threshold': 0.7,
+        'rpn_score_threshold': 0.0,
+        'rpn_min_size_threshold': 0.0,
+        'test_rpn_pre_nms_top_k': 1000,
+        'test_rpn_post_nms_top_k': 1000,
+        'test_rpn_nms_threshold': 0.7,
+        'test_rpn_score_threshold': 0.0,
+        'test_rpn_min_size_threshold': 0.0,
+        'use_batched_nms': False,
+    },
+    'roi_sampling': {
+        'num_samples_per_image': 512,
+        'fg_fraction': 0.25,
+        'fg_iou_thresh': 0.5,
+        'bg_iou_thresh_hi': 0.5,
+        'bg_iou_thresh_lo': 0.0,
+        'mix_gt_boxes': True,
+    },
+    'mask_sampling': {
+        'num_mask_samples_per_image': 128,  # Typically = `num_samples_per_image` * `fg_fraction`.
+    },
+    'postprocess': {
+        'pre_nms_num_boxes': 1000,
+    },
+}, is_strict=False)
+
+
+MASKRCNN_RESTRICTIONS = [
+]
+# pylint: enable=line-too-long

official/legacy/detection/configs/olnmask_config.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.
+
+"""Config template to train Object Localization Network (OLN)."""
+
+from official.legacy.detection.configs import base_config
+from official.modeling.hyperparams import params_dict
+
+
+# pylint: disable=line-too-long
+OLNMASK_CFG = params_dict.ParamsDict(base_config.BASE_CFG)
+OLNMASK_CFG.override({
+    'type': 'olnmask',
+    'eval': {
+        'type': 'oln_xclass_box',
+        'use_category': False,
+        'seen_class': 'voc',
+        'num_images_to_visualize': 0,
+    },
+    'architecture': {
+        'parser': 'olnmask_parser',
+        'min_level': 2,
+        'max_level': 6,
+        'include_rpn_class': False,
+        'include_frcnn_class': False,
+        'include_frcnn_box': True,
+        'include_mask': False,
+        'mask_target_size': 28,
+        'num_classes': 2,
+    },
+    'olnmask_parser': {
+        'output_size': [640, 640],
+        'num_channels': 3,
+        'rpn_match_threshold': 0.7,
+        'rpn_unmatched_threshold': 0.3,
+        'rpn_batch_size_per_im': 256,
+        'rpn_fg_fraction': 0.5,
+        'aug_rand_hflip': True,
+        'aug_scale_min': 0.5,
+        'aug_scale_max': 2.0,
+        'skip_crowd_during_training': True,
+        'max_num_instances': 100,
+        'mask_crop_size': 112,
+        # centerness targets.
+        'has_centerness': True,
+        'rpn_center_match_iou_threshold': 0.3,
+        'rpn_center_unmatched_iou_threshold': 0.1,
+        'rpn_num_center_samples_per_im': 256,
+        # class manipulation.
+        'class_agnostic': True,
+        'train_class': 'voc',
+    },
+    'anchor': {
+        'num_scales': 1,
+        'aspect_ratios': [1.0],
+        'anchor_size': 8,
+    },
+    'rpn_head': {
+        'num_convs': 2,
+        'num_filters': 256,
+        'use_separable_conv': False,
+        'use_batch_norm': False,
+        # RPN-Centerness learning {
+        'has_centerness': True,  # }
+    },
+    'frcnn_head': {
+        'num_convs': 0,
+        'num_filters': 256,
+        'use_separable_conv': False,
+        'num_fcs': 2,
+        'fc_dims': 1024,
+        'use_batch_norm': False,
+        'has_scoring': True,
+    },
+    'mrcnn_head': {
+        'num_convs': 4,
+        'num_filters': 256,
+        'use_separable_conv': False,
+        'use_batch_norm': False,
+        'has_scoring': False,
+    },
+    'rpn_score_loss': {
+        'rpn_batch_size_per_im': 256,
+    },
+    'rpn_box_loss': {
+        'huber_loss_delta': 1.0 / 9.0,
+    },
+    'frcnn_box_loss': {
+        'huber_loss_delta': 1.0,
+    },
+    'frcnn_box_score_loss': {
+        'ignore_threshold': 0.3,
+    },
+    'roi_proposal': {
+        'rpn_pre_nms_top_k': 2000,
+        'rpn_post_nms_top_k': 2000,
+        'rpn_nms_threshold': 0.7,
+        'rpn_score_threshold': 0.0,
+        'rpn_min_size_threshold': 0.0,
+        'test_rpn_pre_nms_top_k': 2000,
+        'test_rpn_post_nms_top_k': 2000,
+        'test_rpn_nms_threshold': 0.7,
+        'test_rpn_score_threshold': 0.0,
+        'test_rpn_min_size_threshold': 0.0,
+        'use_batched_nms': False,
+    },
+    'roi_sampling': {
+        'num_samples_per_image': 512,
+        'fg_fraction': 0.25,
+        'fg_iou_thresh': 0.5,
+        'bg_iou_thresh_hi': 0.5,
+        'bg_iou_thresh_lo': 0.0,
+        'mix_gt_boxes': True,
+    },
+    'mask_sampling': {
+        'num_mask_samples_per_image': 128,  # Typically = `num_samples_per_image` * `fg_fraction`.
+    },
+    'postprocess': {
+        'use_batched_nms': False,
+        'max_total_size': 100,
+        'nms_iou_threshold': 0.5,
+        'score_threshold': 0.00,
+        'pre_nms_num_boxes': 2000,
+    },
+}, is_strict=False)
+
+
+OLNMASK_RESTRICTIONS = [
+    # 'anchor.aspect_ratios == [1.0]',
+    # 'anchor.scales == 1',
+]
+# pylint: enable=line-too-long

official/legacy/detection/configs/retinanet_config.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.
+
+"""Config template to train Retinanet."""
+
+from official.legacy.detection.configs import base_config
+from official.modeling.hyperparams import params_dict
+
+
+# pylint: disable=line-too-long
+RETINANET_CFG = params_dict.ParamsDict(base_config.BASE_CFG)
+RETINANET_CFG.override({
+    'type': 'retinanet',
+    'architecture': {
+        'parser': 'retinanet_parser',
+    },
+    'retinanet_parser': {
+        'output_size': [640, 640],
+        'num_channels': 3,
+        'match_threshold': 0.5,
+        'unmatched_threshold': 0.5,
+        'aug_rand_hflip': True,
+        'aug_scale_min': 1.0,
+        'aug_scale_max': 1.0,
+        'use_autoaugment': False,
+        'autoaugment_policy_name': 'v0',
+        'skip_crowd_during_training': True,
+        'max_num_instances': 100,
+    },
+    'retinanet_head': {
+        'num_convs': 4,
+        'num_filters': 256,
+        'use_separable_conv': False,
+    },
+    'retinanet_loss': {
+        'focal_loss_alpha': 0.25,
+        'focal_loss_gamma': 1.5,
+        'huber_loss_delta': 0.1,
+        'box_loss_weight': 50,
+    },
+    'enable_summary': True,
+}, is_strict=False)
+
+RETINANET_RESTRICTIONS = [
+]
+
+# pylint: enable=line-too-long

official/legacy/detection/configs/shapemask_config.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.
+
+"""Config to train shapemask on COCO."""
+
+from official.legacy.detection.configs import base_config
+from official.modeling.hyperparams import params_dict
+
+SHAPEMASK_RESNET_FROZEN_VAR_PREFIX = r'(conv2d(|_([1-9]|10))|batch_normalization(|_([1-9]|10)))\/'
+
+SHAPEMASK_CFG = params_dict.ParamsDict(base_config.BASE_CFG)
+SHAPEMASK_CFG.override({
+    'type': 'shapemask',
+    'architecture': {
+        'parser': 'shapemask_parser',
+        'backbone': 'resnet',
+        'multilevel_features': 'fpn',
+        'outer_box_scale': 1.25,
+    },
+    'train': {
+        'total_steps': 45000,
+        'learning_rate': {
+            'learning_rate_steps': [30000, 40000],
+        },
+        'frozen_variable_prefix': SHAPEMASK_RESNET_FROZEN_VAR_PREFIX,
+        'regularization_variable_regex': None,
+    },
+    'eval': {
+        'type': 'shapemask_box_and_mask',
+        'mask_eval_class': 'all',  # 'all', 'voc', or 'nonvoc'.
+    },
+    'shapemask_parser': {
+        'output_size': [640, 640],
+        'num_channels': 3,
+        'match_threshold': 0.5,
+        'unmatched_threshold': 0.5,
+        'aug_rand_hflip': True,
+        'aug_scale_min': 0.8,
+        'aug_scale_max': 1.2,
+        'skip_crowd_during_training': True,
+        'max_num_instances': 100,
+        # Shapemask specific parameters
+        'mask_train_class': 'all',  # 'all', 'voc', or 'nonvoc'.
+        'use_category': True,
+        'outer_box_scale': 1.25,
+        'num_sampled_masks': 8,
+        'mask_crop_size': 32,
+        'mask_min_level': 3,
+        'mask_max_level': 5,
+        'box_jitter_scale': 0.025,
+        'upsample_factor': 4,
+    },
+    'retinanet_head': {
+        'num_convs': 4,
+        'num_filters': 256,
+        'use_separable_conv': False,
+        'use_batch_norm': True,
+    },
+    'shapemask_head': {
+        'num_downsample_channels': 128,
+        'mask_crop_size': 32,
+        'use_category_for_mask': True,
+        'num_convs': 4,
+        'upsample_factor': 4,
+        'shape_prior_path': '',
+    },
+    'retinanet_loss': {
+        'focal_loss_alpha': 0.4,
+        'focal_loss_gamma': 1.5,
+        'huber_loss_delta': 0.15,
+        'box_loss_weight': 50,
+    },
+    'shapemask_loss': {
+        'shape_prior_loss_weight': 0.1,
+        'coarse_mask_loss_weight': 1.0,
+        'fine_mask_loss_weight': 1.0,
+    },
+}, is_strict=False)
+
+SHAPEMASK_RESTRICTIONS = [
+    'shapemask_head.mask_crop_size == shapemask_parser.mask_crop_size',
+    'shapemask_head.upsample_factor == shapemask_parser.upsample_factor',
+    'shapemask_parser.outer_box_scale ==  architecture.outer_box_scale',
+]
+
+# pylint: enable=line-too-long

official/legacy/detection/dataloader/__init__.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.
+

official/legacy/detection/dataloader/factory.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.
+
+"""Model architecture factory."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from official.legacy.detection.dataloader import maskrcnn_parser
+from official.legacy.detection.dataloader import olnmask_parser
+from official.legacy.detection.dataloader import retinanet_parser
+from official.legacy.detection.dataloader import shapemask_parser
+
+
+def parser_generator(params, mode):
+  """Generator function for various dataset parser."""
+  if params.architecture.parser == 'retinanet_parser':
+    anchor_params = params.anchor
+    parser_params = params.retinanet_parser
+    parser_fn = retinanet_parser.Parser(
+        output_size=parser_params.output_size,
+        min_level=params.architecture.min_level,
+        max_level=params.architecture.max_level,
+        num_scales=anchor_params.num_scales,
+        aspect_ratios=anchor_params.aspect_ratios,
+        anchor_size=anchor_params.anchor_size,
+        match_threshold=parser_params.match_threshold,
+        unmatched_threshold=parser_params.unmatched_threshold,
+        aug_rand_hflip=parser_params.aug_rand_hflip,
+        aug_scale_min=parser_params.aug_scale_min,
+        aug_scale_max=parser_params.aug_scale_max,
+        use_autoaugment=parser_params.use_autoaugment,
+        autoaugment_policy_name=parser_params.autoaugment_policy_name,
+        skip_crowd_during_training=parser_params.skip_crowd_during_training,
+        max_num_instances=parser_params.max_num_instances,
+        use_bfloat16=params.architecture.use_bfloat16,
+        mode=mode)
+  elif params.architecture.parser == 'maskrcnn_parser':
+    anchor_params = params.anchor
+    parser_params = params.maskrcnn_parser
+    parser_fn = maskrcnn_parser.Parser(
+        output_size=parser_params.output_size,
+        min_level=params.architecture.min_level,
+        max_level=params.architecture.max_level,
+        num_scales=anchor_params.num_scales,
+        aspect_ratios=anchor_params.aspect_ratios,
+        anchor_size=anchor_params.anchor_size,
+        rpn_match_threshold=parser_params.rpn_match_threshold,
+        rpn_unmatched_threshold=parser_params.rpn_unmatched_threshold,
+        rpn_batch_size_per_im=parser_params.rpn_batch_size_per_im,
+        rpn_fg_fraction=parser_params.rpn_fg_fraction,
+        aug_rand_hflip=parser_params.aug_rand_hflip,
+        aug_scale_min=parser_params.aug_scale_min,
+        aug_scale_max=parser_params.aug_scale_max,
+        skip_crowd_during_training=parser_params.skip_crowd_during_training,
+        max_num_instances=parser_params.max_num_instances,
+        include_mask=params.architecture.include_mask,
+        mask_crop_size=parser_params.mask_crop_size,
+        use_bfloat16=params.architecture.use_bfloat16,
+        mode=mode)
+  elif params.architecture.parser == 'olnmask_parser':
+    anchor_params = params.anchor
+    parser_params = params.olnmask_parser
+    parser_fn = olnmask_parser.Parser(
+        output_size=parser_params.output_size,
+        min_level=params.architecture.min_level,
+        max_level=params.architecture.max_level,
+        num_scales=anchor_params.num_scales,
+        aspect_ratios=anchor_params.aspect_ratios,
+        anchor_size=anchor_params.anchor_size,
+        rpn_match_threshold=parser_params.rpn_match_threshold,
+        rpn_unmatched_threshold=parser_params.rpn_unmatched_threshold,
+        rpn_batch_size_per_im=parser_params.rpn_batch_size_per_im,
+        rpn_fg_fraction=parser_params.rpn_fg_fraction,
+        aug_rand_hflip=parser_params.aug_rand_hflip,
+        aug_scale_min=parser_params.aug_scale_min,
+        aug_scale_max=parser_params.aug_scale_max,
+        skip_crowd_during_training=parser_params.skip_crowd_during_training,
+        max_num_instances=parser_params.max_num_instances,
+        include_mask=params.architecture.include_mask,
+        mask_crop_size=parser_params.mask_crop_size,
+        use_bfloat16=params.architecture.use_bfloat16,
+        mode=mode,
+        has_centerness=parser_params.has_centerness,
+        rpn_center_match_iou_threshold=(
+            parser_params.rpn_center_match_iou_threshold),
+        rpn_center_unmatched_iou_threshold=(
+            parser_params.rpn_center_unmatched_iou_threshold),
+        rpn_num_center_samples_per_im=(
+            parser_params.rpn_num_center_samples_per_im),
+        class_agnostic=parser_params.class_agnostic,
+        train_class=parser_params.train_class,)
+  elif params.architecture.parser == 'shapemask_parser':
+    anchor_params = params.anchor
+    parser_params = params.shapemask_parser
+    parser_fn = shapemask_parser.Parser(
+        output_size=parser_params.output_size,
+        min_level=params.architecture.min_level,
+        max_level=params.architecture.max_level,
+        num_scales=anchor_params.num_scales,
+        aspect_ratios=anchor_params.aspect_ratios,
+        anchor_size=anchor_params.anchor_size,
+        use_category=parser_params.use_category,
+        outer_box_scale=parser_params.outer_box_scale,
+        box_jitter_scale=parser_params.box_jitter_scale,
+        num_sampled_masks=parser_params.num_sampled_masks,
+        mask_crop_size=parser_params.mask_crop_size,
+        mask_min_level=parser_params.mask_min_level,
+        mask_max_level=parser_params.mask_max_level,
+        upsample_factor=parser_params.upsample_factor,
+        match_threshold=parser_params.match_threshold,
+        unmatched_threshold=parser_params.unmatched_threshold,
+        aug_rand_hflip=parser_params.aug_rand_hflip,
+        aug_scale_min=parser_params.aug_scale_min,
+        aug_scale_max=parser_params.aug_scale_max,
+        skip_crowd_during_training=parser_params.skip_crowd_during_training,
+        max_num_instances=parser_params.max_num_instances,
+        use_bfloat16=params.architecture.use_bfloat16,
+        mask_train_class=parser_params.mask_train_class,
+        mode=mode)
+  else:
+    raise ValueError('Parser %s is not supported.' % params.architecture.parser)
+
+  return parser_fn

official/legacy/detection/dataloader/input_reader.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.
+
+"""Data loader and input processing."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+from typing import Optional, Text
+import tensorflow as tf
+from official.legacy.detection.dataloader import factory
+from official.legacy.detection.dataloader import mode_keys as ModeKeys
+from official.modeling.hyperparams import params_dict
+
+
+class InputFn(object):
+  """Input function that creates dataset from files."""
+
+  def __init__(self,
+               file_pattern: Text,
+               params: params_dict.ParamsDict,
+               mode: Text,
+               batch_size: int,
+               num_examples: Optional[int] = -1):
+    """Initialize.
+
+    Args:
+      file_pattern: the file pattern for the data example (TFRecords).
+      params: the parameter object for constructing example parser and model.
+      mode: ModeKeys.TRAIN or ModeKeys.Eval
+      batch_size: the data batch size.
+      num_examples: If positive, only takes this number of examples and raise
+        tf.errors.OutOfRangeError after that. If non-positive, it will be
+        ignored.
+    """
+    assert file_pattern is not None
+    assert mode is not None
+    assert batch_size is not None
+    self._file_pattern = file_pattern
+    self._mode = mode
+    self._is_training = (mode == ModeKeys.TRAIN)
+    self._batch_size = batch_size
+    self._num_examples = num_examples
+    self._parser_fn = factory.parser_generator(params, mode)
+    self._dataset_fn = tf.data.TFRecordDataset
+
+    self._input_sharding = (not self._is_training)
+    try:
+      if self._is_training:
+        self._input_sharding = params.train.input_sharding
+      else:
+        self._input_sharding = params.eval.input_sharding
+    except AttributeError:
+      pass
+
+  def __call__(self, ctx=None, batch_size: int = None):
+    """Provides tf.data.Dataset object.
+
+    Args:
+      ctx: context object.
+      batch_size: expected batch size input data.
+
+    Returns:
+      tf.data.Dataset object.
+    """
+    if not batch_size:
+      batch_size = self._batch_size
+    assert batch_size is not None
+    dataset = tf.data.Dataset.list_files(
+        self._file_pattern, shuffle=self._is_training)
+
+    if self._input_sharding and ctx and ctx.num_input_pipelines > 1:
+      dataset = dataset.shard(ctx.num_input_pipelines, ctx.input_pipeline_id)
+    dataset = dataset.cache()
+
+    if self._is_training:
+      dataset = dataset.repeat()
+
+    dataset = dataset.interleave(
+        map_func=self._dataset_fn,
+        cycle_length=32,
+        num_parallel_calls=tf.data.experimental.AUTOTUNE)
+
+    if self._is_training:
+      dataset = dataset.shuffle(1000)
+    if self._num_examples > 0:
+      dataset = dataset.take(self._num_examples)
+
+    # Parses the fetched records to input tensors for model function.
+    dataset = dataset.map(
+        self._parser_fn, num_parallel_calls=tf.data.experimental.AUTOTUNE)
+    dataset = dataset.batch(batch_size, drop_remainder=True)
+    dataset = dataset.prefetch(tf.data.experimental.AUTOTUNE)
+    return dataset

official/legacy/detection/dataloader/maskrcnn_parser.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.
+
+"""Data parser and processing for Mask R-CNN."""
+
+import tensorflow as tf
+
+from official.legacy.detection.dataloader import anchor
+from official.legacy.detection.dataloader import mode_keys as ModeKeys
+from official.legacy.detection.dataloader import tf_example_decoder
+from official.legacy.detection.utils import box_utils
+from official.legacy.detection.utils import dataloader_utils
+from official.legacy.detection.utils import input_utils
+
+
+class Parser(object):
+  """Parser to parse an image and its annotations into a dictionary of tensors."""
+
+  def __init__(self,
+               output_size,
+               min_level,
+               max_level,
+               num_scales,
+               aspect_ratios,
+               anchor_size,
+               rpn_match_threshold=0.7,
+               rpn_unmatched_threshold=0.3,
+               rpn_batch_size_per_im=256,
+               rpn_fg_fraction=0.5,
+               aug_rand_hflip=False,
+               aug_scale_min=1.0,
+               aug_scale_max=1.0,
+               skip_crowd_during_training=True,
+               max_num_instances=100,
+               include_mask=False,
+               mask_crop_size=112,
+               use_bfloat16=True,
+               mode=None):
+    """Initializes parameters for parsing annotations in the dataset.
+
+    Args:
+      output_size: `Tensor` or `list` for [height, width] of output image. The
+        output_size should be divided by the largest feature stride 2^max_level.
+      min_level: `int` number of minimum level of the output feature pyramid.
+      max_level: `int` number of maximum level of the output feature pyramid.
+      num_scales: `int` number representing intermediate scales added
+        on each level. For instances, num_scales=2 adds one additional
+        intermediate anchor scales [2^0, 2^0.5] on each level.
+      aspect_ratios: `list` of float numbers representing the aspect raito
+        anchors added on each level. The number indicates the ratio of width to
+        height. For instances, aspect_ratios=[1.0, 2.0, 0.5] adds three anchors
+        on each scale level.
+      anchor_size: `float` number representing the scale of size of the base
+        anchor to the feature stride 2^level.
+      rpn_match_threshold:
+      rpn_unmatched_threshold:
+      rpn_batch_size_per_im:
+      rpn_fg_fraction:
+      aug_rand_hflip: `bool`, if True, augment training with random
+        horizontal flip.
+      aug_scale_min: `float`, the minimum scale applied to `output_size` for
+        data augmentation during training.
+      aug_scale_max: `float`, the maximum scale applied to `output_size` for
+        data augmentation during training.
+      skip_crowd_during_training: `bool`, if True, skip annotations labeled with
+        `is_crowd` equals to 1.
+      max_num_instances: `int` number of maximum number of instances in an
+        image. The groundtruth data will be padded to `max_num_instances`.
+      include_mask: a bool to indicate whether parse mask groundtruth.
+      mask_crop_size: the size which groundtruth mask is cropped to.
+      use_bfloat16: `bool`, if True, cast output image to tf.bfloat16.
+      mode: a ModeKeys. Specifies if this is training, evaluation, prediction
+        or prediction with groundtruths in the outputs.
+    """
+    self._mode = mode
+    self._max_num_instances = max_num_instances
+    self._skip_crowd_during_training = skip_crowd_during_training
+    self._is_training = (mode == ModeKeys.TRAIN)
+
+    self._example_decoder = tf_example_decoder.TfExampleDecoder(
+        include_mask=include_mask)
+
+    # Anchor.
+    self._output_size = output_size
+    self._min_level = min_level
+    self._max_level = max_level
+    self._num_scales = num_scales
+    self._aspect_ratios = aspect_ratios
+    self._anchor_size = anchor_size
+
+    # Target assigning.
+    self._rpn_match_threshold = rpn_match_threshold
+    self._rpn_unmatched_threshold = rpn_unmatched_threshold
+    self._rpn_batch_size_per_im = rpn_batch_size_per_im
+    self._rpn_fg_fraction = rpn_fg_fraction
+
+    # Data augmentation.
+    self._aug_rand_hflip = aug_rand_hflip
+    self._aug_scale_min = aug_scale_min
+    self._aug_scale_max = aug_scale_max
+
+    # Mask.
+    self._include_mask = include_mask
+    self._mask_crop_size = mask_crop_size
+
+    # Device.
+    self._use_bfloat16 = use_bfloat16
+
+    # Data is parsed depending on the model Modekey.
+    if mode == ModeKeys.TRAIN:
+      self._parse_fn = self._parse_train_data
+    elif mode == ModeKeys.EVAL:
+      self._parse_fn = self._parse_eval_data
+    elif mode == ModeKeys.PREDICT or mode == ModeKeys.PREDICT_WITH_GT:
+      self._parse_fn = self._parse_predict_data
+    else:
+      raise ValueError('mode is not defined.')
+
+  def __call__(self, value):
+    """Parses data to an image and associated training labels.
+
+    Args:
+      value: a string tensor holding a serialized tf.Example proto.
+
+    Returns:
+      image, labels: if mode == ModeKeys.TRAIN. see _parse_train_data.
+      {'images': image, 'labels': labels}: if mode == ModeKeys.PREDICT
+        or ModeKeys.PREDICT_WITH_GT.
+    """
+    with tf.name_scope('parser'):
+      data = self._example_decoder.decode(value)
+      return self._parse_fn(data)
+
+  def _parse_train_data(self, data):
+    """Parses data for training.
+
+    Args:
+      data: the decoded tensor dictionary from TfExampleDecoder.
+
+    Returns:
+      image: image tensor that is preproessed to have normalized value and
+        dimension [output_size[0], output_size[1], 3]
+      labels: a dictionary of tensors used for training. The following describes
+        {key: value} pairs in the dictionary.
+        image_info: a 2D `Tensor` that encodes the information of the image and
+          the applied preprocessing. It is in the format of
+          [[original_height, original_width], [scaled_height, scaled_width],
+        anchor_boxes: ordered dictionary with keys
+          [min_level, min_level+1, ..., max_level]. The values are tensor with
+          shape [height_l, width_l, 4] representing anchor boxes at each level.
+        rpn_score_targets: ordered dictionary with keys
+          [min_level, min_level+1, ..., max_level]. The values are tensor with
+          shape [height_l, width_l, anchors_per_location]. The height_l and
+          width_l represent the dimension of class logits at l-th level.
+        rpn_box_targets: ordered dictionary with keys
+          [min_level, min_level+1, ..., max_level]. The values are tensor with
+          shape [height_l, width_l, anchors_per_location * 4]. The height_l and
+          width_l represent the dimension of bounding box regression output at
+          l-th level.
+        gt_boxes: Groundtruth bounding box annotations. The box is represented
+           in [y1, x1, y2, x2] format. The coordinates are w.r.t the scaled
+           image that is fed to the network. The tennsor is padded with -1 to
+           the fixed dimension [self._max_num_instances, 4].
+        gt_classes: Groundtruth classes annotations. The tennsor is padded
+          with -1 to the fixed dimension [self._max_num_instances].
+        gt_masks: groundtrugh masks cropped by the bounding box and
+          resized to a fixed size determined by mask_crop_size.
+    """
+    classes = data['groundtruth_classes']
+    boxes = data['groundtruth_boxes']
+    if self._include_mask:
+      masks = data['groundtruth_instance_masks']
+
+    is_crowds = data['groundtruth_is_crowd']
+    # Skips annotations with `is_crowd` = True.
+    if self._skip_crowd_during_training and self._is_training:
+      num_groundtruths = tf.shape(classes)[0]
+      with tf.control_dependencies([num_groundtruths, is_crowds]):
+        indices = tf.cond(
+            tf.greater(tf.size(is_crowds), 0),
+            lambda: tf.where(tf.logical_not(is_crowds))[:, 0],
+            lambda: tf.cast(tf.range(num_groundtruths), tf.int64))
+      classes = tf.gather(classes, indices)
+      boxes = tf.gather(boxes, indices)
+      if self._include_mask:
+        masks = tf.gather(masks, indices)
+
+    # Gets original image and its size.
+    image = data['image']
+    image_shape = tf.shape(image)[0:2]
+
+    # Normalizes image with mean and std pixel values.
+    image = input_utils.normalize_image(image)
+
+    # Flips image randomly during training.
+    if self._aug_rand_hflip:
+      if self._include_mask:
+        image, boxes, masks = input_utils.random_horizontal_flip(
+            image, boxes, masks)
+      else:
+        image, boxes = input_utils.random_horizontal_flip(
+            image, boxes)
+
+    # Converts boxes from normalized coordinates to pixel coordinates.
+    # Now the coordinates of boxes are w.r.t. the original image.
+    boxes = box_utils.denormalize_boxes(boxes, image_shape)
+
+    # Resizes and crops image.
+    image, image_info = input_utils.resize_and_crop_image(
+        image,
+        self._output_size,
+        padded_size=input_utils.compute_padded_size(
+            self._output_size, 2 ** self._max_level),
+        aug_scale_min=self._aug_scale_min,
+        aug_scale_max=self._aug_scale_max)
+    image_height, image_width, _ = image.get_shape().as_list()
+
+    # Resizes and crops boxes.
+    # Now the coordinates of boxes are w.r.t the scaled image.
+    image_scale = image_info[2, :]
+    offset = image_info[3, :]
+    boxes = input_utils.resize_and_crop_boxes(
+        boxes, image_scale, image_info[1, :], offset)
+
+    # Filters out ground truth boxes that are all zeros.
+    indices = box_utils.get_non_empty_box_indices(boxes)
+    boxes = tf.gather(boxes, indices)
+    classes = tf.gather(classes, indices)
+    if self._include_mask:
+      masks = tf.gather(masks, indices)
+      # Transfer boxes to the original image space and do normalization.
+      cropped_boxes = boxes + tf.tile(tf.expand_dims(offset, axis=0), [1, 2])
+      cropped_boxes /= tf.tile(tf.expand_dims(image_scale, axis=0), [1, 2])
+      cropped_boxes = box_utils.normalize_boxes(cropped_boxes, image_shape)
+      num_masks = tf.shape(masks)[0]
+      masks = tf.image.crop_and_resize(
+          tf.expand_dims(masks, axis=-1),
+          cropped_boxes,
+          box_indices=tf.range(num_masks, dtype=tf.int32),
+          crop_size=[self._mask_crop_size, self._mask_crop_size],
+          method='bilinear')
+      masks = tf.squeeze(masks, axis=-1)
+
+    # Assigns anchor targets.
+    # Note that after the target assignment, box targets are absolute pixel
+    # offsets w.r.t. the scaled image.
+    input_anchor = anchor.Anchor(
+        self._min_level,
+        self._max_level,
+        self._num_scales,
+        self._aspect_ratios,
+        self._anchor_size,
+        (image_height, image_width))
+    anchor_labeler = anchor.RpnAnchorLabeler(
+        input_anchor,
+        self._rpn_match_threshold,
+        self._rpn_unmatched_threshold,
+        self._rpn_batch_size_per_im,
+        self._rpn_fg_fraction)
+    rpn_score_targets, rpn_box_targets = anchor_labeler.label_anchors(
+        boxes, tf.cast(tf.expand_dims(classes, axis=-1), dtype=tf.float32))
+
+    # If bfloat16 is used, casts input image to tf.bfloat16.
+    if self._use_bfloat16:
+      image = tf.cast(image, dtype=tf.bfloat16)
+
+    inputs = {
+        'image': image,
+        'image_info': image_info,
+    }
+    # Packs labels for model_fn outputs.
+    labels = {
+        'anchor_boxes': input_anchor.multilevel_boxes,
+        'image_info': image_info,
+        'rpn_score_targets': rpn_score_targets,
+        'rpn_box_targets': rpn_box_targets,
+    }
+    inputs['gt_boxes'] = input_utils.pad_to_fixed_size(boxes,
+                                                       self._max_num_instances,
+                                                       -1)
+    inputs['gt_classes'] = input_utils.pad_to_fixed_size(
+        classes, self._max_num_instances, -1)
+    if self._include_mask:
+      inputs['gt_masks'] = input_utils.pad_to_fixed_size(
+          masks, self._max_num_instances, -1)
+
+    return inputs, labels
+
+  def _parse_eval_data(self, data):
+    """Parses data for evaluation."""
+    raise NotImplementedError('Not implemented!')
+
+  def _parse_predict_data(self, data):
+    """Parses data for prediction.
+
+    Args:
+      data: the decoded tensor dictionary from TfExampleDecoder.
+
+    Returns:
+      A dictionary of {'images': image, 'labels': labels} where
+        image: image tensor that is preproessed to have normalized value and
+          dimension [output_size[0], output_size[1], 3]
+        labels: a dictionary of tensors used for training. The following
+          describes {key: value} pairs in the dictionary.
+          source_ids: Source image id. Default value -1 if the source id is
+            empty in the groundtruth annotation.
+          image_info: a 2D `Tensor` that encodes the information of the image
+            and the applied preprocessing. It is in the format of
+            [[original_height, original_width], [scaled_height, scaled_width],
+          anchor_boxes: ordered dictionary with keys
+            [min_level, min_level+1, ..., max_level]. The values are tensor with
+            shape [height_l, width_l, 4] representing anchor boxes at each
+            level.
+    """
+    # Gets original image and its size.
+    image = data['image']
+    image_shape = tf.shape(image)[0:2]
+
+    # Normalizes image with mean and std pixel values.
+    image = input_utils.normalize_image(image)
+
+    # Resizes and crops image.
+    image, image_info = input_utils.resize_and_crop_image(
+        image,
+        self._output_size,
+        padded_size=input_utils.compute_padded_size(
+            self._output_size, 2 ** self._max_level),
+        aug_scale_min=1.0,
+        aug_scale_max=1.0)
+    image_height, image_width, _ = image.get_shape().as_list()
+
+    # If bfloat16 is used, casts input image to tf.bfloat16.
+    if self._use_bfloat16:
+      image = tf.cast(image, dtype=tf.bfloat16)
+
+    # Compute Anchor boxes.
+    _ = anchor.Anchor(self._min_level, self._max_level, self._num_scales,
+                      self._aspect_ratios, self._anchor_size,
+                      (image_height, image_width))
+
+    labels = {
+        'image_info': image_info,
+    }
+
+    if self._mode == ModeKeys.PREDICT_WITH_GT:
+      # Converts boxes from normalized coordinates to pixel coordinates.
+      boxes = box_utils.denormalize_boxes(
+          data['groundtruth_boxes'], image_shape)
+      groundtruths = {
+          'source_id': data['source_id'],
+          'height': data['height'],
+          'width': data['width'],
+          'num_detections': tf.shape(data['groundtruth_classes']),
+          'boxes': boxes,
+          'classes': data['groundtruth_classes'],
+          'areas': data['groundtruth_area'],
+          'is_crowds': tf.cast(data['groundtruth_is_crowd'], tf.int32),
+      }
+      groundtruths['source_id'] = dataloader_utils.process_source_id(
+          groundtruths['source_id'])
+      groundtruths = dataloader_utils.pad_groundtruths_to_fixed_size(
+          groundtruths, self._max_num_instances)
+      # TODO(yeqing):  Remove the `groundtrtuh` layer key (no longer needed).
+      labels['groundtruths'] = groundtruths
+    inputs = {
+        'image': image,
+        'image_info': image_info,
+    }
+
+    return inputs, labels

official/legacy/detection/dataloader/mode_keys.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.
+
+"""Standard names for input dataloader modes.
+
+The following standard keys are defined:
+
+* `TRAIN`: training mode.
+* `EVAL`: evaluation mode.
+* `PREDICT`: prediction mode.
+* `PREDICT_WITH_GT`: prediction mode with groundtruths in returned variables.
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+
+TRAIN = 'train'
+EVAL = 'eval'
+PREDICT = 'predict'
+PREDICT_WITH_GT = 'predict_with_gt'

official/legacy/detection/dataloader/olnmask_parser.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.
+
+"""Data parser and processing for Mask R-CNN."""
+
+import tensorflow as tf
+
+from official.legacy.detection.dataloader import anchor
+from official.legacy.detection.dataloader.maskrcnn_parser import Parser as MaskrcnnParser
+from official.legacy.detection.utils import box_utils
+from official.legacy.detection.utils import class_utils
+from official.legacy.detection.utils import input_utils
+
+
+class Parser(MaskrcnnParser):
+  """Parser to parse an image and its annotations into a dictionary of tensors."""
+
+  def __init__(self,
+               output_size,
+               min_level,
+               max_level,
+               num_scales,
+               aspect_ratios,
+               anchor_size,
+               rpn_match_threshold=0.7,
+               rpn_unmatched_threshold=0.3,
+               rpn_batch_size_per_im=256,
+               rpn_fg_fraction=0.5,
+               aug_rand_hflip=False,
+               aug_scale_min=1.0,
+               aug_scale_max=1.0,
+               skip_crowd_during_training=True,
+               max_num_instances=100,
+               include_mask=False,
+               mask_crop_size=112,
+               use_bfloat16=True,
+               mode=None,
+               # for centerness learning.
+               has_centerness=False,
+               rpn_center_match_iou_threshold=0.3,
+               rpn_center_unmatched_iou_threshold=0.1,
+               rpn_num_center_samples_per_im=256,
+               # for class manipulation.
+               class_agnostic=False,
+               train_class='all',
+               ):
+    """Initializes parameters for parsing annotations in the dataset.
+
+    Args:
+      output_size: `Tensor` or `list` for [height, width] of output image. The
+        output_size should be divided by the largest feature stride 2^max_level.
+      min_level: `int` number of minimum level of the output feature pyramid.
+      max_level: `int` number of maximum level of the output feature pyramid.
+      num_scales: `int` number representing intermediate scales added
+        on each level. For instances, num_scales=2 adds one additional
+        intermediate anchor scales [2^0, 2^0.5] on each level.
+      aspect_ratios: `list` of float numbers representing the aspect raito
+        anchors added on each level. The number indicates the ratio of width to
+        height. For instances, aspect_ratios=[1.0, 2.0, 0.5] adds three anchors
+        on each scale level.
+      anchor_size: `float` number representing the scale of size of the base
+        anchor to the feature stride 2^level.
+      rpn_match_threshold:
+      rpn_unmatched_threshold:
+      rpn_batch_size_per_im:
+      rpn_fg_fraction:
+      aug_rand_hflip: `bool`, if True, augment training with random
+        horizontal flip.
+      aug_scale_min: `float`, the minimum scale applied to `output_size` for
+        data augmentation during training.
+      aug_scale_max: `float`, the maximum scale applied to `output_size` for
+        data augmentation during training.
+      skip_crowd_during_training: `bool`, if True, skip annotations labeled with
+        `is_crowd` equals to 1.
+      max_num_instances: `int` number of maximum number of instances in an
+        image. The groundtruth data will be padded to `max_num_instances`.
+      include_mask: a bool to indicate whether parse mask groundtruth.
+      mask_crop_size: the size which groundtruth mask is cropped to.
+      use_bfloat16: `bool`, if True, cast output image to tf.bfloat16.
+      mode: a ModeKeys. Specifies if this is training, evaluation, prediction
+        or prediction with groundtruths in the outputs.
+      has_centerness: whether to create centerness targets
+      rpn_center_match_iou_threshold: iou threshold for valid centerness samples
+        ,set to 0.3 by default.
+      rpn_center_unmatched_iou_threshold: iou threshold for invalid centerness
+        samples, set to 0.1 by default.
+      rpn_num_center_samples_per_im: number of centerness samples per image,
+        256 by default.
+      class_agnostic: whether to merge class ids into one foreground(=1) class,
+        False by default.
+      train_class: 'all' or 'voc' or 'nonvoc', 'all' by default.
+    """
+    super(Parser, self).__init__(
+        output_size=output_size,
+        min_level=min_level,
+        max_level=max_level,
+        num_scales=num_scales,
+        aspect_ratios=aspect_ratios,
+        anchor_size=anchor_size,
+        rpn_match_threshold=rpn_match_threshold,
+        rpn_unmatched_threshold=rpn_unmatched_threshold,
+        rpn_batch_size_per_im=rpn_batch_size_per_im,
+        rpn_fg_fraction=rpn_fg_fraction,
+        aug_rand_hflip=aug_rand_hflip,
+        aug_scale_min=aug_scale_min,
+        aug_scale_max=aug_scale_max,
+        skip_crowd_during_training=skip_crowd_during_training,
+        max_num_instances=max_num_instances,
+        include_mask=include_mask,
+        mask_crop_size=mask_crop_size,
+        use_bfloat16=use_bfloat16,
+        mode=mode,)
+
+    # Centerness target assigning.
+    self._has_centerness = has_centerness
+    self._rpn_center_match_iou_threshold = rpn_center_match_iou_threshold
+    self._rpn_center_unmatched_iou_threshold = (
+        rpn_center_unmatched_iou_threshold)
+    self._rpn_num_center_samples_per_im = rpn_num_center_samples_per_im
+
+    # Class manipulation.
+    self._class_agnostic = class_agnostic
+    self._train_class = train_class
+
+  def _parse_train_data(self, data):
+    """Parses data for training.
+
+    Args:
+      data: the decoded tensor dictionary from TfExampleDecoder.
+
+    Returns:
+      image: image tensor that is preproessed to have normalized value and
+        dimension [output_size[0], output_size[1], 3]
+      labels: a dictionary of tensors used for training. The following describes
+        {key: value} pairs in the dictionary.
+        image_info: a 2D `Tensor` that encodes the information of the image and
+          the applied preprocessing. It is in the format of
+          [[original_height, original_width], [scaled_height, scaled_width],
+        anchor_boxes: ordered dictionary with keys
+          [min_level, min_level+1, ..., max_level]. The values are tensor with
+          shape [height_l, width_l, 4] representing anchor boxes at each level.
+        rpn_score_targets: ordered dictionary with keys
+          [min_level, min_level+1, ..., max_level]. The values are tensor with
+          shape [height_l, width_l, anchors_per_location]. The height_l and
+          width_l represent the dimension of class logits at l-th level.
+        rpn_box_targets: ordered dictionary with keys
+          [min_level, min_level+1, ..., max_level]. The values are tensor with
+          shape [height_l, width_l, anchors_per_location * 4]. The height_l and
+          width_l represent the dimension of bounding box regression output at
+          l-th level.
+        gt_boxes: Groundtruth bounding box annotations. The box is represented
+           in [y1, x1, y2, x2] format. The coordinates are w.r.t the scaled
+           image that is fed to the network. The tennsor is padded with -1 to
+           the fixed dimension [self._max_num_instances, 4].
+        gt_classes: Groundtruth classes annotations. The tennsor is padded
+          with -1 to the fixed dimension [self._max_num_instances].
+        gt_masks: groundtrugh masks cropped by the bounding box and
+          resized to a fixed size determined by mask_crop_size.
+    """
+    classes = data['groundtruth_classes']
+    boxes = data['groundtruth_boxes']
+    if self._include_mask:
+      masks = data['groundtruth_instance_masks']
+
+    is_crowds = data['groundtruth_is_crowd']
+    # Skips annotations with `is_crowd` = True.
+    if self._skip_crowd_during_training and self._is_training:
+      num_groundtruths = tf.shape(classes)[0]
+      with tf.control_dependencies([num_groundtruths, is_crowds]):
+        indices = tf.cond(
+            tf.greater(tf.size(is_crowds), 0),
+            lambda: tf.where(tf.logical_not(is_crowds))[:, 0],
+            lambda: tf.cast(tf.range(num_groundtruths), tf.int64))
+      classes = tf.gather(classes, indices)
+      boxes = tf.gather(boxes, indices)
+      if self._include_mask:
+        masks = tf.gather(masks, indices)
+
+    # Gets original image and its size.
+    image = data['image']
+    image_shape = tf.shape(image)[0:2]
+
+    # Normalizes image with mean and std pixel values.
+    image = input_utils.normalize_image(image)
+
+    # Flips image randomly during training.
+    if self._aug_rand_hflip:
+      if self._include_mask:
+        image, boxes, masks = input_utils.random_horizontal_flip(
+            image, boxes, masks)
+      else:
+        image, boxes = input_utils.random_horizontal_flip(
+            image, boxes)
+
+    # Converts boxes from normalized coordinates to pixel coordinates.
+    # Now the coordinates of boxes are w.r.t. the original image.
+    boxes = box_utils.denormalize_boxes(boxes, image_shape)
+
+    # Resizes and crops image.
+    image, image_info = input_utils.resize_and_crop_image(
+        image,
+        self._output_size,
+        padded_size=input_utils.compute_padded_size(
+            self._output_size, 2 ** self._max_level),
+        aug_scale_min=self._aug_scale_min,
+        aug_scale_max=self._aug_scale_max)
+    image_height, image_width, _ = image.get_shape().as_list()
+
+    # Resizes and crops boxes.
+    # Now the coordinates of boxes are w.r.t the scaled image.
+    image_scale = image_info[2, :]
+    offset = image_info[3, :]
+    boxes = input_utils.resize_and_crop_boxes(
+        boxes, image_scale, image_info[1, :], offset)
+
+    # Filters out ground truth boxes that are all zeros.
+    indices = box_utils.get_non_empty_box_indices(boxes)
+    boxes = tf.gather(boxes, indices)
+    classes = tf.gather(classes, indices)
+    if self._include_mask:
+      masks = tf.gather(masks, indices)
+      # Transfer boxes to the original image space and do normalization.
+      cropped_boxes = boxes + tf.tile(tf.expand_dims(offset, axis=0), [1, 2])
+      cropped_boxes /= tf.tile(tf.expand_dims(image_scale, axis=0), [1, 2])
+      cropped_boxes = box_utils.normalize_boxes(cropped_boxes, image_shape)
+      num_masks = tf.shape(masks)[0]
+      masks = tf.image.crop_and_resize(
+          tf.expand_dims(masks, axis=-1),
+          cropped_boxes,
+          box_indices=tf.range(num_masks, dtype=tf.int32),
+          crop_size=[self._mask_crop_size, self._mask_crop_size],
+          method='bilinear')
+      masks = tf.squeeze(masks, axis=-1)
+
+    # Class manipulation.
+    # Filter out novel split classes from training.
+    if self._train_class != 'all':
+      valid_classes = tf.cast(
+          class_utils.coco_split_class_ids(self._train_class),
+          dtype=classes.dtype)
+      match = tf.reduce_any(tf.equal(
+          tf.expand_dims(valid_classes, 1),
+          tf.expand_dims(classes, 0)), 0)
+      # kill novel split classes and boxes.
+      boxes = tf.gather(boxes, tf.where(match)[:, 0])
+      classes = tf.gather(classes, tf.where(match)[:, 0])
+      if self._include_mask:
+        masks = tf.gather(masks, tf.where(match)[:, 0])
+
+    # Assigns anchor targets.
+    # Note that after the target assignment, box targets are absolute pixel
+    # offsets w.r.t. the scaled image.
+    input_anchor = anchor.Anchor(
+        self._min_level,
+        self._max_level,
+        self._num_scales,
+        self._aspect_ratios,
+        self._anchor_size,
+        (image_height, image_width))
+    anchor_labeler = anchor.OlnAnchorLabeler(
+        input_anchor,
+        self._rpn_match_threshold,
+        self._rpn_unmatched_threshold,
+        self._rpn_batch_size_per_im,
+        self._rpn_fg_fraction,
+        # for centerness target.
+        self._has_centerness,
+        self._rpn_center_match_iou_threshold,
+        self._rpn_center_unmatched_iou_threshold,
+        self._rpn_num_center_samples_per_im,)
+
+    if self._has_centerness:
+      rpn_score_targets, _, rpn_lrtb_targets, rpn_center_targets = (
+          anchor_labeler.label_anchors_lrtb(
+              gt_boxes=boxes,
+              gt_labels=tf.cast(
+                  tf.expand_dims(classes, axis=-1), dtype=tf.float32)))
+    else:
+      rpn_score_targets, rpn_box_targets = anchor_labeler.label_anchors(
+          boxes, tf.cast(tf.expand_dims(classes, axis=-1), dtype=tf.float32))
+      # For base rpn, dummy placeholder for centerness target.
+      rpn_center_targets = rpn_score_targets.copy()
+
+    # If bfloat16 is used, casts input image to tf.bfloat16.
+    if self._use_bfloat16:
+      image = tf.cast(image, dtype=tf.bfloat16)
+
+    inputs = {
+        'image': image,
+        'image_info': image_info,
+    }
+    # Packs labels for model_fn outputs.
+    labels = {
+        'anchor_boxes': input_anchor.multilevel_boxes,
+        'image_info': image_info,
+        'rpn_score_targets': rpn_score_targets,
+        'rpn_box_targets': (rpn_lrtb_targets if self._has_centerness
+                            else rpn_box_targets),
+        'rpn_center_targets': rpn_center_targets,
+    }
+    # If class_agnostic, convert to binary classes.
+    if self._class_agnostic:
+      classes = tf.where(tf.greater(classes, 0),
+                         tf.ones_like(classes),
+                         tf.zeros_like(classes))
+
+    inputs['gt_boxes'] = input_utils.pad_to_fixed_size(boxes,
+                                                       self._max_num_instances,
+                                                       -1)
+    inputs['gt_classes'] = input_utils.pad_to_fixed_size(
+        classes, self._max_num_instances, -1)
+    if self._include_mask:
+      inputs['gt_masks'] = input_utils.pad_to_fixed_size(
+          masks, self._max_num_instances, -1)
+
+    return inputs, labels

official/legacy/detection/dataloader/retinanet_parser.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.
+
+"""Data parser and processing.
+
+Parse image and ground truths in a dataset to training targets and package them
+into (image, labels) tuple for RetinaNet.
+
+T.-Y. Lin, P. Goyal, R. Girshick, K. He,  and P. Dollar
+Focal Loss for Dense Object Detection. arXiv:1708.02002
+"""
+
+import tensorflow as tf
+
+from official.legacy.detection.dataloader import anchor
+from official.legacy.detection.dataloader import mode_keys as ModeKeys
+from official.legacy.detection.dataloader import tf_example_decoder
+from official.legacy.detection.utils import box_utils
+from official.legacy.detection.utils import input_utils
+
+
+def process_source_id(source_id):
+  """Processes source_id to the right format."""
+  if source_id.dtype == tf.string:
+    source_id = tf.cast(tf.strings.to_number(source_id), tf.int32)
+  with tf.control_dependencies([source_id]):
+    source_id = tf.cond(
+        pred=tf.equal(tf.size(input=source_id), 0),
+        true_fn=lambda: tf.cast(tf.constant(-1), tf.int32),
+        false_fn=lambda: tf.identity(source_id))
+  return source_id
+
+
+def pad_groundtruths_to_fixed_size(gt, n):
+  """Pads the first dimension of groundtruths labels to the fixed size."""
+  gt['boxes'] = input_utils.pad_to_fixed_size(gt['boxes'], n, -1)
+  gt['is_crowds'] = input_utils.pad_to_fixed_size(gt['is_crowds'], n, 0)
+  gt['areas'] = input_utils.pad_to_fixed_size(gt['areas'], n, -1)
+  gt['classes'] = input_utils.pad_to_fixed_size(gt['classes'], n, -1)
+  return gt
+
+
+class Parser(object):
+  """Parser to parse an image and its annotations into a dictionary of tensors."""
+
+  def __init__(self,
+               output_size,
+               min_level,
+               max_level,
+               num_scales,
+               aspect_ratios,
+               anchor_size,
+               match_threshold=0.5,
+               unmatched_threshold=0.5,
+               aug_rand_hflip=False,
+               aug_scale_min=1.0,
+               aug_scale_max=1.0,
+               use_autoaugment=False,
+               autoaugment_policy_name='v0',
+               skip_crowd_during_training=True,
+               max_num_instances=100,
+               use_bfloat16=True,
+               mode=None):
+    """Initializes parameters for parsing annotations in the dataset.
+
+    Args:
+      output_size: `Tensor` or `list` for [height, width] of output image. The
+        output_size should be divided by the largest feature stride 2^max_level.
+      min_level: `int` number of minimum level of the output feature pyramid.
+      max_level: `int` number of maximum level of the output feature pyramid.
+      num_scales: `int` number representing intermediate scales added on each
+        level. For instances, num_scales=2 adds one additional intermediate
+        anchor scales [2^0, 2^0.5] on each level.
+      aspect_ratios: `list` of float numbers representing the aspect raito
+        anchors added on each level. The number indicates the ratio of width to
+        height. For instances, aspect_ratios=[1.0, 2.0, 0.5] adds three anchors
+        on each scale level.
+      anchor_size: `float` number representing the scale of size of the base
+        anchor to the feature stride 2^level.
+      match_threshold: `float` number between 0 and 1 representing the
+        lower-bound threshold to assign positive labels for anchors. An anchor
+        with a score over the threshold is labeled positive.
+      unmatched_threshold: `float` number between 0 and 1 representing the
+        upper-bound threshold to assign negative labels for anchors. An anchor
+        with a score below the threshold is labeled negative.
+      aug_rand_hflip: `bool`, if True, augment training with random horizontal
+        flip.
+      aug_scale_min: `float`, the minimum scale applied to `output_size` for
+        data augmentation during training.
+      aug_scale_max: `float`, the maximum scale applied to `output_size` for
+        data augmentation during training.
+      use_autoaugment: `bool`, if True, use the AutoAugment augmentation policy
+        during training.
+      autoaugment_policy_name: `string` that specifies the name of the
+        AutoAugment policy that will be used during training.
+      skip_crowd_during_training: `bool`, if True, skip annotations labeled with
+        `is_crowd` equals to 1.
+      max_num_instances: `int` number of maximum number of instances in an
+        image. The groundtruth data will be padded to `max_num_instances`.
+      use_bfloat16: `bool`, if True, cast output image to tf.bfloat16.
+      mode: a ModeKeys. Specifies if this is training, evaluation, prediction or
+        prediction with groundtruths in the outputs.
+    """
+    self._mode = mode
+    self._max_num_instances = max_num_instances
+    self._skip_crowd_during_training = skip_crowd_during_training
+    self._is_training = (mode == ModeKeys.TRAIN)
+
+    self._example_decoder = tf_example_decoder.TfExampleDecoder(
+        include_mask=False)
+
+    # Anchor.
+    self._output_size = output_size
+    self._min_level = min_level
+    self._max_level = max_level
+    self._num_scales = num_scales
+    self._aspect_ratios = aspect_ratios
+    self._anchor_size = anchor_size
+    self._match_threshold = match_threshold
+    self._unmatched_threshold = unmatched_threshold
+
+    # Data augmentation.
+    self._aug_rand_hflip = aug_rand_hflip
+    self._aug_scale_min = aug_scale_min
+    self._aug_scale_max = aug_scale_max
+
+    # Data Augmentation with AutoAugment.
+    self._use_autoaugment = use_autoaugment
+    self._autoaugment_policy_name = autoaugment_policy_name
+
+    # Device.
+    self._use_bfloat16 = use_bfloat16
+
+    # Data is parsed depending on the model Modekey.
+    if mode == ModeKeys.TRAIN:
+      self._parse_fn = self._parse_train_data
+    elif mode == ModeKeys.EVAL:
+      self._parse_fn = self._parse_eval_data
+    elif mode == ModeKeys.PREDICT or mode == ModeKeys.PREDICT_WITH_GT:
+      self._parse_fn = self._parse_predict_data
+    else:
+      raise ValueError('mode is not defined.')
+
+  def __call__(self, value):
+    """Parses data to an image and associated training labels.
+
+    Args:
+      value: a string tensor holding a serialized tf.Example proto.
+
+    Returns:
+      image: image tensor that is preproessed to have normalized value and
+        dimension [output_size[0], output_size[1], 3]
+      labels:
+        cls_targets: ordered dictionary with keys
+          [min_level, min_level+1, ..., max_level]. The values are tensor with
+          shape [height_l, width_l, anchors_per_location]. The height_l and
+          width_l represent the dimension of class logits at l-th level.
+        box_targets: ordered dictionary with keys
+          [min_level, min_level+1, ..., max_level]. The values are tensor with
+          shape [height_l, width_l, anchors_per_location * 4]. The height_l and
+          width_l represent the dimension of bounding box regression output at
+          l-th level.
+        num_positives: number of positive anchors in the image.
+        anchor_boxes: ordered dictionary with keys
+          [min_level, min_level+1, ..., max_level]. The values are tensor with
+          shape [height_l, width_l, 4] representing anchor boxes at each level.
+        image_info: a 2D `Tensor` that encodes the information of the image and
+          the applied preprocessing. It is in the format of
+          [[original_height, original_width], [scaled_height, scaled_width],
+           [y_scale, x_scale], [y_offset, x_offset]].
+        groundtruths:
+          source_id: source image id. Default value -1 if the source id is empty
+            in the groundtruth annotation.
+          boxes: groundtruth bounding box annotations. The box is represented in
+            [y1, x1, y2, x2] format. The tennsor is padded with -1 to the fixed
+            dimension [self._max_num_instances, 4].
+          classes: groundtruth classes annotations. The tennsor is padded with
+            -1 to the fixed dimension [self._max_num_instances].
+          areas: groundtruth areas annotations. The tennsor is padded with -1
+            to the fixed dimension [self._max_num_instances].
+          is_crowds: groundtruth annotations to indicate if an annotation
+            represents a group of instances by value {0, 1}. The tennsor is
+            padded with 0 to the fixed dimension [self._max_num_instances].
+    """
+    with tf.name_scope('parser'):
+      data = self._example_decoder.decode(value)
+      return self._parse_fn(data)
+
+  def _parse_train_data(self, data):
+    """Parses data for training and evaluation."""
+    classes = data['groundtruth_classes']
+    boxes = data['groundtruth_boxes']
+    is_crowds = data['groundtruth_is_crowd']
+    # Skips annotations with `is_crowd` = True.
+    if self._skip_crowd_during_training and self._is_training:
+      num_groundtrtuhs = tf.shape(input=classes)[0]
+      with tf.control_dependencies([num_groundtrtuhs, is_crowds]):
+        indices = tf.cond(
+            pred=tf.greater(tf.size(input=is_crowds), 0),
+            true_fn=lambda: tf.where(tf.logical_not(is_crowds))[:, 0],
+            false_fn=lambda: tf.cast(tf.range(num_groundtrtuhs), tf.int64))
+      classes = tf.gather(classes, indices)
+      boxes = tf.gather(boxes, indices)
+
+    # Gets original image and its size.
+    image = data['image']
+
+    image_shape = tf.shape(input=image)[0:2]
+
+    # Normalizes image with mean and std pixel values.
+    image = input_utils.normalize_image(image)
+
+    # Flips image randomly during training.
+    if self._aug_rand_hflip:
+      image, boxes = input_utils.random_horizontal_flip(image, boxes)
+
+    # Converts boxes from normalized coordinates to pixel coordinates.
+    boxes = box_utils.denormalize_boxes(boxes, image_shape)
+
+    # Resizes and crops image.
+    image, image_info = input_utils.resize_and_crop_image(
+        image,
+        self._output_size,
+        padded_size=input_utils.compute_padded_size(self._output_size,
+                                                    2**self._max_level),
+        aug_scale_min=self._aug_scale_min,
+        aug_scale_max=self._aug_scale_max)
+    image_height, image_width, _ = image.get_shape().as_list()
+
+    # Resizes and crops boxes.
+    image_scale = image_info[2, :]
+    offset = image_info[3, :]
+    boxes = input_utils.resize_and_crop_boxes(boxes, image_scale,
+                                              image_info[1, :], offset)
+    # Filters out ground truth boxes that are all zeros.
+    indices = box_utils.get_non_empty_box_indices(boxes)
+    boxes = tf.gather(boxes, indices)
+    classes = tf.gather(classes, indices)
+
+    # Assigns anchors.
+    input_anchor = anchor.Anchor(self._min_level, self._max_level,
+                                 self._num_scales, self._aspect_ratios,
+                                 self._anchor_size, (image_height, image_width))
+    anchor_labeler = anchor.AnchorLabeler(input_anchor, self._match_threshold,
+                                          self._unmatched_threshold)
+    (cls_targets, box_targets, num_positives) = anchor_labeler.label_anchors(
+        boxes, tf.cast(tf.expand_dims(classes, axis=1), tf.float32))
+
+    # If bfloat16 is used, casts input image to tf.bfloat16.
+    if self._use_bfloat16:
+      image = tf.cast(image, dtype=tf.bfloat16)
+
+    # Packs labels for model_fn outputs.
+    labels = {
+        'cls_targets': cls_targets,
+        'box_targets': box_targets,
+        'anchor_boxes': input_anchor.multilevel_boxes,
+        'num_positives': num_positives,
+        'image_info': image_info,
+    }
+    return image, labels
+
+  def _parse_eval_data(self, data):
+    """Parses data for training and evaluation."""
+    groundtruths = {}
+    classes = data['groundtruth_classes']
+    boxes = data['groundtruth_boxes']
+
+    # Gets original image and its size.
+    image = data['image']
+    image_shape = tf.shape(input=image)[0:2]
+
+    # Normalizes image with mean and std pixel values.
+    image = input_utils.normalize_image(image)
+
+    # Converts boxes from normalized coordinates to pixel coordinates.
+    boxes = box_utils.denormalize_boxes(boxes, image_shape)
+
+    # Resizes and crops image.
+    image, image_info = input_utils.resize_and_crop_image(
+        image,
+        self._output_size,
+        padded_size=input_utils.compute_padded_size(self._output_size,
+                                                    2**self._max_level),
+        aug_scale_min=1.0,
+        aug_scale_max=1.0)
+    image_height, image_width, _ = image.get_shape().as_list()
+
+    # Resizes and crops boxes.
+    image_scale = image_info[2, :]
+    offset = image_info[3, :]
+    boxes = input_utils.resize_and_crop_boxes(boxes, image_scale,
+                                              image_info[1, :], offset)
+    # Filters out ground truth boxes that are all zeros.
+    indices = box_utils.get_non_empty_box_indices(boxes)
+    boxes = tf.gather(boxes, indices)
+    classes = tf.gather(classes, indices)
+
+    # Assigns anchors.
+    input_anchor = anchor.Anchor(self._min_level, self._max_level,
+                                 self._num_scales, self._aspect_ratios,
+                                 self._anchor_size, (image_height, image_width))
+    anchor_labeler = anchor.AnchorLabeler(input_anchor, self._match_threshold,
+                                          self._unmatched_threshold)
+    (cls_targets, box_targets, num_positives) = anchor_labeler.label_anchors(
+        boxes, tf.cast(tf.expand_dims(classes, axis=1), tf.float32))
+
+    # If bfloat16 is used, casts input image to tf.bfloat16.
+    if self._use_bfloat16:
+      image = tf.cast(image, dtype=tf.bfloat16)
+
+    # Sets up groundtruth data for evaluation.
+    groundtruths = {
+        'source_id':
+            data['source_id'],
+        'num_groundtrtuhs':
+            tf.shape(data['groundtruth_classes']),
+        'image_info':
+            image_info,
+        'boxes':
+            box_utils.denormalize_boxes(data['groundtruth_boxes'], image_shape),
+        'classes':
+            data['groundtruth_classes'],
+        'areas':
+            data['groundtruth_area'],
+        'is_crowds':
+            tf.cast(data['groundtruth_is_crowd'], tf.int32),
+    }
+    groundtruths['source_id'] = process_source_id(groundtruths['source_id'])
+    groundtruths = pad_groundtruths_to_fixed_size(groundtruths,
+                                                  self._max_num_instances)
+
+    # Packs labels for model_fn outputs.
+    labels = {
+        'cls_targets': cls_targets,
+        'box_targets': box_targets,
+        'anchor_boxes': input_anchor.multilevel_boxes,
+        'num_positives': num_positives,
+        'image_info': image_info,
+        'groundtruths': groundtruths,
+    }
+    return image, labels
+
+  def _parse_predict_data(self, data):
+    """Parses data for prediction."""
+    # Gets original image and its size.
+    image = data['image']
+    image_shape = tf.shape(input=image)[0:2]
+
+    # Normalizes image with mean and std pixel values.
+    image = input_utils.normalize_image(image)
+
+    # Resizes and crops image.
+    image, image_info = input_utils.resize_and_crop_image(
+        image,
+        self._output_size,
+        padded_size=input_utils.compute_padded_size(self._output_size,
+                                                    2**self._max_level),
+        aug_scale_min=1.0,
+        aug_scale_max=1.0)
+    image_height, image_width, _ = image.get_shape().as_list()
+
+    # If bfloat16 is used, casts input image to tf.bfloat16.
+    if self._use_bfloat16:
+      image = tf.cast(image, dtype=tf.bfloat16)
+
+    # Compute Anchor boxes.
+    input_anchor = anchor.Anchor(self._min_level, self._max_level,
+                                 self._num_scales, self._aspect_ratios,
+                                 self._anchor_size, (image_height, image_width))
+
+    labels = {
+        'anchor_boxes': input_anchor.multilevel_boxes,
+        'image_info': image_info,
+    }
+    # If mode is PREDICT_WITH_GT, returns groundtruths and training targets
+    # in labels.
+    if self._mode == ModeKeys.PREDICT_WITH_GT:
+      # Converts boxes from normalized coordinates to pixel coordinates.
+      boxes = box_utils.denormalize_boxes(data['groundtruth_boxes'],
+                                          image_shape)
+      groundtruths = {
+          'source_id': data['source_id'],
+          'num_detections': tf.shape(data['groundtruth_classes']),
+          'boxes': boxes,
+          'classes': data['groundtruth_classes'],
+          'areas': data['groundtruth_area'],
+          'is_crowds': tf.cast(data['groundtruth_is_crowd'], tf.int32),
+      }
+      groundtruths['source_id'] = process_source_id(groundtruths['source_id'])
+      groundtruths = pad_groundtruths_to_fixed_size(groundtruths,
+                                                    self._max_num_instances)
+      labels['groundtruths'] = groundtruths
+
+      # Computes training objective for evaluation loss.
+      classes = data['groundtruth_classes']
+
+      image_scale = image_info[2, :]
+      offset = image_info[3, :]
+      boxes = input_utils.resize_and_crop_boxes(boxes, image_scale,
+                                                image_info[1, :], offset)
+      # Filters out ground truth boxes that are all zeros.
+      indices = box_utils.get_non_empty_box_indices(boxes)
+      boxes = tf.gather(boxes, indices)
+
+      # Assigns anchors.
+      anchor_labeler = anchor.AnchorLabeler(input_anchor, self._match_threshold,
+                                            self._unmatched_threshold)
+      (cls_targets, box_targets, num_positives) = anchor_labeler.label_anchors(
+          boxes, tf.cast(tf.expand_dims(classes, axis=1), tf.float32))
+      labels['cls_targets'] = cls_targets
+      labels['box_targets'] = box_targets
+      labels['num_positives'] = num_positives
+    return image, labels

official/legacy/detection/dataloader/tf_example_decoder.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.
+
+
+"""Tensorflow Example proto decoder for object detection.
+
+A decoder to decode string tensors containing serialized tensorflow.Example
+protos for object detection.
+"""
+import tensorflow as tf
+
+
+class TfExampleDecoder(object):
+  """Tensorflow Example proto decoder."""
+
+  def __init__(self, include_mask=False):
+    self._include_mask = include_mask
+    self._keys_to_features = {
+        'image/encoded':
+            tf.io.FixedLenFeature((), tf.string),
+        'image/source_id':
+            tf.io.FixedLenFeature((), tf.string),
+        'image/height':
+            tf.io.FixedLenFeature((), tf.int64),
+        'image/width':
+            tf.io.FixedLenFeature((), tf.int64),
+        'image/object/bbox/xmin':
+            tf.io.VarLenFeature(tf.float32),
+        'image/object/bbox/xmax':
+            tf.io.VarLenFeature(tf.float32),
+        'image/object/bbox/ymin':
+            tf.io.VarLenFeature(tf.float32),
+        'image/object/bbox/ymax':
+            tf.io.VarLenFeature(tf.float32),
+        'image/object/class/label':
+            tf.io.VarLenFeature(tf.int64),
+        'image/object/area':
+            tf.io.VarLenFeature(tf.float32),
+        'image/object/is_crowd':
+            tf.io.VarLenFeature(tf.int64),
+    }
+    if include_mask:
+      self._keys_to_features.update({
+          'image/object/mask':
+              tf.io.VarLenFeature(tf.string),
+      })
+
+  def _decode_image(self, parsed_tensors):
+    """Decodes the image and set its static shape."""
+    image = tf.io.decode_image(parsed_tensors['image/encoded'], channels=3)
+    image.set_shape([None, None, 3])
+    return image
+
+  def _decode_boxes(self, parsed_tensors):
+    """Concat box coordinates in the format of [ymin, xmin, ymax, xmax]."""
+    xmin = parsed_tensors['image/object/bbox/xmin']
+    xmax = parsed_tensors['image/object/bbox/xmax']
+    ymin = parsed_tensors['image/object/bbox/ymin']
+    ymax = parsed_tensors['image/object/bbox/ymax']
+    return tf.stack([ymin, xmin, ymax, xmax], axis=-1)
+
+  def _decode_masks(self, parsed_tensors):
+    """Decode a set of PNG masks to the tf.float32 tensors."""
+    def _decode_png_mask(png_bytes):
+      mask = tf.squeeze(
+          tf.io.decode_png(png_bytes, channels=1, dtype=tf.uint8), axis=-1)
+      mask = tf.cast(mask, dtype=tf.float32)
+      mask.set_shape([None, None])
+      return mask
+
+    height = parsed_tensors['image/height']
+    width = parsed_tensors['image/width']
+    masks = parsed_tensors['image/object/mask']
+    return tf.cond(
+        pred=tf.greater(tf.size(input=masks), 0),
+        true_fn=lambda: tf.map_fn(_decode_png_mask, masks, dtype=tf.float32),
+        false_fn=lambda: tf.zeros([0, height, width], dtype=tf.float32))
+
+  def _decode_areas(self, parsed_tensors):
+    xmin = parsed_tensors['image/object/bbox/xmin']
+    xmax = parsed_tensors['image/object/bbox/xmax']
+    ymin = parsed_tensors['image/object/bbox/ymin']
+    ymax = parsed_tensors['image/object/bbox/ymax']
+    return tf.cond(
+        tf.greater(tf.shape(parsed_tensors['image/object/area'])[0], 0),
+        lambda: parsed_tensors['image/object/area'],
+        lambda: (xmax - xmin) * (ymax - ymin))
+
+  def decode(self, serialized_example):
+    """Decode the serialized example.
+
+    Args:
+      serialized_example: a single serialized tf.Example string.
+
+    Returns:
+      decoded_tensors: a dictionary of tensors with the following fields:
+        - image: a uint8 tensor of shape [None, None, 3].
+        - source_id: a string scalar tensor.
+        - height: an integer scalar tensor.
+        - width: an integer scalar tensor.
+        - groundtruth_classes: a int64 tensor of shape [None].
+        - groundtruth_is_crowd: a bool tensor of shape [None].
+        - groundtruth_area: a float32 tensor of shape [None].
+        - groundtruth_boxes: a float32 tensor of shape [None, 4].
+        - groundtruth_instance_masks: a float32 tensor of shape
+            [None, None, None].
+        - groundtruth_instance_masks_png: a string tensor of shape [None].
+    """
+    parsed_tensors = tf.io.parse_single_example(
+        serialized=serialized_example, features=self._keys_to_features)
+    for k in parsed_tensors:
+      if isinstance(parsed_tensors[k], tf.SparseTensor):
+        if parsed_tensors[k].dtype == tf.string:
+          parsed_tensors[k] = tf.sparse.to_dense(
+              parsed_tensors[k], default_value='')
+        else:
+          parsed_tensors[k] = tf.sparse.to_dense(
+              parsed_tensors[k], default_value=0)
+
+    image = self._decode_image(parsed_tensors)
+    boxes = self._decode_boxes(parsed_tensors)
+    areas = self._decode_areas(parsed_tensors)
+    is_crowds = tf.cond(
+        tf.greater(tf.shape(parsed_tensors['image/object/is_crowd'])[0], 0),
+        lambda: tf.cast(parsed_tensors['image/object/is_crowd'], dtype=tf.bool),
+        lambda: tf.zeros_like(parsed_tensors['image/object/class/label'], dtype=tf.bool))  # pylint: disable=line-too-long
+    if self._include_mask:
+      masks = self._decode_masks(parsed_tensors)
+
+    decoded_tensors = {
+        'image': image,
+        'source_id': parsed_tensors['image/source_id'],
+        'height': parsed_tensors['image/height'],
+        'width': parsed_tensors['image/width'],
+        'groundtruth_classes': parsed_tensors['image/object/class/label'],
+        'groundtruth_is_crowd': is_crowds,
+        'groundtruth_area': areas,
+        'groundtruth_boxes': boxes,
+    }
+    if self._include_mask:
+      decoded_tensors.update({
+          'groundtruth_instance_masks': masks,
+          'groundtruth_instance_masks_png': parsed_tensors['image/object/mask'],
+      })
+    return decoded_tensors

official/legacy/detection/evaluation/__init__.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.
+