inputs.py 27.1 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24
# Copyright 2017 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 input function for tf-learn object detection model."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import functools

import tensorflow as tf
from object_detection.builders import dataset_builder
25 26
from object_detection.builders import image_resizer_builder
from object_detection.builders import model_builder
27
from object_detection.builders import preprocessor_builder
28
from object_detection.core import preprocessor
29 30 31 32
from object_detection.core import standard_fields as fields
from object_detection.data_decoders import tf_example_decoder
from object_detection.protos import eval_pb2
from object_detection.protos import input_reader_pb2
33
from object_detection.protos import model_pb2
34
from object_detection.protos import train_pb2
35
from object_detection.utils import config_util
36
from object_detection.utils import ops as util_ops
37
from object_detection.utils import shape_utils
38

39 40
HASH_KEY = 'hash'
HASH_BINS = 1 << 31
41 42
SERVING_FED_EXAMPLE_KEY = 'serialized_example'

43 44 45 46 47
# A map of names to methods that help build the input pipeline.
INPUT_BUILDER_UTIL_MAP = {
    'dataset_build': dataset_builder.build,
}

48

49 50 51 52 53 54
def transform_input_data(tensor_dict,
                         model_preprocess_fn,
                         image_resizer_fn,
                         num_classes,
                         data_augmentation_fn=None,
                         merge_multiple_boxes=False,
55 56
                         retain_original_image=False,
                         use_bfloat16=False):
57 58 59
  """A single function that is responsible for all input data transformations.

  Data transformation functions are applied in the following order.
60 61 62 63 64 65
  1. If key fields.InputDataFields.image_additional_channels is present in
     tensor_dict, the additional channels will be merged into
     fields.InputDataFields.image.
  2. data_augmentation_fn (optional): applied on tensor_dict.
  3. model_preprocess_fn: applied only on image tensor in tensor_dict.
  4. image_resizer_fn: applied on original image and instance mask tensor in
66
     tensor_dict.
67 68
  5. one_hot_encoding: applied to classes tensor in tensor_dict.
  6. merge_multiple_boxes (optional): when groundtruth boxes are exactly the
69 70 71 72 73 74 75 76 77
     same they can be merged into a single box with an associated k-hot class
     label.

  Args:
    tensor_dict: dictionary containing input tensors keyed by
      fields.InputDataFields.
    model_preprocess_fn: model's preprocess function to apply on image tensor.
      This function must take in a 4-D float tensor and return a 4-D preprocess
      float tensor and a tensor containing the true image shape.
78 79 80 81
    image_resizer_fn: image resizer function to apply on groundtruth instance
      `masks. This function must take a 3-D float tensor of an image and a 3-D
      tensor of instance masks and return a resized version of these along with
      the true shapes.
82 83 84 85 86 87 88 89
    num_classes: number of max classes to one-hot (or k-hot) encode the class
      labels.
    data_augmentation_fn: (optional) data augmentation function to apply on
      input `tensor_dict`.
    merge_multiple_boxes: (optional) whether to merge multiple groundtruth boxes
      and classes for a given image if the boxes are exactly the same.
    retain_original_image: (optional) whether to retain original image in the
      output dictionary.
90
    use_bfloat16: (optional) a bool, whether to use bfloat16 in training.
91 92 93 94 95

  Returns:
    A dictionary keyed by fields.InputDataFields containing the tensors obtained
    after applying all the transformations.
  """
96 97 98
  if fields.InputDataFields.groundtruth_boxes in tensor_dict:
    tensor_dict = util_ops.filter_groundtruth_with_nan_box_coordinates(
        tensor_dict)
99 100 101 102 103
  if fields.InputDataFields.image_additional_channels in tensor_dict:
    channels = tensor_dict[fields.InputDataFields.image_additional_channels]
    tensor_dict[fields.InputDataFields.image] = tf.concat(
        [tensor_dict[fields.InputDataFields.image], channels], axis=2)

104
  if retain_original_image:
105
    tensor_dict[fields.InputDataFields.original_image] = tf.cast(
P
pkulzc 已提交
106 107
        image_resizer_fn(tensor_dict[fields.InputDataFields.image], None)[0],
        tf.uint8)
108 109 110 111 112 113

  # Apply data augmentation ops.
  if data_augmentation_fn is not None:
    tensor_dict = data_augmentation_fn(tensor_dict)

  # Apply model preprocessing ops and resize instance masks.
114 115 116
  image = tensor_dict[fields.InputDataFields.image]
  preprocessed_resized_image, true_image_shape = model_preprocess_fn(
      tf.expand_dims(tf.to_float(image), axis=0))
117 118 119
  if use_bfloat16:
    preprocessed_resized_image = tf.cast(
        preprocessed_resized_image, tf.bfloat16)
120 121 122 123 124 125 126
  tensor_dict[fields.InputDataFields.image] = tf.squeeze(
      preprocessed_resized_image, axis=0)
  tensor_dict[fields.InputDataFields.true_image_shape] = tf.squeeze(
      true_image_shape, axis=0)
  if fields.InputDataFields.groundtruth_instance_masks in tensor_dict:
    masks = tensor_dict[fields.InputDataFields.groundtruth_instance_masks]
    _, resized_masks, _ = image_resizer_fn(image, masks)
P
pkulzc 已提交
127 128
    if use_bfloat16:
      resized_masks = tf.cast(resized_masks, tf.bfloat16)
129 130 131 132 133 134 135 136 137 138
    tensor_dict[fields.InputDataFields.
                groundtruth_instance_masks] = resized_masks

  # Transform groundtruth classes to one hot encodings.
  label_offset = 1
  zero_indexed_groundtruth_classes = tensor_dict[
      fields.InputDataFields.groundtruth_classes] - label_offset
  tensor_dict[fields.InputDataFields.groundtruth_classes] = tf.one_hot(
      zero_indexed_groundtruth_classes, num_classes)

139 140 141
  if fields.InputDataFields.groundtruth_confidences in tensor_dict:
    groundtruth_confidences = tensor_dict[
        fields.InputDataFields.groundtruth_confidences]
142
    # Map the confidences to the one-hot encoding of classes
143
    tensor_dict[fields.InputDataFields.groundtruth_confidences] = (
144 145
        tf.reshape(groundtruth_confidences, [-1, 1]) *
        tensor_dict[fields.InputDataFields.groundtruth_classes])
146 147 148 149 150 151
  else:
    groundtruth_confidences = tf.ones_like(
        zero_indexed_groundtruth_classes, dtype=tf.float32)
    tensor_dict[fields.InputDataFields.groundtruth_confidences] = (
        tensor_dict[fields.InputDataFields.groundtruth_classes])

152
  if merge_multiple_boxes:
153 154 155 156 157 158
    merged_boxes, merged_classes, merged_confidences, _ = (
        util_ops.merge_boxes_with_multiple_labels(
            tensor_dict[fields.InputDataFields.groundtruth_boxes],
            zero_indexed_groundtruth_classes,
            groundtruth_confidences,
            num_classes))
159
    merged_classes = tf.cast(merged_classes, tf.float32)
160 161
    tensor_dict[fields.InputDataFields.groundtruth_boxes] = merged_boxes
    tensor_dict[fields.InputDataFields.groundtruth_classes] = merged_classes
162 163
    tensor_dict[fields.InputDataFields.groundtruth_confidences] = (
        merged_confidences)
P
pkulzc 已提交
164 165 166
  if fields.InputDataFields.groundtruth_boxes in tensor_dict:
    tensor_dict[fields.InputDataFields.num_groundtruth_boxes] = tf.shape(
        tensor_dict[fields.InputDataFields.groundtruth_boxes])[0]
167 168 169 170

  return tensor_dict


171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200
def pad_input_data_to_static_shapes(tensor_dict, max_num_boxes, num_classes,
                                    spatial_image_shape=None):
  """Pads input tensors to static shapes.

  Args:
    tensor_dict: Tensor dictionary of input data
    max_num_boxes: Max number of groundtruth boxes needed to compute shapes for
      padding.
    num_classes: Number of classes in the dataset needed to compute shapes for
      padding.
    spatial_image_shape: A list of two integers of the form [height, width]
      containing expected spatial shape of the image.

  Returns:
    A dictionary keyed by fields.InputDataFields containing padding shapes for
    tensors in the dataset.

  Raises:
    ValueError: If groundtruth classes is neither rank 1 nor rank 2.
  """

  if not spatial_image_shape or spatial_image_shape == [-1, -1]:
    height, width = None, None
  else:
    height, width = spatial_image_shape  # pylint: disable=unpacking-non-sequence

  num_additional_channels = 0
  if fields.InputDataFields.image_additional_channels in tensor_dict:
    num_additional_channels = tensor_dict[
        fields.InputDataFields.image_additional_channels].shape[2].value
201 202 203 204
  num_image_channels = 3
  if fields.InputDataFields.image in tensor_dict:
    num_image_channels = tensor_dict[fields.InputDataFields
                                     .image].shape[2].value
205 206 207
  padding_shapes = {
      # Additional channels are merged before batching.
      fields.InputDataFields.image: [
208
          height, width, num_image_channels + num_additional_channels
209
      ],
P
pkulzc 已提交
210
      fields.InputDataFields.original_image_spatial_shape: [2],
211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226
      fields.InputDataFields.image_additional_channels: [
          height, width, num_additional_channels
      ],
      fields.InputDataFields.source_id: [],
      fields.InputDataFields.filename: [],
      fields.InputDataFields.key: [],
      fields.InputDataFields.groundtruth_difficult: [max_num_boxes],
      fields.InputDataFields.groundtruth_boxes: [max_num_boxes, 4],
      fields.InputDataFields.groundtruth_classes: [max_num_boxes, num_classes],
      fields.InputDataFields.groundtruth_instance_masks: [
          max_num_boxes, height, width
      ],
      fields.InputDataFields.groundtruth_is_crowd: [max_num_boxes],
      fields.InputDataFields.groundtruth_group_of: [max_num_boxes],
      fields.InputDataFields.groundtruth_area: [max_num_boxes],
      fields.InputDataFields.groundtruth_weights: [max_num_boxes],
227 228 229
      fields.InputDataFields.groundtruth_confidences: [
          max_num_boxes, num_classes
      ],
230 231
      fields.InputDataFields.num_groundtruth_boxes: [],
      fields.InputDataFields.groundtruth_label_types: [max_num_boxes],
232
      fields.InputDataFields.groundtruth_label_weights: [max_num_boxes],
233 234 235 236 237
      fields.InputDataFields.true_image_shape: [3],
      fields.InputDataFields.multiclass_scores: [
          max_num_boxes, num_classes + 1 if num_classes is not None else None
      ],
      fields.InputDataFields.groundtruth_image_classes: [num_classes],
238
      fields.InputDataFields.groundtruth_image_confidences: [num_classes],
239 240 241 242
  }

  if fields.InputDataFields.original_image in tensor_dict:
    padding_shapes[fields.InputDataFields.original_image] = [
243
        height, width, num_image_channels + num_additional_channels
244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259
    ]
  if fields.InputDataFields.groundtruth_keypoints in tensor_dict:
    tensor_shape = (
        tensor_dict[fields.InputDataFields.groundtruth_keypoints].shape)
    padding_shape = [max_num_boxes, tensor_shape[1].value,
                     tensor_shape[2].value]
    padding_shapes[fields.InputDataFields.groundtruth_keypoints] = padding_shape
  if fields.InputDataFields.groundtruth_keypoint_visibilities in tensor_dict:
    tensor_shape = tensor_dict[fields.InputDataFields.
                               groundtruth_keypoint_visibilities].shape
    padding_shape = [max_num_boxes, tensor_shape[1].value]
    padding_shapes[fields.InputDataFields.
                   groundtruth_keypoint_visibilities] = padding_shape

  padded_tensor_dict = {}
  for tensor_name in tensor_dict:
260 261
    padded_tensor_dict[tensor_name] = shape_utils.pad_or_clip_nd(
        tensor_dict[tensor_name], padding_shapes[tensor_name])
262 263 264 265 266 267 268 269

  # Make sure that the number of groundtruth boxes now reflects the
  # padded/clipped tensors.
  if fields.InputDataFields.num_groundtruth_boxes in padded_tensor_dict:
    padded_tensor_dict[fields.InputDataFields.num_groundtruth_boxes] = (
        tf.minimum(
            padded_tensor_dict[fields.InputDataFields.num_groundtruth_boxes],
            max_num_boxes))
270 271 272
  return padded_tensor_dict


273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292
def augment_input_data(tensor_dict, data_augmentation_options):
  """Applies data augmentation ops to input tensors.

  Args:
    tensor_dict: A dictionary of input tensors keyed by fields.InputDataFields.
    data_augmentation_options: A list of tuples, where each tuple contains a
      function and a dictionary that contains arguments and their values.
      Usually, this is the output of core/preprocessor.build.

  Returns:
    A dictionary of tensors obtained by applying data augmentation ops to the
    input tensor dictionary.
  """
  tensor_dict[fields.InputDataFields.image] = tf.expand_dims(
      tf.to_float(tensor_dict[fields.InputDataFields.image]), 0)

  include_instance_masks = (fields.InputDataFields.groundtruth_instance_masks
                            in tensor_dict)
  include_keypoints = (fields.InputDataFields.groundtruth_keypoints
                       in tensor_dict)
293 294 295 296
  include_label_weights = (fields.InputDataFields.groundtruth_weights
                           in tensor_dict)
  include_label_confidences = (fields.InputDataFields.groundtruth_confidences
                               in tensor_dict)
297 298 299
  tensor_dict = preprocessor.preprocess(
      tensor_dict, data_augmentation_options,
      func_arg_map=preprocessor.get_default_func_arg_map(
300 301
          include_label_weights=include_label_weights,
          include_label_confidences=include_label_confidences,
302 303 304 305 306 307 308
          include_instance_masks=include_instance_masks,
          include_keypoints=include_keypoints))
  tensor_dict[fields.InputDataFields.image] = tf.squeeze(
      tensor_dict[fields.InputDataFields.image], axis=0)
  return tensor_dict


309 310 311 312 313 314
def _get_labels_dict(input_dict):
  """Extracts labels dict from input dict."""
  required_label_keys = [
      fields.InputDataFields.num_groundtruth_boxes,
      fields.InputDataFields.groundtruth_boxes,
      fields.InputDataFields.groundtruth_classes,
315
      fields.InputDataFields.groundtruth_weights,
316 317 318 319 320 321
  ]
  labels_dict = {}
  for key in required_label_keys:
    labels_dict[key] = input_dict[key]

  optional_label_keys = [
322
      fields.InputDataFields.groundtruth_confidences,
323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338
      fields.InputDataFields.groundtruth_keypoints,
      fields.InputDataFields.groundtruth_instance_masks,
      fields.InputDataFields.groundtruth_area,
      fields.InputDataFields.groundtruth_is_crowd,
      fields.InputDataFields.groundtruth_difficult
  ]

  for key in optional_label_keys:
    if key in input_dict:
      labels_dict[key] = input_dict[key]
  if fields.InputDataFields.groundtruth_difficult in labels_dict:
    labels_dict[fields.InputDataFields.groundtruth_difficult] = tf.cast(
        labels_dict[fields.InputDataFields.groundtruth_difficult], tf.int32)
  return labels_dict


339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367
def _replace_empty_string_with_random_number(string_tensor):
  """Returns string unchanged if non-empty, and random string tensor otherwise.

  The random string is an integer 0 and 2**63 - 1, casted as string.


  Args:
    string_tensor: A tf.tensor of dtype string.

  Returns:
    out_string: A tf.tensor of dtype string. If string_tensor contains the empty
      string, out_string will contain a random integer casted to a string.
      Otherwise string_tensor is returned unchanged.

  """

  empty_string = tf.constant('', dtype=tf.string, name='EmptyString')

  random_source_id = tf.as_string(
      tf.random_uniform(shape=[], maxval=2**63 - 1, dtype=tf.int64))

  out_string = tf.cond(
      tf.equal(string_tensor, empty_string),
      true_fn=lambda: random_source_id,
      false_fn=lambda: string_tensor)

  return out_string


368 369
def _get_features_dict(input_dict):
  """Extracts features dict from input dict."""
370 371 372 373 374

  source_id = _replace_empty_string_with_random_number(
      input_dict[fields.InputDataFields.source_id])

  hash_from_source_id = tf.string_to_hash_bucket_fast(source_id, HASH_BINS)
375 376 377 378 379
  features = {
      fields.InputDataFields.image:
          input_dict[fields.InputDataFields.image],
      HASH_KEY: tf.cast(hash_from_source_id, tf.int32),
      fields.InputDataFields.true_image_shape:
P
pkulzc 已提交
380 381 382
          input_dict[fields.InputDataFields.true_image_shape],
      fields.InputDataFields.original_image_spatial_shape:
          input_dict[fields.InputDataFields.original_image_spatial_shape]
383 384 385 386 387 388 389
  }
  if fields.InputDataFields.original_image in input_dict:
    features[fields.InputDataFields.original_image] = input_dict[
        fields.InputDataFields.original_image]
  return features


390 391
def create_train_input_fn(train_config, train_input_config,
                          model_config):
392 393 394 395 396
  """Creates a train `input` function for `Estimator`.

  Args:
    train_config: A train_pb2.TrainConfig.
    train_input_config: An input_reader_pb2.InputReader.
397
    model_config: A model_pb2.DetectionModel.
398 399 400 401 402

  Returns:
    `input_fn` for `Estimator` in TRAIN mode.
  """

403
  def _train_input_fn(params=None):
404 405
    """Returns `features` and `labels` tensor dictionaries for training.

406 407 408
    Args:
      params: Parameter dictionary passed from the estimator.

409
    Returns:
410 411
      A tf.data.Dataset that holds (features, labels) tuple.

412
      features: Dictionary of feature tensors.
413 414 415 416 417 418 419
        features[fields.InputDataFields.image] is a [batch_size, H, W, C]
          float32 tensor with preprocessed images.
        features[HASH_KEY] is a [batch_size] int32 tensor representing unique
          identifiers for the images.
        features[fields.InputDataFields.true_image_shape] is a [batch_size, 3]
          int32 tensor representing the true image shapes, as preprocessed
          images could be padded.
420
        features[fields.InputDataFields.original_image] (optional) is a
421
          [batch_size, H, W, C] float32 tensor with original images.
422
      labels: Dictionary of groundtruth tensors.
423 424 425 426 427 428 429 430 431 432 433 434 435 436 437 438 439 440
        labels[fields.InputDataFields.num_groundtruth_boxes] is a [batch_size]
          int32 tensor indicating the number of groundtruth boxes.
        labels[fields.InputDataFields.groundtruth_boxes] is a
          [batch_size, num_boxes, 4] float32 tensor containing the corners of
          the groundtruth boxes.
        labels[fields.InputDataFields.groundtruth_classes] is a
          [batch_size, num_boxes, num_classes] float32 one-hot tensor of
          classes.
        labels[fields.InputDataFields.groundtruth_weights] is a
          [batch_size, num_boxes] float32 tensor containing groundtruth weights
          for the boxes.
        -- Optional --
        labels[fields.InputDataFields.groundtruth_instance_masks] is a
          [batch_size, num_boxes, H, W] float32 tensor containing only binary
          values, which represent instance masks for objects.
        labels[fields.InputDataFields.groundtruth_keypoints] is a
          [batch_size, num_boxes, num_keypoints, 2] float32 tensor containing
          keypoints for each box.
441 442

    Raises:
443 444
      TypeError: if the `train_config`, `train_input_config` or `model_config`
        are not of the correct type.
445 446 447 448 449 450 451
    """
    if not isinstance(train_config, train_pb2.TrainConfig):
      raise TypeError('For training mode, the `train_config` must be a '
                      'train_pb2.TrainConfig.')
    if not isinstance(train_input_config, input_reader_pb2.InputReader):
      raise TypeError('The `train_input_config` must be a '
                      'input_reader_pb2.InputReader.')
452 453 454
    if not isinstance(model_config, model_pb2.DetectionModel):
      raise TypeError('The `model_config` must be a '
                      'model_pb2.DetectionModel.')
455

456 457 458 459 460 461 462 463 464 465 466 467 468 469 470 471 472 473
    def transform_and_pad_input_data_fn(tensor_dict):
      """Combines transform and pad operation."""
      data_augmentation_options = [
          preprocessor_builder.build(step)
          for step in train_config.data_augmentation_options
      ]
      data_augmentation_fn = functools.partial(
          augment_input_data,
          data_augmentation_options=data_augmentation_options)
      model = model_builder.build(model_config, is_training=True)
      image_resizer_config = config_util.get_image_resizer_config(model_config)
      image_resizer_fn = image_resizer_builder.build(image_resizer_config)
      transform_data_fn = functools.partial(
          transform_input_data, model_preprocess_fn=model.preprocess,
          image_resizer_fn=image_resizer_fn,
          num_classes=config_util.get_number_of_classes(model_config),
          data_augmentation_fn=data_augmentation_fn,
          merge_multiple_boxes=train_config.merge_multiple_label_boxes,
474 475
          retain_original_image=train_config.retain_original_images,
          use_bfloat16=train_config.use_bfloat16)
476 477 478 479 480 481 482 483

      tensor_dict = pad_input_data_to_static_shapes(
          tensor_dict=transform_data_fn(tensor_dict),
          max_num_boxes=train_input_config.max_number_of_boxes,
          num_classes=config_util.get_number_of_classes(model_config),
          spatial_image_shape=config_util.get_spatial_image_size(
              image_resizer_config))
      return (_get_features_dict(tensor_dict), _get_labels_dict(tensor_dict))
484

485
    dataset = INPUT_BUILDER_UTIL_MAP['dataset_build'](
486
        train_input_config,
487 488 489
        transform_input_data_fn=transform_and_pad_input_data_fn,
        batch_size=params['batch_size'] if params else train_config.batch_size)
    return dataset
490 491 492 493

  return _train_input_fn


494
def create_eval_input_fn(eval_config, eval_input_config, model_config):
495 496 497 498 499
  """Creates an eval `input` function for `Estimator`.

  Args:
    eval_config: An eval_pb2.EvalConfig.
    eval_input_config: An input_reader_pb2.InputReader.
500
    model_config: A model_pb2.DetectionModel.
501 502 503 504 505

  Returns:
    `input_fn` for `Estimator` in EVAL mode.
  """

506
  def _eval_input_fn(params=None):
507 508
    """Returns `features` and `labels` tensor dictionaries for evaluation.

509 510 511
    Args:
      params: Parameter dictionary passed from the estimator.

512
    Returns:
513 514
      A tf.data.Dataset that holds (features, labels) tuple.

515
      features: Dictionary of feature tensors.
516 517 518 519 520 521 522 523 524
        features[fields.InputDataFields.image] is a [1, H, W, C] float32 tensor
          with preprocessed images.
        features[HASH_KEY] is a [1] int32 tensor representing unique
          identifiers for the images.
        features[fields.InputDataFields.true_image_shape] is a [1, 3]
          int32 tensor representing the true image shapes, as preprocessed
          images could be padded.
        features[fields.InputDataFields.original_image] is a [1, H', W', C]
          float32 tensor with the original image.
525
      labels: Dictionary of groundtruth tensors.
526 527 528 529 530 531 532 533 534 535 536 537 538 539
        labels[fields.InputDataFields.groundtruth_boxes] is a [1, num_boxes, 4]
          float32 tensor containing the corners of the groundtruth boxes.
        labels[fields.InputDataFields.groundtruth_classes] is a
          [num_boxes, num_classes] float32 one-hot tensor of classes.
        labels[fields.InputDataFields.groundtruth_area] is a [1, num_boxes]
          float32 tensor containing object areas.
        labels[fields.InputDataFields.groundtruth_is_crowd] is a [1, num_boxes]
          bool tensor indicating if the boxes enclose a crowd.
        labels[fields.InputDataFields.groundtruth_difficult] is a [1, num_boxes]
          int32 tensor indicating if the boxes represent difficult instances.
        -- Optional --
        labels[fields.InputDataFields.groundtruth_instance_masks] is a
          [1, num_boxes, H, W] float32 tensor containing only binary values,
          which represent instance masks for objects.
540 541

    Raises:
542 543
      TypeError: if the `eval_config`, `eval_input_config` or `model_config`
        are not of the correct type.
544
    """
545
    params = params or {}
546 547
    if not isinstance(eval_config, eval_pb2.EvalConfig):
      raise TypeError('For eval mode, the `eval_config` must be a '
548
                      'train_pb2.EvalConfig.')
549 550 551
    if not isinstance(eval_input_config, input_reader_pb2.InputReader):
      raise TypeError('The `eval_input_config` must be a '
                      'input_reader_pb2.InputReader.')
552 553 554 555
    if not isinstance(model_config, model_pb2.DetectionModel):
      raise TypeError('The `model_config` must be a '
                      'model_pb2.DetectionModel.')

556 557 558 559 560 561 562 563 564 565 566 567 568 569 570 571 572 573 574 575
    def transform_and_pad_input_data_fn(tensor_dict):
      """Combines transform and pad operation."""
      num_classes = config_util.get_number_of_classes(model_config)
      model = model_builder.build(model_config, is_training=False)
      image_resizer_config = config_util.get_image_resizer_config(model_config)
      image_resizer_fn = image_resizer_builder.build(image_resizer_config)

      transform_data_fn = functools.partial(
          transform_input_data, model_preprocess_fn=model.preprocess,
          image_resizer_fn=image_resizer_fn,
          num_classes=num_classes,
          data_augmentation_fn=None,
          retain_original_image=eval_config.retain_original_images)
      tensor_dict = pad_input_data_to_static_shapes(
          tensor_dict=transform_data_fn(tensor_dict),
          max_num_boxes=eval_input_config.max_number_of_boxes,
          num_classes=config_util.get_number_of_classes(model_config),
          spatial_image_shape=config_util.get_spatial_image_size(
              image_resizer_config))
      return (_get_features_dict(tensor_dict), _get_labels_dict(tensor_dict))
576 577
    dataset = INPUT_BUILDER_UTIL_MAP['dataset_build'](
        eval_input_config,
P
pkulzc 已提交
578
        batch_size=params['batch_size'] if params else eval_config.batch_size,
579 580
        transform_input_data_fn=transform_and_pad_input_data_fn)
    return dataset
581 582 583 584

  return _eval_input_fn


585
def create_predict_input_fn(model_config, predict_input_config):
586 587
  """Creates a predict `input` function for `Estimator`.

588 589
  Args:
    model_config: A model_pb2.DetectionModel.
590
    predict_input_config: An input_reader_pb2.InputReader.
591

592 593 594 595
  Returns:
    `input_fn` for `Estimator` in PREDICT mode.
  """

596
  def _predict_input_fn(params=None):
597 598
    """Decodes serialized tf.Examples and returns `ServingInputReceiver`.

599 600 601
    Args:
      params: Parameter dictionary passed from the estimator.

602 603 604
    Returns:
      `ServingInputReceiver`.
    """
605
    del params
606
    example = tf.placeholder(dtype=tf.string, shape=[], name='tf_example')
607

608 609 610 611
    num_classes = config_util.get_number_of_classes(model_config)
    model = model_builder.build(model_config, is_training=False)
    image_resizer_config = config_util.get_image_resizer_config(model_config)
    image_resizer_fn = image_resizer_builder.build(image_resizer_config)
612

613 614 615 616 617 618
    transform_fn = functools.partial(
        transform_input_data, model_preprocess_fn=model.preprocess,
        image_resizer_fn=image_resizer_fn,
        num_classes=num_classes,
        data_augmentation_fn=None)

619 620 621
    decoder = tf_example_decoder.TfExampleDecoder(
        load_instance_masks=False,
        num_additional_channels=predict_input_config.num_additional_channels)
622
    input_dict = transform_fn(decoder.decode(example))
623 624
    images = tf.to_float(input_dict[fields.InputDataFields.image])
    images = tf.expand_dims(images, axis=0)
625 626
    true_image_shape = tf.expand_dims(
        input_dict[fields.InputDataFields.true_image_shape], axis=0)
627 628

    return tf.estimator.export.ServingInputReceiver(
629 630 631
        features={
            fields.InputDataFields.image: images,
            fields.InputDataFields.true_image_shape: true_image_shape},
632 633 634
        receiver_tensors={SERVING_FED_EXAMPLE_KEY: example})

  return _predict_input_fn