inputs.py 53 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
# 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

23
import tensorflow.compat.v1 as tf
24
from tensorflow.compat.v1 import estimator as tf_estimator
25
from object_detection.builders import dataset_builder
26
27
from object_detection.builders import image_resizer_builder
from object_detection.builders import model_builder
28
from object_detection.builders import preprocessor_builder
29
30
from object_detection.core import box_list
from object_detection.core import box_list_ops
31
from object_detection.core import densepose_ops
32
from object_detection.core import keypoint_ops
33
from object_detection.core import preprocessor
34
35
36
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
37
from object_detection.protos import image_resizer_pb2
38
from object_detection.protos import input_reader_pb2
39
from object_detection.protos import model_pb2
40
from object_detection.protos import train_pb2
41
from object_detection.utils import config_util
42
from object_detection.utils import ops as util_ops
43
from object_detection.utils import shape_utils
44

45
46
HASH_KEY = 'hash'
HASH_BINS = 1 << 31
47
SERVING_FED_EXAMPLE_KEY = 'serialized_example'
48
_LABEL_OFFSET = 1
49

50
51
52
# A map of names to methods that help build the input pipeline.
INPUT_BUILDER_UTIL_MAP = {
    'dataset_build': dataset_builder.build,
53
    'model_build': model_builder.build,
54
55
}

56

pkulzc's avatar
pkulzc committed
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
def _multiclass_scores_or_one_hot_labels(multiclass_scores,
                                         groundtruth_boxes,
                                         groundtruth_classes, num_classes):
  """Returns one-hot encoding of classes when multiclass_scores is empty."""
  # Replace groundtruth_classes tensor with multiclass_scores tensor when its
  # non-empty. If multiclass_scores is empty fall back on groundtruth_classes
  # tensor.
  def true_fn():
    return tf.reshape(multiclass_scores,
                      [tf.shape(groundtruth_boxes)[0], num_classes])
  def false_fn():
    return tf.one_hot(groundtruth_classes, num_classes)
  return tf.cond(tf.size(multiclass_scores) > 0, true_fn, false_fn)


Rich Munoz's avatar
Rich Munoz committed
72
73
74
def convert_labeled_classes_to_k_hot(groundtruth_labeled_classes,
                                     num_classes,
                                     map_empty_to_ones=False):
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
  """Returns k-hot encoding of the labeled classes.

  If map_empty_to_ones is enabled and the input labeled_classes is empty,
  this function assumes all classes are exhaustively labeled, thus returning
  an all-one encoding.

  Args:
    groundtruth_labeled_classes: a Tensor holding a sparse representation of
      labeled classes.
    num_classes: an integer representing the number of classes
    map_empty_to_ones: boolean (default: False).  Set this to be True to default
    to an all-ones result if given an empty `groundtruth_labeled_classes`.
  Returns:
    A k-hot (and 0-indexed) tensor representation of
    `groundtruth_labeled_classes`.
  """
91
92
93
94
95
96
97
98
99
100
101
102

  # If the input labeled_classes is empty, it assumes all classes are
  # exhaustively labeled, thus returning an all-one encoding.
  def true_fn():
    return tf.sparse_to_dense(
        groundtruth_labeled_classes - _LABEL_OFFSET, [num_classes],
        tf.constant(1, dtype=tf.float32),
        validate_indices=False)

  def false_fn():
    return tf.ones(num_classes, dtype=tf.float32)

103
104
105
  if map_empty_to_ones:
    return tf.cond(tf.size(groundtruth_labeled_classes) > 0, true_fn, false_fn)
  return true_fn()
106
107
108
109
110


def _remove_unrecognized_classes(class_ids, unrecognized_label):
  """Returns class ids with unrecognized classes filtered out."""

111
112
  recognized_indices = tf.squeeze(
      tf.where(tf.greater(class_ids, unrecognized_label)), -1)
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
  return tf.gather(class_ids, recognized_indices)


def assert_or_prune_invalid_boxes(boxes):
  """Makes sure boxes have valid sizes (ymax >= ymin, xmax >= xmin).

  When the hardware supports assertions, the function raises an error when
  boxes have an invalid size. If assertions are not supported (e.g. on TPU),
  boxes with invalid sizes are filtered out.

  Args:
    boxes: float tensor of shape [num_boxes, 4]

  Returns:
    boxes: float tensor of shape [num_valid_boxes, 4] with invalid boxes
      filtered out.

  Raises:
    tf.errors.InvalidArgumentError: When we detect boxes with invalid size.
      This is not supported on TPUs.
  """

  ymin, xmin, ymax, xmax = tf.split(
      boxes, num_or_size_splits=4, axis=1)

  height_check = tf.Assert(tf.reduce_all(ymax >= ymin), [ymin, ymax])
  width_check = tf.Assert(tf.reduce_all(xmax >= xmin), [xmin, xmax])

  with tf.control_dependencies([height_check, width_check]):
    boxes_tensor = tf.concat([ymin, xmin, ymax, xmax], axis=1)
    boxlist = box_list.BoxList(boxes_tensor)
    # TODO(b/149221748) Remove pruning when XLA supports assertions.
    boxlist = box_list_ops.prune_small_boxes(boxlist, 0)

  return boxlist.get()


150
151
152
153
154
155
def transform_input_data(tensor_dict,
                         model_preprocess_fn,
                         image_resizer_fn,
                         num_classes,
                         data_augmentation_fn=None,
                         merge_multiple_boxes=False,
156
                         retain_original_image=False,
157
                         use_multiclass_scores=False,
158
                         use_bfloat16=False,
159
160
                         retain_original_image_additional_channels=False,
                         keypoint_type_weight=None):
161
162
163
  """A single function that is responsible for all input data transformations.

  Data transformation functions are applied in the following order.
164
165
166
167
168
  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.
169
170
171
172
173
174
  4. keypoint_type_weight (optional): If groundtruth keypoints are in
     the tensor dictionary, per-keypoint weights are produced. These weights are
     initialized by `keypoint_type_weight` (or ones if left None).
     Then, for all keypoints that are not visible, the weights are set to 0 (to
     avoid penalizing the model in a loss function).
  5. image_resizer_fn: applied on original image and instance mask tensor in
175
     tensor_dict.
176
177
  6. one_hot_encoding: applied to classes tensor in tensor_dict.
  7. merge_multiple_boxes (optional): when groundtruth boxes are exactly the
178
179
180
181
182
183
184
185
186
     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.
187
188
189
190
    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.
191
192
193
194
195
196
197
198
    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.
pkulzc's avatar
pkulzc committed
199
200
201
202
    use_multiclass_scores: whether to use multiclass scores as class targets
      instead of one-hot encoding of `groundtruth_classes`. When
      this is True and multiclass_scores is empty, one-hot encoding of
      `groundtruth_classes` is used as a fallback.
203
    use_bfloat16: (optional) a bool, whether to use bfloat16 in training.
204
205
    retain_original_image_additional_channels: (optional) Whether to retain
      original image additional channels in the output dictionary.
206
207
208
    keypoint_type_weight: A list (of length num_keypoints) containing
      groundtruth loss weights to use for each keypoint. If None, will use a
      weight of 1.
209
210
211
212

  Returns:
    A dictionary keyed by fields.InputDataFields containing the tensors obtained
    after applying all the transformations.
213
214
215
216
217

  Raises:
    KeyError: If both groundtruth_labeled_classes and groundtruth_image_classes
      are provided by the decoder in tensor_dict since both fields are
      considered to contain the same information.
218
  """
pkulzc's avatar
pkulzc committed
219
  out_tensor_dict = tensor_dict.copy()
220

221
222
223
224
225
226
  input_fields = fields.InputDataFields
  labeled_classes_field = input_fields.groundtruth_labeled_classes
  image_classes_field = input_fields.groundtruth_image_classes
  verified_neg_classes_field = input_fields.groundtruth_verified_neg_classes
  not_exhaustive_field = input_fields.groundtruth_not_exhaustive_classes

227
228
229
230
231
  if (labeled_classes_field in out_tensor_dict and
      image_classes_field in out_tensor_dict):
    raise KeyError('groundtruth_labeled_classes and groundtruth_image_classes'
                   'are provided by the decoder, but only one should be set.')

232
233
234
235
236
  for field, map_empty_to_ones in [
      (labeled_classes_field, True),
      (image_classes_field, True),
      (verified_neg_classes_field, False),
      (not_exhaustive_field, False)]:
237
238
239
    if field in out_tensor_dict:
      out_tensor_dict[field] = _remove_unrecognized_classes(
          out_tensor_dict[field], unrecognized_label=-1)
Rich Munoz's avatar
Rich Munoz committed
240
      out_tensor_dict[field] = convert_labeled_classes_to_k_hot(
241
          out_tensor_dict[field], num_classes, map_empty_to_ones)
242
243

  if input_fields.multiclass_scores in out_tensor_dict:
pkulzc's avatar
pkulzc committed
244
    out_tensor_dict[
245
        input_fields
pkulzc's avatar
pkulzc committed
246
        .multiclass_scores] = _multiclass_scores_or_one_hot_labels(
247
248
249
            out_tensor_dict[input_fields.multiclass_scores],
            out_tensor_dict[input_fields.groundtruth_boxes],
            out_tensor_dict[input_fields.groundtruth_classes],
pkulzc's avatar
pkulzc committed
250
251
            num_classes)

252
  if input_fields.groundtruth_boxes in out_tensor_dict:
pkulzc's avatar
pkulzc committed
253
254
255
    out_tensor_dict = util_ops.filter_groundtruth_with_nan_box_coordinates(
        out_tensor_dict)
    out_tensor_dict = util_ops.filter_unrecognized_classes(out_tensor_dict)
256

257
  if retain_original_image:
258
259
    out_tensor_dict[input_fields.original_image] = tf.cast(
        image_resizer_fn(out_tensor_dict[input_fields.image],
pkulzc's avatar
pkulzc committed
260
                         None)[0], tf.uint8)
261

262
263
264
265
  if input_fields.image_additional_channels in out_tensor_dict:
    channels = out_tensor_dict[input_fields.image_additional_channels]
    out_tensor_dict[input_fields.image] = tf.concat(
        [out_tensor_dict[input_fields.image], channels], axis=2)
266
267
    if retain_original_image_additional_channels:
      out_tensor_dict[
268
          input_fields.image_additional_channels] = tf.cast(
269
              image_resizer_fn(channels, None)[0], tf.uint8)
270

271
272
  # Apply data augmentation ops.
  if data_augmentation_fn is not None:
pkulzc's avatar
pkulzc committed
273
    out_tensor_dict = data_augmentation_fn(out_tensor_dict)
274
275

  # Apply model preprocessing ops and resize instance masks.
276
  image = out_tensor_dict[input_fields.image]
277
  preprocessed_resized_image, true_image_shape = model_preprocess_fn(
278
      tf.expand_dims(tf.cast(image, dtype=tf.float32), axis=0))
279
280
281
282
283
284
285
286
287
288

  preprocessed_shape = tf.shape(preprocessed_resized_image)
  new_height, new_width = preprocessed_shape[1], preprocessed_shape[2]

  im_box = tf.stack([
      0.0, 0.0,
      tf.to_float(new_height) / tf.to_float(true_image_shape[0, 0]),
      tf.to_float(new_width) / tf.to_float(true_image_shape[0, 1])
  ])

289
290
  if input_fields.groundtruth_boxes in tensor_dict:
    bboxes = out_tensor_dict[input_fields.groundtruth_boxes]
291
292
    boxlist = box_list.BoxList(bboxes)
    realigned_bboxes = box_list_ops.change_coordinate_frame(boxlist, im_box)
293
294
295

    realigned_boxes_tensor = realigned_bboxes.get()
    valid_boxes_tensor = assert_or_prune_invalid_boxes(realigned_boxes_tensor)
296
    out_tensor_dict[
297
        input_fields.groundtruth_boxes] = valid_boxes_tensor
298

299
300
  if input_fields.groundtruth_keypoints in tensor_dict:
    keypoints = out_tensor_dict[input_fields.groundtruth_keypoints]
301
302
303
    realigned_keypoints = keypoint_ops.change_coordinate_frame(keypoints,
                                                               im_box)
    out_tensor_dict[
304
305
306
307
        input_fields.groundtruth_keypoints] = realigned_keypoints
    flds_gt_kpt = input_fields.groundtruth_keypoints
    flds_gt_kpt_vis = input_fields.groundtruth_keypoint_visibilities
    flds_gt_kpt_weights = input_fields.groundtruth_keypoint_weights
308
309
310
311
    if flds_gt_kpt_vis not in out_tensor_dict:
      out_tensor_dict[flds_gt_kpt_vis] = tf.ones_like(
          out_tensor_dict[flds_gt_kpt][:, :, 0],
          dtype=tf.bool)
312
313
314
315
316
317
318
319
    flds_gt_kpt_depth = fields.InputDataFields.groundtruth_keypoint_depths
    flds_gt_kpt_depth_weight = (
        fields.InputDataFields.groundtruth_keypoint_depth_weights)
    if flds_gt_kpt_depth in out_tensor_dict:
      out_tensor_dict[flds_gt_kpt_depth] = out_tensor_dict[flds_gt_kpt_depth]
      out_tensor_dict[flds_gt_kpt_depth_weight] = out_tensor_dict[
          flds_gt_kpt_depth_weight]

320
321
322
323
    out_tensor_dict[flds_gt_kpt_weights] = (
        keypoint_ops.keypoint_weights_from_visibilities(
            out_tensor_dict[flds_gt_kpt_vis],
            keypoint_type_weight))
324

325
  dp_surface_coords_fld = input_fields.groundtruth_dp_surface_coords
326
327
328
329
330
331
  if dp_surface_coords_fld in tensor_dict:
    dp_surface_coords = out_tensor_dict[dp_surface_coords_fld]
    realigned_dp_surface_coords = densepose_ops.change_coordinate_frame(
        dp_surface_coords, im_box)
    out_tensor_dict[dp_surface_coords_fld] = realigned_dp_surface_coords

332
333
334
  if use_bfloat16:
    preprocessed_resized_image = tf.cast(
        preprocessed_resized_image, tf.bfloat16)
335
336
337
338
    if input_fields.context_features in out_tensor_dict:
      out_tensor_dict[input_fields.context_features] = tf.cast(
          out_tensor_dict[input_fields.context_features], tf.bfloat16)
  out_tensor_dict[input_fields.image] = tf.squeeze(
339
      preprocessed_resized_image, axis=0)
340
  out_tensor_dict[input_fields.true_image_shape] = tf.squeeze(
341
      true_image_shape, axis=0)
342
343
  if input_fields.groundtruth_instance_masks in out_tensor_dict:
    masks = out_tensor_dict[input_fields.groundtruth_instance_masks]
344
    _, resized_masks, _ = image_resizer_fn(image, masks)
345
346
    if use_bfloat16:
      resized_masks = tf.cast(resized_masks, tf.bfloat16)
pkulzc's avatar
pkulzc committed
347
    out_tensor_dict[
348
        input_fields.groundtruth_instance_masks] = resized_masks
349

pkulzc's avatar
pkulzc committed
350
  zero_indexed_groundtruth_classes = out_tensor_dict[
351
      input_fields.groundtruth_classes] - _LABEL_OFFSET
352
  if use_multiclass_scores:
pkulzc's avatar
pkulzc committed
353
    out_tensor_dict[
354
355
        input_fields.groundtruth_classes] = out_tensor_dict[
            input_fields.multiclass_scores]
pkulzc's avatar
pkulzc committed
356
  else:
357
    out_tensor_dict[input_fields.groundtruth_classes] = tf.one_hot(
pkulzc's avatar
pkulzc committed
358
        zero_indexed_groundtruth_classes, num_classes)
359
  out_tensor_dict.pop(input_fields.multiclass_scores, None)
360

361
  if input_fields.groundtruth_confidences in out_tensor_dict:
pkulzc's avatar
pkulzc committed
362
    groundtruth_confidences = out_tensor_dict[
363
        input_fields.groundtruth_confidences]
364
    # Map the confidences to the one-hot encoding of classes
365
    out_tensor_dict[input_fields.groundtruth_confidences] = (
366
        tf.reshape(groundtruth_confidences, [-1, 1]) *
367
        out_tensor_dict[input_fields.groundtruth_classes])
368
369
370
  else:
    groundtruth_confidences = tf.ones_like(
        zero_indexed_groundtruth_classes, dtype=tf.float32)
371
372
    out_tensor_dict[input_fields.groundtruth_confidences] = (
        out_tensor_dict[input_fields.groundtruth_classes])
373

374
  if merge_multiple_boxes:
375
376
    merged_boxes, merged_classes, merged_confidences, _ = (
        util_ops.merge_boxes_with_multiple_labels(
377
            out_tensor_dict[input_fields.groundtruth_boxes],
378
379
380
            zero_indexed_groundtruth_classes,
            groundtruth_confidences,
            num_classes))
381
    merged_classes = tf.cast(merged_classes, tf.float32)
382
383
384
    out_tensor_dict[input_fields.groundtruth_boxes] = merged_boxes
    out_tensor_dict[input_fields.groundtruth_classes] = merged_classes
    out_tensor_dict[input_fields.groundtruth_confidences] = (
385
        merged_confidences)
386
387
388
  if input_fields.groundtruth_boxes in out_tensor_dict:
    out_tensor_dict[input_fields.num_groundtruth_boxes] = tf.shape(
        out_tensor_dict[input_fields.groundtruth_boxes])[0]
389

pkulzc's avatar
pkulzc committed
390
  return out_tensor_dict
391
392


393
394
395
396
397
def pad_input_data_to_static_shapes(tensor_dict,
                                    max_num_boxes,
                                    num_classes,
                                    spatial_image_shape=None,
                                    max_num_context_features=None,
398
399
                                    context_feature_length=None,
                                    max_dp_points=336):
400
401
  """Pads input tensors to static shapes.

402
403
404
  In case num_additional_channels > 0, we assume that the additional channels
  have already been concatenated to the base image.

405
406
407
408
409
410
411
412
  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.
413
414
415
    max_num_context_features (optional): The maximum number of context
      features needed to compute shapes padding.
    context_feature_length (optional): The length of the context feature.
416
417
418
419
420
    max_dp_points (optional): The maximum number of DensePose sampled points per
      instance. The default (336) is selected since the original DensePose paper
      (https://arxiv.org/pdf/1802.00434.pdf) indicates that the maximum number
      of samples per part is 14, and therefore 24 * 14 = 336 is the maximum
      sampler per instance.
421
422
423
424
425
426

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

  Raises:
427
    ValueError: If groundtruth classes is neither rank 1 nor rank 2, or if we
428
429
430
      detect that additional channels have not been concatenated yet, or if
      max_num_context_features is not specified and context_features is in the
      tensor dict.
431
432
433
434
435
436
  """
  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

437
  input_fields = fields.InputDataFields
438
  num_additional_channels = 0
439
  if input_fields.image_additional_channels in tensor_dict:
440
    num_additional_channels = shape_utils.get_dim_as_int(tensor_dict[
441
        input_fields.image_additional_channels].shape[2])
442
443
444
445

  # We assume that if num_additional_channels > 0, then it has already been
  # concatenated to the base image (but not the ground truth).
  num_channels = 3
446
  if input_fields.image in tensor_dict:
447
    num_channels = shape_utils.get_dim_as_int(
448
        tensor_dict[input_fields.image].shape[2])
449
450
451
452
453
454

  if num_additional_channels:
    if num_additional_channels >= num_channels:
      raise ValueError(
          'Image must be already concatenated with additional channels.')

455
    if (input_fields.original_image in tensor_dict and
456
        shape_utils.get_dim_as_int(
457
            tensor_dict[input_fields.original_image].shape[2]) ==
458
459
460
461
        num_channels):
      raise ValueError(
          'Image must be already concatenated with additional channels.')

462
  if input_fields.context_features in tensor_dict and (
463
464
465
466
467
      max_num_context_features is None):
    raise ValueError('max_num_context_features must be specified in the model '
                     'config if include_context is specified in the input '
                     'config')

468
  padding_shapes = {
469
470
471
      input_fields.image: [height, width, num_channels],
      input_fields.original_image_spatial_shape: [2],
      input_fields.image_additional_channels: [
472
473
          height, width, num_additional_channels
      ],
474
475
476
477
478
479
480
      input_fields.source_id: [],
      input_fields.filename: [],
      input_fields.key: [],
      input_fields.groundtruth_difficult: [max_num_boxes],
      input_fields.groundtruth_boxes: [max_num_boxes, 4],
      input_fields.groundtruth_classes: [max_num_boxes, num_classes],
      input_fields.groundtruth_instance_masks: [
481
482
          max_num_boxes, height, width
      ],
483
      input_fields.groundtruth_instance_mask_weights: [max_num_boxes],
484
485
486
487
488
      input_fields.groundtruth_is_crowd: [max_num_boxes],
      input_fields.groundtruth_group_of: [max_num_boxes],
      input_fields.groundtruth_area: [max_num_boxes],
      input_fields.groundtruth_weights: [max_num_boxes],
      input_fields.groundtruth_confidences: [
489
490
          max_num_boxes, num_classes
      ],
491
492
493
494
495
496
497
      input_fields.num_groundtruth_boxes: [],
      input_fields.groundtruth_label_types: [max_num_boxes],
      input_fields.groundtruth_label_weights: [max_num_boxes],
      input_fields.true_image_shape: [3],
      input_fields.groundtruth_image_classes: [num_classes],
      input_fields.groundtruth_image_confidences: [num_classes],
      input_fields.groundtruth_labeled_classes: [num_classes],
498
499
  }

500
501
  if input_fields.original_image in tensor_dict:
    padding_shapes[input_fields.original_image] = [
502
        height, width,
503
        shape_utils.get_dim_as_int(tensor_dict[input_fields.
504
                                               original_image].shape[2])
505
    ]
506
  if input_fields.groundtruth_keypoints in tensor_dict:
507
    tensor_shape = (
508
        tensor_dict[input_fields.groundtruth_keypoints].shape)
509
510
511
    padding_shape = [max_num_boxes,
                     shape_utils.get_dim_as_int(tensor_shape[1]),
                     shape_utils.get_dim_as_int(tensor_shape[2])]
512
513
514
    padding_shapes[input_fields.groundtruth_keypoints] = padding_shape
  if input_fields.groundtruth_keypoint_visibilities in tensor_dict:
    tensor_shape = tensor_dict[input_fields.
515
                               groundtruth_keypoint_visibilities].shape
516
    padding_shape = [max_num_boxes, shape_utils.get_dim_as_int(tensor_shape[1])]
517
    padding_shapes[input_fields.
518
519
                   groundtruth_keypoint_visibilities] = padding_shape

520
521
522
523
524
525
526
527
528
  if fields.InputDataFields.groundtruth_keypoint_depths in tensor_dict:
    tensor_shape = tensor_dict[fields.InputDataFields.
                               groundtruth_keypoint_depths].shape
    padding_shape = [max_num_boxes, shape_utils.get_dim_as_int(tensor_shape[1])]
    padding_shapes[fields.InputDataFields.
                   groundtruth_keypoint_depths] = padding_shape
    padding_shapes[fields.InputDataFields.
                   groundtruth_keypoint_depth_weights] = padding_shape

529
  if input_fields.groundtruth_keypoint_weights in tensor_dict:
530
    tensor_shape = (
531
        tensor_dict[input_fields.groundtruth_keypoint_weights].shape)
532
    padding_shape = [max_num_boxes, shape_utils.get_dim_as_int(tensor_shape[1])]
533
    padding_shapes[input_fields.
534
                   groundtruth_keypoint_weights] = padding_shape
535
  if input_fields.groundtruth_dp_num_points in tensor_dict:
536
    padding_shapes[
537
        input_fields.groundtruth_dp_num_points] = [max_num_boxes]
538
    padding_shapes[
539
        input_fields.groundtruth_dp_part_ids] = [
540
541
            max_num_boxes, max_dp_points]
    padding_shapes[
542
        input_fields.groundtruth_dp_surface_coords] = [
543
            max_num_boxes, max_dp_points, 4]
544
545
546
547
548
549
550
551
  if input_fields.groundtruth_track_ids in tensor_dict:
    padding_shapes[
        input_fields.groundtruth_track_ids] = [max_num_boxes]

  if input_fields.groundtruth_verified_neg_classes in tensor_dict:
    padding_shapes[
        input_fields.groundtruth_verified_neg_classes] = [num_classes]
  if input_fields.groundtruth_not_exhaustive_classes in tensor_dict:
552
    padding_shapes[
553
        input_fields.groundtruth_not_exhaustive_classes] = [num_classes]
554
555

  # Prepare for ContextRCNN related fields.
556
  if input_fields.context_features in tensor_dict:
557
    padding_shape = [max_num_context_features, context_feature_length]
558
    padding_shapes[input_fields.context_features] = padding_shape
559
560

    tensor_shape = tf.shape(
561
562
563
564
565
566
567
568
        tensor_dict[fields.InputDataFields.context_features])
    tensor_dict[fields.InputDataFields.valid_context_size] = tensor_shape[0]
    padding_shapes[fields.InputDataFields.valid_context_size] = []
  if fields.InputDataFields.context_feature_length in tensor_dict:
    padding_shapes[fields.InputDataFields.context_feature_length] = []
  if fields.InputDataFields.context_features_image_id_list in tensor_dict:
    padding_shapes[fields.InputDataFields.context_features_image_id_list] = [
        max_num_context_features]
569

570
571
  if input_fields.is_annotated in tensor_dict:
    padding_shapes[input_fields.is_annotated] = []
572

573
574
  padded_tensor_dict = {}
  for tensor_name in tensor_dict:
575
576
    padded_tensor_dict[tensor_name] = shape_utils.pad_or_clip_nd(
        tensor_dict[tensor_name], padding_shapes[tensor_name])
577
578
579

  # Make sure that the number of groundtruth boxes now reflects the
  # padded/clipped tensors.
580
581
  if input_fields.num_groundtruth_boxes in padded_tensor_dict:
    padded_tensor_dict[input_fields.num_groundtruth_boxes] = (
582
        tf.minimum(
583
            padded_tensor_dict[input_fields.num_groundtruth_boxes],
584
            max_num_boxes))
585
586
587
  return padded_tensor_dict


588
589
590
591
592
593
594
595
596
597
598
599
600
601
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(
602
      tf.cast(tensor_dict[fields.InputDataFields.image], dtype=tf.float32), 0)
603
604
605

  include_instance_masks = (fields.InputDataFields.groundtruth_instance_masks
                            in tensor_dict)
606
607
  include_instance_mask_weights = (
      fields.InputDataFields.groundtruth_instance_mask_weights in tensor_dict)
608
609
  include_keypoints = (fields.InputDataFields.groundtruth_keypoints
                       in tensor_dict)
610
611
  include_keypoint_visibilities = (
      fields.InputDataFields.groundtruth_keypoint_visibilities in tensor_dict)
612
613
  include_keypoint_depths = (
      fields.InputDataFields.groundtruth_keypoint_depths in tensor_dict)
614
615
616
617
  include_label_weights = (fields.InputDataFields.groundtruth_weights
                           in tensor_dict)
  include_label_confidences = (fields.InputDataFields.groundtruth_confidences
                               in tensor_dict)
618
619
  include_multiclass_scores = (fields.InputDataFields.multiclass_scores in
                               tensor_dict)
620
621
622
623
  dense_pose_fields = [fields.InputDataFields.groundtruth_dp_num_points,
                       fields.InputDataFields.groundtruth_dp_part_ids,
                       fields.InputDataFields.groundtruth_dp_surface_coords]
  include_dense_pose = all(field in tensor_dict for field in dense_pose_fields)
624
625
626
  tensor_dict = preprocessor.preprocess(
      tensor_dict, data_augmentation_options,
      func_arg_map=preprocessor.get_default_func_arg_map(
627
628
          include_label_weights=include_label_weights,
          include_label_confidences=include_label_confidences,
629
          include_multiclass_scores=include_multiclass_scores,
630
          include_instance_masks=include_instance_masks,
631
          include_instance_mask_weights=include_instance_mask_weights,
632
          include_keypoints=include_keypoints,
633
          include_keypoint_visibilities=include_keypoint_visibilities,
634
635
          include_dense_pose=include_dense_pose,
          include_keypoint_depths=include_keypoint_depths))
636
637
638
639
640
  tensor_dict[fields.InputDataFields.image] = tf.squeeze(
      tensor_dict[fields.InputDataFields.image], axis=0)
  return tensor_dict


641
642
643
644
645
646
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,
647
      fields.InputDataFields.groundtruth_weights,
648
649
650
651
652
653
  ]
  labels_dict = {}
  for key in required_label_keys:
    labels_dict[key] = input_dict[key]

  optional_label_keys = [
654
      fields.InputDataFields.groundtruth_confidences,
655
      fields.InputDataFields.groundtruth_labeled_classes,
656
      fields.InputDataFields.groundtruth_keypoints,
657
658
      fields.InputDataFields.groundtruth_keypoint_depths,
      fields.InputDataFields.groundtruth_keypoint_depth_weights,
659
      fields.InputDataFields.groundtruth_instance_masks,
660
      fields.InputDataFields.groundtruth_instance_mask_weights,
661
662
      fields.InputDataFields.groundtruth_area,
      fields.InputDataFields.groundtruth_is_crowd,
663
      fields.InputDataFields.groundtruth_group_of,
664
665
666
      fields.InputDataFields.groundtruth_difficult,
      fields.InputDataFields.groundtruth_keypoint_visibilities,
      fields.InputDataFields.groundtruth_keypoint_weights,
667
668
      fields.InputDataFields.groundtruth_dp_num_points,
      fields.InputDataFields.groundtruth_dp_part_ids,
669
      fields.InputDataFields.groundtruth_dp_surface_coords,
670
671
      fields.InputDataFields.groundtruth_track_ids,
      fields.InputDataFields.groundtruth_verified_neg_classes,
672
673
      fields.InputDataFields.groundtruth_not_exhaustive_classes,
      fields.InputDataFields.groundtruth_image_classes,
674
675
676
677
678
679
680
681
682
683
684
  ]

  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


685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
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


714
def _get_features_dict(input_dict, include_source_id=False):
715
  """Extracts features dict from input dict."""
716
717
718
719
720

  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)
721
722
723
724
725
  features = {
      fields.InputDataFields.image:
          input_dict[fields.InputDataFields.image],
      HASH_KEY: tf.cast(hash_from_source_id, tf.int32),
      fields.InputDataFields.true_image_shape:
pkulzc's avatar
pkulzc committed
726
727
728
          input_dict[fields.InputDataFields.true_image_shape],
      fields.InputDataFields.original_image_spatial_shape:
          input_dict[fields.InputDataFields.original_image_spatial_shape]
729
  }
730
731
  if include_source_id:
    features[fields.InputDataFields.source_id] = source_id
732
733
734
  if fields.InputDataFields.original_image in input_dict:
    features[fields.InputDataFields.original_image] = input_dict[
        fields.InputDataFields.original_image]
735
736
737
  if fields.InputDataFields.image_additional_channels in input_dict:
    features[fields.InputDataFields.image_additional_channels] = input_dict[
        fields.InputDataFields.image_additional_channels]
738
739
740
741
742
743
  if fields.InputDataFields.context_features in input_dict:
    features[fields.InputDataFields.context_features] = input_dict[
        fields.InputDataFields.context_features]
  if fields.InputDataFields.valid_context_size in input_dict:
    features[fields.InputDataFields.valid_context_size] = input_dict[
        fields.InputDataFields.valid_context_size]
744
745
746
  if fields.InputDataFields.context_features_image_id_list in input_dict:
    features[fields.InputDataFields.context_features_image_id_list] = (
        input_dict[fields.InputDataFields.context_features_image_id_list])
747
748
749
  return features


750
751
def create_train_input_fn(train_config, train_input_config,
                          model_config):
752
753
754
755
756
  """Creates a train `input` function for `Estimator`.

  Args:
    train_config: A train_pb2.TrainConfig.
    train_input_config: An input_reader_pb2.InputReader.
757
    model_config: A model_pb2.DetectionModel.
758
759
760
761
762

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

763
  def _train_input_fn(params=None):
764
765
    return train_input(train_config, train_input_config, model_config,
                       params=params)
766

767
  return _train_input_fn
768

769

770
def train_input(train_config, train_input_config,
771
                model_config, model=None, params=None, input_context=None):
772
773
774
775
776
777
778
779
780
  """Returns `features` and `labels` tensor dictionaries for training.

  Args:
    train_config: A train_pb2.TrainConfig.
    train_input_config: An input_reader_pb2.InputReader.
    model_config: A model_pb2.DetectionModel.
    model: A pre-constructed Detection Model.
      If None, one will be created from the config.
    params: Parameter dictionary passed from the estimator.
781
782
783
    input_context: optional, A tf.distribute.InputContext object used to
      shard filenames and compute per-replica batch_size when this function
      is being called per-replica.
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813

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

    features: Dictionary of feature tensors.
      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.
      features[fields.InputDataFields.original_image] (optional) is a
        [batch_size, H, W, C] float32 tensor with original images.
    labels: Dictionary of groundtruth tensors.
      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.
814
815
816
      labels[fields.InputDataFields.groundtruth_instance_mask_weights] is a
        [batch_size, num_boxes] float32 tensor containing groundtruth weights
        for each instance mask.
817
818
819
      labels[fields.InputDataFields.groundtruth_keypoints] is a
        [batch_size, num_boxes, num_keypoints, 2] float32 tensor containing
        keypoints for each box.
820
821
822
823
824
825
      labels[fields.InputDataFields.groundtruth_weights] is a
        [batch_size, num_boxes, num_keypoints] float32 tensor containing
        groundtruth weights for the keypoints.
      labels[fields.InputDataFields.groundtruth_visibilities] is a
        [batch_size, num_boxes, num_keypoints] bool tensor containing
        groundtruth visibilities for each keypoint.
826
827
      labels[fields.InputDataFields.groundtruth_labeled_classes] is a
        [batch_size, num_classes] float32 k-hot tensor of classes.
828
829
830
831
832
833
834
835
836
837
838
      labels[fields.InputDataFields.groundtruth_dp_num_points] is a
        [batch_size, num_boxes] int32 tensor with the number of sampled
        DensePose points per object.
      labels[fields.InputDataFields.groundtruth_dp_part_ids] is a
        [batch_size, num_boxes, max_sampled_points] int32 tensor with the
        DensePose part ids (0-indexed) per object.
      labels[fields.InputDataFields.groundtruth_dp_surface_coords] is a
        [batch_size, num_boxes, max_sampled_points, 4] float32 tensor with the
        DensePose surface coordinates. The format is (y, x, v, u), where (y, x)
        are normalized image coordinates and (v, u) are normalized surface part
        coordinates.
839
840
      labels[fields.InputDataFields.groundtruth_track_ids] is a
        [batch_size, num_boxes] int32 tensor with the track ID for each object.
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861

  Raises:
    TypeError: if the `train_config`, `train_input_config` or `model_config`
      are not of the correct type.
  """
  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.')
  if not isinstance(model_config, model_pb2.DetectionModel):
    raise TypeError('The `model_config` must be a '
                    'model_pb2.DetectionModel.')

  if model is None:
    model_preprocess_fn = INPUT_BUILDER_UTIL_MAP['model_build'](
        model_config, is_training=True).preprocess
  else:
    model_preprocess_fn = model.preprocess

862
863
  num_classes = config_util.get_number_of_classes(model_config)

864
865
866
867
868
869
870
871
872
873
874
875
  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)

    image_resizer_config = config_util.get_image_resizer_config(model_config)
    image_resizer_fn = image_resizer_builder.build(image_resizer_config)
876
    keypoint_type_weight = train_input_config.keypoint_type_weight or None
877
878
879
    transform_data_fn = functools.partial(
        transform_input_data, model_preprocess_fn=model_preprocess_fn,
        image_resizer_fn=image_resizer_fn,
880
        num_classes=num_classes,
881
882
883
884
        data_augmentation_fn=data_augmentation_fn,
        merge_multiple_boxes=train_config.merge_multiple_label_boxes,
        retain_original_image=train_config.retain_original_images,
        use_multiclass_scores=train_config.use_multiclass_scores,
885
886
        use_bfloat16=train_config.use_bfloat16,
        keypoint_type_weight=keypoint_type_weight)
887
888
889
890

    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,
891
        num_classes=num_classes,
892
        spatial_image_shape=config_util.get_spatial_image_size(
893
894
895
896
897
898
899
900
            image_resizer_config),
        max_num_context_features=config_util.get_max_num_context_features(
            model_config),
        context_feature_length=config_util.get_context_feature_length(
            model_config))
    include_source_id = train_input_config.include_source_id
    return (_get_features_dict(tensor_dict, include_source_id),
            _get_labels_dict(tensor_dict))
901
  reduce_to_frame_fn = get_reduce_to_frame_fn(train_input_config, True)
902
903
904
905

  dataset = INPUT_BUILDER_UTIL_MAP['dataset_build'](
      train_input_config,
      transform_input_data_fn=transform_and_pad_input_data_fn,
906
      batch_size=params['batch_size'] if params else train_config.batch_size,
907
908
      input_context=input_context,
      reduce_to_frame_fn=reduce_to_frame_fn)
909
  return dataset
910
911


912
def create_eval_input_fn(eval_config, eval_input_config, model_config):
913
914
915
916
917
  """Creates an eval `input` function for `Estimator`.

  Args:
    eval_config: An eval_pb2.EvalConfig.
    eval_input_config: An input_reader_pb2.InputReader.
918
    model_config: A model_pb2.DetectionModel.
919
920
921
922
923

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

924
  def _eval_input_fn(params=None):
925
926
    return eval_input(eval_config, eval_input_config, model_config,
                      params=params)
927

928
  return _eval_input_fn
929

930

931
def eval_input(eval_config, eval_input_config, model_config,
932
               model=None, params=None, input_context=None):
933
934
935
936
937
938
939
940
941
  """Returns `features` and `labels` tensor dictionaries for evaluation.

  Args:
    eval_config: An eval_pb2.EvalConfig.
    eval_input_config: An input_reader_pb2.InputReader.
    model_config: A model_pb2.DetectionModel.
    model: A pre-constructed Detection Model.
      If None, one will be created from the config.
    params: Parameter dictionary passed from the estimator.
942
943
944
    input_context: optional, A tf.distribute.InputContext object used to
      shard filenames and compute per-replica batch_size when this function
      is being called per-replica.
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973

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

    features: Dictionary of feature tensors.
      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.
    labels: Dictionary of groundtruth tensors.
      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.
974
975
976
      labels[fields.InputDataFields.groundtruth_instance_mask_weights] is a
        [1, num_boxes] float32 tensor containing groundtruth weights for each
        instance mask.
977
978
979
980
981
982
      labels[fields.InputDataFields.groundtruth_weights] is a
        [batch_size, num_boxes, num_keypoints] float32 tensor containing
        groundtruth weights for the keypoints.
      labels[fields.InputDataFields.groundtruth_visibilities] is a
        [batch_size, num_boxes, num_keypoints] bool tensor containing
        groundtruth visibilities for each keypoint.
983
984
985
986
987
      labels[fields.InputDataFields.groundtruth_group_of] is a [1, num_boxes]
        bool tensor indicating if the box covers more than 5 instances of the
        same class which heavily occlude each other.
      labels[fields.InputDataFields.groundtruth_labeled_classes] is a
        [num_boxes, num_classes] float32 k-hot tensor of classes.
988
989
990
991
992
993
994
995
996
997
998
      labels[fields.InputDataFields.groundtruth_dp_num_points] is a
        [batch_size, num_boxes] int32 tensor with the number of sampled
        DensePose points per object.
      labels[fields.InputDataFields.groundtruth_dp_part_ids] is a
        [batch_size, num_boxes, max_sampled_points] int32 tensor with the
        DensePose part ids (0-indexed) per object.
      labels[fields.InputDataFields.groundtruth_dp_surface_coords] is a
        [batch_size, num_boxes, max_sampled_points, 4] float32 tensor with the
        DensePose surface coordinates. The format is (y, x, v, u), where (y, x)
        are normalized image coordinates and (v, u) are normalized surface part
        coordinates.
999
1000
      labels[fields.InputDataFields.groundtruth_track_ids] is a
        [batch_size, num_boxes] int32 tensor with the track ID for each object.
1001
1002
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016

  Raises:
    TypeError: if the `eval_config`, `eval_input_config` or `model_config`
      are not of the correct type.
  """
  params = params or {}
  if not isinstance(eval_config, eval_pb2.EvalConfig):
    raise TypeError('For eval mode, the `eval_config` must be a '
                    'train_pb2.EvalConfig.')
  if not isinstance(eval_input_config, input_reader_pb2.InputReader):
    raise TypeError('The `eval_input_config` must be a '
                    'input_reader_pb2.InputReader.')
  if not isinstance(model_config, model_pb2.DetectionModel):
    raise TypeError('The `model_config` must be a '
                    'model_pb2.DetectionModel.')

1017
1018
1019
1020
1021
1022
1023
1024
  if eval_config.force_no_resize:
    arch = model_config.WhichOneof('model')
    arch_config = getattr(model_config, arch)
    image_resizer_proto = image_resizer_pb2.ImageResizer()
    image_resizer_proto.identity_resizer.CopyFrom(
        image_resizer_pb2.IdentityResizer())
    arch_config.image_resizer.CopyFrom(image_resizer_proto)

1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
  if model is None:
    model_preprocess_fn = INPUT_BUILDER_UTIL_MAP['model_build'](
        model_config, is_training=False).preprocess
  else:
    model_preprocess_fn = model.preprocess

  def transform_and_pad_input_data_fn(tensor_dict):
    """Combines transform and pad operation."""
    num_classes = config_util.get_number_of_classes(model_config)

    image_resizer_config = config_util.get_image_resizer_config(model_config)
    image_resizer_fn = image_resizer_builder.build(image_resizer_config)
1037
    keypoint_type_weight = eval_input_config.keypoint_type_weight or None
1038
1039
1040
1041
1042
1043

    transform_data_fn = functools.partial(
        transform_input_data, model_preprocess_fn=model_preprocess_fn,
        image_resizer_fn=image_resizer_fn,
        num_classes=num_classes,
        data_augmentation_fn=None,
1044
1045
        retain_original_image=eval_config.retain_original_images,
        retain_original_image_additional_channels=
1046
1047
        eval_config.retain_original_image_additional_channels,
        keypoint_type_weight=keypoint_type_weight)
1048
1049
1050
1051
1052
    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(
1053
1054
1055
1056
1057
1058
1059
1060
            image_resizer_config),
        max_num_context_features=config_util.get_max_num_context_features(
            model_config),
        context_feature_length=config_util.get_context_feature_length(
            model_config))
    include_source_id = eval_input_config.include_source_id
    return (_get_features_dict(tensor_dict, include_source_id),
            _get_labels_dict(tensor_dict))
1061
1062
1063

  reduce_to_frame_fn = get_reduce_to_frame_fn(eval_input_config, False)

1064
1065
1066
  dataset = INPUT_BUILDER_UTIL_MAP['dataset_build'](
      eval_input_config,
      batch_size=params['batch_size'] if params else eval_config.batch_size,
1067
      transform_input_data_fn=transform_and_pad_input_data_fn,
1068
      input_context=input_context,
1069
      reduce_to_frame_fn=reduce_to_frame_fn)
1070
  return dataset
1071
1072


1073
def create_predict_input_fn(model_config, predict_input_config):
1074
1075
  """Creates a predict `input` function for `Estimator`.

1076
1077
  Args:
    model_config: A model_pb2.DetectionModel.
1078
    predict_input_config: An input_reader_pb2.InputReader.
1079

1080
1081
1082
1083
  Returns:
    `input_fn` for `Estimator` in PREDICT mode.
  """

1084
  def _predict_input_fn(params=None):
1085
1086
    """Decodes serialized tf.Examples and returns `ServingInputReceiver`.

1087
1088
1089
    Args:
      params: Parameter dictionary passed from the estimator.

1090
1091
1092
    Returns:
      `ServingInputReceiver`.
    """
1093
    del params
1094
    example = tf.placeholder(dtype=tf.string, shape=[], name='tf_example')
1095

1096
    num_classes = config_util.get_number_of_classes(model_config)
1097
1098
1099
    model_preprocess_fn = INPUT_BUILDER_UTIL_MAP['model_build'](
        model_config, is_training=False).preprocess

1100
1101
    image_resizer_config = config_util.get_image_resizer_config(model_config)
    image_resizer_fn = image_resizer_builder.build(image_resizer_config)
1102

1103
    transform_fn = functools.partial(
1104
        transform_input_data, model_preprocess_fn=model_preprocess_fn,
1105
1106
1107
1108
        image_resizer_fn=image_resizer_fn,
        num_classes=num_classes,
        data_augmentation_fn=None)

1109
1110
1111
    decoder = tf_example_decoder.TfExampleDecoder(
        load_instance_masks=False,
        num_additional_channels=predict_input_config.num_additional_channels)
1112
    input_dict = transform_fn(decoder.decode(example))
1113
    images = tf.cast(input_dict[fields.InputDataFields.image], dtype=tf.float32)
1114
    images = tf.expand_dims(images, axis=0)
1115
1116
    true_image_shape = tf.expand_dims(
        input_dict[fields.InputDataFields.true_image_shape], axis=0)
1117

1118
    return tf_estimator.export.ServingInputReceiver(
1119
1120
1121
        features={
            fields.InputDataFields.image: images,
            fields.InputDataFields.true_image_shape: true_image_shape},
1122
1123
1124
        receiver_tensors={SERVING_FED_EXAMPLE_KEY: example})

  return _predict_input_fn
1125
1126
1127
1128
1129
1130
1131
1132
1133
1134
1135
1136
1137
1138
1139
1140
1141
1142
1143
1144
1145


def get_reduce_to_frame_fn(input_reader_config, is_training):
  """Returns a function reducing sequence tensors to single frame tensors.

  If the input type is not TF_SEQUENCE_EXAMPLE, the tensors are passed through
  this function unchanged. Otherwise, when in training mode, a single frame is
  selected at random from the sequence example, and the tensors for that frame
  are converted to single frame tensors, with all associated context features.
  In evaluation mode all frames are converted to single frame tensors with
  copied context tensors. After the sequence example tensors are converted into
  one or many single frame tensors, the images from each frame are decoded.

  Args:
    input_reader_config: An input_reader_pb2.InputReader.
    is_training: Whether we are in training mode.

  Returns:
    `reduce_to_frame_fn` for the dataset builder
  """
  if input_reader_config.input_type != (
1146
1147
      input_reader_pb2.InputType.Value('TF_SEQUENCE_EXAMPLE')):
    return lambda dataset, dataset_map_fn, batch_size, config: dataset
1148
  else:
1149
1150
    def reduce_to_frame(dataset, dataset_map_fn, batch_size,
                        input_reader_config):
1151
1152
1153
1154
      """Returns a function reducing sequence tensors to single frame tensors.

      Args:
        dataset: A tf dataset containing sequence tensors.
1155
1156
1157
1158
1159
1160
        dataset_map_fn: A function that handles whether to
          map_with_legacy_function for this dataset
        batch_size: used if map_with_legacy_function is true to determine
          num_parallel_calls
        input_reader_config: used if map_with_legacy_function is true to
          determine num_parallel_calls
1161
1162
1163
1164
1165
1166
1167
1168
1169
1170
1171
1172
1173
1174
1175
1176
1177
1178
1179
1180
1181

      Returns:
        A tf dataset containing single frame tensors.
      """
      if is_training:
        def get_single_frame(tensor_dict):
          """Returns a random frame from a sequence.

          Picks a random frame and returns slices of sequence tensors
          corresponding to the random frame. Returns non-sequence tensors
          unchanged.

          Args:
            tensor_dict: A dictionary containing sequence tensors.

          Returns:
            Tensors for a single random frame within the sequence.
          """
          num_frames = tf.cast(
              tf.shape(tensor_dict[fields.InputDataFields.source_id])[0],
              dtype=tf.int32)
1182
1183
1184
1185
1186
1187
          if input_reader_config.frame_index == -1:
            frame_index = tf.random.uniform((), minval=0, maxval=num_frames,
                                            dtype=tf.int32)
          else:
            frame_index = tf.constant(input_reader_config.frame_index,
                                      dtype=tf.int32)
1188
1189
1190
1191
1192
1193
1194
1195
1196
          out_tensor_dict = {}
          for key in tensor_dict:
            if key in fields.SEQUENCE_FIELDS:
              # Slice random frame from sequence tensors
              out_tensor_dict[key] = tensor_dict[key][frame_index]
            else:
              # Copy all context tensors.
              out_tensor_dict[key] = tensor_dict[key]
          return out_tensor_dict
1197
1198
        dataset = dataset_map_fn(dataset, get_single_frame, batch_size,
                                 input_reader_config)
1199
      else:
1200
1201
        dataset = dataset_map_fn(dataset, util_ops.tile_context_tensors,
                                 batch_size, input_reader_config)
1202
1203
        dataset = dataset.unbatch()
      # Decode frame here as SequenceExample tensors contain encoded images.
1204
1205
      dataset = dataset_map_fn(dataset, util_ops.decode_image, batch_size,
                               input_reader_config)
1206
1207
      return dataset
    return reduce_to_frame