inputs.py 52.2 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 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
29
from object_detection.core import box_list
from object_detection.core import box_list_ops
30
from object_detection.core import densepose_ops
31
from object_detection.core import keypoint_ops
32
from object_detection.core import preprocessor
33
34
35
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
36
from object_detection.protos import image_resizer_pb2
37
from object_detection.protos import input_reader_pb2
38
from object_detection.protos import model_pb2
39
from object_detection.protos import train_pb2
40
from object_detection.utils import config_util
41
from object_detection.utils import ops as util_ops
42
from object_detection.utils import shape_utils
43

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

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

55

pkulzc's avatar
pkulzc committed
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
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
71
72
73
def convert_labeled_classes_to_k_hot(groundtruth_labeled_classes,
                                     num_classes,
                                     map_empty_to_ones=False):
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
  """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`.
  """
90
91
92
93
94
95
96
97
98
99
100
101

  # 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)

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


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

110
111
  recognized_indices = tf.squeeze(
      tf.where(tf.greater(class_ids, unrecognized_label)), -1)
112
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
  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()


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

  Data transformation functions are applied in the following order.
163
164
165
166
167
  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.
168
169
170
171
172
173
  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
174
     tensor_dict.
175
176
  6. one_hot_encoding: applied to classes tensor in tensor_dict.
  7. merge_multiple_boxes (optional): when groundtruth boxes are exactly the
177
178
179
180
181
182
183
184
185
     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.
186
187
188
189
    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.
190
191
192
193
194
195
196
197
    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
198
199
200
201
    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.
202
    use_bfloat16: (optional) a bool, whether to use bfloat16 in training.
203
204
    retain_original_image_additional_channels: (optional) Whether to retain
      original image additional channels in the output dictionary.
205
206
207
    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.
208
209
210
211

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

  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.
217
  """
pkulzc's avatar
pkulzc committed
218
  out_tensor_dict = tensor_dict.copy()
219

220
221
222
223
224
225
  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

226
227
228
229
230
  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.')

231
232
233
234
235
  for field, map_empty_to_ones in [
      (labeled_classes_field, True),
      (image_classes_field, True),
      (verified_neg_classes_field, False),
      (not_exhaustive_field, False)]:
236
237
238
    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
239
      out_tensor_dict[field] = convert_labeled_classes_to_k_hot(
240
          out_tensor_dict[field], num_classes, map_empty_to_ones)
241
242

  if input_fields.multiclass_scores in out_tensor_dict:
pkulzc's avatar
pkulzc committed
243
    out_tensor_dict[
244
        input_fields
pkulzc's avatar
pkulzc committed
245
        .multiclass_scores] = _multiclass_scores_or_one_hot_labels(
246
247
248
            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
249
250
            num_classes)

251
  if input_fields.groundtruth_boxes in out_tensor_dict:
pkulzc's avatar
pkulzc committed
252
253
254
    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)
255

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

261
262
263
264
  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)
265
266
    if retain_original_image_additional_channels:
      out_tensor_dict[
267
          input_fields.image_additional_channels] = tf.cast(
268
              image_resizer_fn(channels, None)[0], tf.uint8)
269

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

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

  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])
  ])

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

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

298
299
  if input_fields.groundtruth_keypoints in tensor_dict:
    keypoints = out_tensor_dict[input_fields.groundtruth_keypoints]
300
301
302
    realigned_keypoints = keypoint_ops.change_coordinate_frame(keypoints,
                                                               im_box)
    out_tensor_dict[
303
304
305
306
        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
307
308
309
310
    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)
311
312
313
314
315
316
317
318
    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]

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

324
  dp_surface_coords_fld = input_fields.groundtruth_dp_surface_coords
325
326
327
328
329
330
  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

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

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

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

373
  if merge_multiple_boxes:
374
375
    merged_boxes, merged_classes, merged_confidences, _ = (
        util_ops.merge_boxes_with_multiple_labels(
376
            out_tensor_dict[input_fields.groundtruth_boxes],
377
378
379
            zero_indexed_groundtruth_classes,
            groundtruth_confidences,
            num_classes))
380
    merged_classes = tf.cast(merged_classes, tf.float32)
381
382
383
    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] = (
384
        merged_confidences)
385
386
387
  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]
388

pkulzc's avatar
pkulzc committed
389
  return out_tensor_dict
390
391


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

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

404
405
406
407
408
409
410
411
  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.
412
413
414
    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.
415
416
417
418
419
    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.
420
421
422
423
424
425

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

  Raises:
426
    ValueError: If groundtruth classes is neither rank 1 nor rank 2, or if we
427
428
429
      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.
430
431
432
433
434
435
  """
  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

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

  # 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
445
  if input_fields.image in tensor_dict:
446
    num_channels = shape_utils.get_dim_as_int(
447
        tensor_dict[input_fields.image].shape[2])
448
449
450
451
452
453

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

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

461
  if input_fields.context_features in tensor_dict and (
462
463
464
465
466
      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')

467
  padding_shapes = {
468
469
470
      input_fields.image: [height, width, num_channels],
      input_fields.original_image_spatial_shape: [2],
      input_fields.image_additional_channels: [
471
472
          height, width, num_additional_channels
      ],
473
474
475
476
477
478
479
      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: [
480
481
          max_num_boxes, height, width
      ],
482
483
484
485
486
      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: [
487
488
          max_num_boxes, num_classes
      ],
489
490
491
492
493
494
495
      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],
496
497
  }

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

518
519
520
521
522
523
524
525
526
  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

527
  if input_fields.groundtruth_keypoint_weights in tensor_dict:
528
    tensor_shape = (
529
        tensor_dict[input_fields.groundtruth_keypoint_weights].shape)
530
    padding_shape = [max_num_boxes, shape_utils.get_dim_as_int(tensor_shape[1])]
531
    padding_shapes[input_fields.
532
                   groundtruth_keypoint_weights] = padding_shape
533
  if input_fields.groundtruth_dp_num_points in tensor_dict:
534
    padding_shapes[
535
        input_fields.groundtruth_dp_num_points] = [max_num_boxes]
536
    padding_shapes[
537
        input_fields.groundtruth_dp_part_ids] = [
538
539
            max_num_boxes, max_dp_points]
    padding_shapes[
540
        input_fields.groundtruth_dp_surface_coords] = [
541
            max_num_boxes, max_dp_points, 4]
542
543
544
545
546
547
548
549
  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:
550
    padding_shapes[
551
        input_fields.groundtruth_not_exhaustive_classes] = [num_classes]
552
553

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

    tensor_shape = tf.shape(
559
560
561
562
563
564
565
566
        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]
567

568
569
  if input_fields.is_annotated in tensor_dict:
    padding_shapes[input_fields.is_annotated] = []
570

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

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


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

  include_instance_masks = (fields.InputDataFields.groundtruth_instance_masks
                            in tensor_dict)
  include_keypoints = (fields.InputDataFields.groundtruth_keypoints
                       in tensor_dict)
606
607
  include_keypoint_visibilities = (
      fields.InputDataFields.groundtruth_keypoint_visibilities in tensor_dict)
608
609
  include_keypoint_depths = (
      fields.InputDataFields.groundtruth_keypoint_depths in tensor_dict)
610
611
612
613
  include_label_weights = (fields.InputDataFields.groundtruth_weights
                           in tensor_dict)
  include_label_confidences = (fields.InputDataFields.groundtruth_confidences
                               in tensor_dict)
614
615
  include_multiclass_scores = (fields.InputDataFields.multiclass_scores in
                               tensor_dict)
616
617
618
619
  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)
620
621
622
  tensor_dict = preprocessor.preprocess(
      tensor_dict, data_augmentation_options,
      func_arg_map=preprocessor.get_default_func_arg_map(
623
624
          include_label_weights=include_label_weights,
          include_label_confidences=include_label_confidences,
625
          include_multiclass_scores=include_multiclass_scores,
626
          include_instance_masks=include_instance_masks,
627
          include_keypoints=include_keypoints,
628
          include_keypoint_visibilities=include_keypoint_visibilities,
629
630
          include_dense_pose=include_dense_pose,
          include_keypoint_depths=include_keypoint_depths))
631
632
633
634
635
  tensor_dict[fields.InputDataFields.image] = tf.squeeze(
      tensor_dict[fields.InputDataFields.image], axis=0)
  return tensor_dict


636
637
638
639
640
641
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,
642
      fields.InputDataFields.groundtruth_weights,
643
644
645
646
647
648
  ]
  labels_dict = {}
  for key in required_label_keys:
    labels_dict[key] = input_dict[key]

  optional_label_keys = [
649
      fields.InputDataFields.groundtruth_confidences,
650
      fields.InputDataFields.groundtruth_labeled_classes,
651
      fields.InputDataFields.groundtruth_keypoints,
652
653
      fields.InputDataFields.groundtruth_keypoint_depths,
      fields.InputDataFields.groundtruth_keypoint_depth_weights,
654
655
656
      fields.InputDataFields.groundtruth_instance_masks,
      fields.InputDataFields.groundtruth_area,
      fields.InputDataFields.groundtruth_is_crowd,
657
      fields.InputDataFields.groundtruth_group_of,
658
659
660
      fields.InputDataFields.groundtruth_difficult,
      fields.InputDataFields.groundtruth_keypoint_visibilities,
      fields.InputDataFields.groundtruth_keypoint_weights,
661
662
      fields.InputDataFields.groundtruth_dp_num_points,
      fields.InputDataFields.groundtruth_dp_part_ids,
663
      fields.InputDataFields.groundtruth_dp_surface_coords,
664
665
666
      fields.InputDataFields.groundtruth_track_ids,
      fields.InputDataFields.groundtruth_verified_neg_classes,
      fields.InputDataFields.groundtruth_not_exhaustive_classes
667
668
669
670
671
672
673
674
675
676
677
  ]

  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


678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
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


707
def _get_features_dict(input_dict, include_source_id=False):
708
  """Extracts features dict from input dict."""
709
710
711
712
713

  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)
714
715
716
717
718
  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
719
720
721
          input_dict[fields.InputDataFields.true_image_shape],
      fields.InputDataFields.original_image_spatial_shape:
          input_dict[fields.InputDataFields.original_image_spatial_shape]
722
  }
723
724
  if include_source_id:
    features[fields.InputDataFields.source_id] = source_id
725
726
727
  if fields.InputDataFields.original_image in input_dict:
    features[fields.InputDataFields.original_image] = input_dict[
        fields.InputDataFields.original_image]
728
729
730
  if fields.InputDataFields.image_additional_channels in input_dict:
    features[fields.InputDataFields.image_additional_channels] = input_dict[
        fields.InputDataFields.image_additional_channels]
731
732
733
734
735
736
  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]
737
738
739
  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])
740
741
742
  return features


743
744
def create_train_input_fn(train_config, train_input_config,
                          model_config):
745
746
747
748
749
  """Creates a train `input` function for `Estimator`.

  Args:
    train_config: A train_pb2.TrainConfig.
    train_input_config: An input_reader_pb2.InputReader.
750
    model_config: A model_pb2.DetectionModel.
751
752
753
754
755

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

756
  def _train_input_fn(params=None):
757
758
    return train_input(train_config, train_input_config, model_config,
                       params=params)
759

760
  return _train_input_fn
761

762

763
def train_input(train_config, train_input_config,
764
                model_config, model=None, params=None, input_context=None):
765
766
767
768
769
770
771
772
773
  """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.
774
775
776
    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.
777
778
779
780
781
782
783
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

  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.
      labels[fields.InputDataFields.groundtruth_keypoints] is a
        [batch_size, num_boxes, num_keypoints, 2] float32 tensor containing
        keypoints for each box.
810
811
812
813
814
815
      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.
816
817
      labels[fields.InputDataFields.groundtruth_labeled_classes] is a
        [batch_size, num_classes] float32 k-hot tensor of classes.
818
819
820
821
822
823
824
825
826
827
828
      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.
829
830
      labels[fields.InputDataFields.groundtruth_track_ids] is a
        [batch_size, num_boxes] int32 tensor with the track ID for each object.
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851

  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

852
853
  num_classes = config_util.get_number_of_classes(model_config)

854
855
856
857
858
859
860
861
862
863
864
865
  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)
866
    keypoint_type_weight = train_input_config.keypoint_type_weight or None
867
868
869
    transform_data_fn = functools.partial(
        transform_input_data, model_preprocess_fn=model_preprocess_fn,
        image_resizer_fn=image_resizer_fn,
870
        num_classes=num_classes,
871
872
873
874
        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,
875
876
        use_bfloat16=train_config.use_bfloat16,
        keypoint_type_weight=keypoint_type_weight)
877
878
879
880

    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,
881
        num_classes=num_classes,
882
        spatial_image_shape=config_util.get_spatial_image_size(
883
884
885
886
887
888
889
890
            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))
891
  reduce_to_frame_fn = get_reduce_to_frame_fn(train_input_config, True)
892
893
894
895

  dataset = INPUT_BUILDER_UTIL_MAP['dataset_build'](
      train_input_config,
      transform_input_data_fn=transform_and_pad_input_data_fn,
896
      batch_size=params['batch_size'] if params else train_config.batch_size,
897
898
      input_context=input_context,
      reduce_to_frame_fn=reduce_to_frame_fn)
899
  return dataset
900
901


902
def create_eval_input_fn(eval_config, eval_input_config, model_config):
903
904
905
906
907
  """Creates an eval `input` function for `Estimator`.

  Args:
    eval_config: An eval_pb2.EvalConfig.
    eval_input_config: An input_reader_pb2.InputReader.
908
    model_config: A model_pb2.DetectionModel.
909
910
911
912
913

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

914
  def _eval_input_fn(params=None):
915
916
    return eval_input(eval_config, eval_input_config, model_config,
                      params=params)
917

918
  return _eval_input_fn
919

920

921
def eval_input(eval_config, eval_input_config, model_config,
922
               model=None, params=None, input_context=None):
923
924
925
926
927
928
929
930
931
  """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.
932
933
934
    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.
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963

  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.
964
965
966
967
968
969
      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.
970
971
972
973
974
      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.
975
976
977
978
979
980
981
982
983
984
985
      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.
986
987
      labels[fields.InputDataFields.groundtruth_track_ids] is a
        [batch_size, num_boxes] int32 tensor with the track ID for each object.
988
989
990
991
992
993
994
995
996
997
998
999
1000
1001
1002
1003

  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.')

1004
1005
1006
1007
1008
1009
1010
1011
  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)

1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
  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)
1024
    keypoint_type_weight = eval_input_config.keypoint_type_weight or None
1025
1026
1027
1028
1029
1030

    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,
1031
1032
        retain_original_image=eval_config.retain_original_images,
        retain_original_image_additional_channels=
1033
1034
        eval_config.retain_original_image_additional_channels,
        keypoint_type_weight=keypoint_type_weight)
1035
1036
1037
1038
1039
    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(
1040
1041
1042
1043
1044
1045
1046
1047
            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))
1048
1049
1050

  reduce_to_frame_fn = get_reduce_to_frame_fn(eval_input_config, False)

1051
1052
1053
  dataset = INPUT_BUILDER_UTIL_MAP['dataset_build'](
      eval_input_config,
      batch_size=params['batch_size'] if params else eval_config.batch_size,
1054
      transform_input_data_fn=transform_and_pad_input_data_fn,
1055
      input_context=input_context,
1056
      reduce_to_frame_fn=reduce_to_frame_fn)
1057
  return dataset
1058
1059


1060
def create_predict_input_fn(model_config, predict_input_config):
1061
1062
  """Creates a predict `input` function for `Estimator`.

1063
1064
  Args:
    model_config: A model_pb2.DetectionModel.
1065
    predict_input_config: An input_reader_pb2.InputReader.
1066

1067
1068
1069
1070
  Returns:
    `input_fn` for `Estimator` in PREDICT mode.
  """

1071
  def _predict_input_fn(params=None):
1072
1073
    """Decodes serialized tf.Examples and returns `ServingInputReceiver`.

1074
1075
1076
    Args:
      params: Parameter dictionary passed from the estimator.

1077
1078
1079
    Returns:
      `ServingInputReceiver`.
    """
1080
    del params
1081
    example = tf.placeholder(dtype=tf.string, shape=[], name='tf_example')
1082

1083
    num_classes = config_util.get_number_of_classes(model_config)
1084
1085
1086
    model_preprocess_fn = INPUT_BUILDER_UTIL_MAP['model_build'](
        model_config, is_training=False).preprocess

1087
1088
    image_resizer_config = config_util.get_image_resizer_config(model_config)
    image_resizer_fn = image_resizer_builder.build(image_resizer_config)
1089

1090
    transform_fn = functools.partial(
1091
        transform_input_data, model_preprocess_fn=model_preprocess_fn,
1092
1093
1094
1095
        image_resizer_fn=image_resizer_fn,
        num_classes=num_classes,
        data_augmentation_fn=None)

1096
1097
1098
    decoder = tf_example_decoder.TfExampleDecoder(
        load_instance_masks=False,
        num_additional_channels=predict_input_config.num_additional_channels)
1099
    input_dict = transform_fn(decoder.decode(example))
1100
    images = tf.cast(input_dict[fields.InputDataFields.image], dtype=tf.float32)
1101
    images = tf.expand_dims(images, axis=0)
1102
1103
    true_image_shape = tf.expand_dims(
        input_dict[fields.InputDataFields.true_image_shape], axis=0)
1104
1105

    return tf.estimator.export.ServingInputReceiver(
1106
1107
1108
        features={
            fields.InputDataFields.image: images,
            fields.InputDataFields.true_image_shape: true_image_shape},
1109
1110
1111
        receiver_tensors={SERVING_FED_EXAMPLE_KEY: example})

  return _predict_input_fn
1112
1113
1114
1115
1116
1117
1118
1119
1120
1121
1122
1123
1124
1125
1126
1127
1128
1129
1130
1131
1132


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 != (
1133
1134
      input_reader_pb2.InputType.Value('TF_SEQUENCE_EXAMPLE')):
    return lambda dataset, dataset_map_fn, batch_size, config: dataset
1135
  else:
1136
1137
    def reduce_to_frame(dataset, dataset_map_fn, batch_size,
                        input_reader_config):
1138
1139
1140
1141
      """Returns a function reducing sequence tensors to single frame tensors.

      Args:
        dataset: A tf dataset containing sequence tensors.
1142
1143
1144
1145
1146
1147
        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
1148
1149
1150
1151
1152
1153
1154
1155
1156
1157
1158
1159
1160
1161
1162
1163
1164
1165
1166
1167
1168

      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)
1169
1170
1171
1172
1173
1174
          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)
1175
1176
1177
1178
1179
1180
1181
1182
1183
          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
1184
1185
        dataset = dataset_map_fn(dataset, get_single_frame, batch_size,
                                 input_reader_config)
1186
      else:
1187
1188
        dataset = dataset_map_fn(dataset, util_ops.tile_context_tensors,
                                 batch_size, input_reader_config)
1189
1190
        dataset = dataset.unbatch()
      # Decode frame here as SequenceExample tensors contain encoded images.
1191
1192
      dataset = dataset_map_fn(dataset, util_ops.decode_image, batch_size,
                               input_reader_config)
1193
1194
      return dataset
    return reduce_to_frame