inputs.py 51 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)


71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
def _convert_labeled_classes_to_k_hot(groundtruth_labeled_classes, num_classes,
                                      map_empty_to_ones=False):
  """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`.
  """
89
90
91
92
93
94
95
96
97
98
99
100

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

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


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

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


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

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

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

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

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

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

230
231
232
233
234
  for field, map_empty_to_ones in [
      (labeled_classes_field, True),
      (image_classes_field, True),
      (verified_neg_classes_field, False),
      (not_exhaustive_field, False)]:
235
236
237
238
    if field in out_tensor_dict:
      out_tensor_dict[field] = _remove_unrecognized_classes(
          out_tensor_dict[field], unrecognized_label=-1)
      out_tensor_dict[field] = _convert_labeled_classes_to_k_hot(
239
          out_tensor_dict[field], num_classes, map_empty_to_ones)
240
241

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

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

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

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

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

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

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

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

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

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

315
  dp_surface_coords_fld = input_fields.groundtruth_dp_surface_coords
316
317
318
319
320
321
  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

322
323
324
  if use_bfloat16:
    preprocessed_resized_image = tf.cast(
        preprocessed_resized_image, tf.bfloat16)
325
326
327
328
    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(
329
      preprocessed_resized_image, axis=0)
330
  out_tensor_dict[input_fields.true_image_shape] = tf.squeeze(
331
      true_image_shape, axis=0)
332
333
  if input_fields.groundtruth_instance_masks in out_tensor_dict:
    masks = out_tensor_dict[input_fields.groundtruth_instance_masks]
334
    _, resized_masks, _ = image_resizer_fn(image, masks)
335
336
    if use_bfloat16:
      resized_masks = tf.cast(resized_masks, tf.bfloat16)
pkulzc's avatar
pkulzc committed
337
    out_tensor_dict[
338
        input_fields.groundtruth_instance_masks] = resized_masks
339

pkulzc's avatar
pkulzc committed
340
  zero_indexed_groundtruth_classes = out_tensor_dict[
341
      input_fields.groundtruth_classes] - _LABEL_OFFSET
342
  if use_multiclass_scores:
pkulzc's avatar
pkulzc committed
343
    out_tensor_dict[
344
345
        input_fields.groundtruth_classes] = out_tensor_dict[
            input_fields.multiclass_scores]
pkulzc's avatar
pkulzc committed
346
  else:
347
    out_tensor_dict[input_fields.groundtruth_classes] = tf.one_hot(
pkulzc's avatar
pkulzc committed
348
        zero_indexed_groundtruth_classes, num_classes)
349
  out_tensor_dict.pop(input_fields.multiclass_scores, None)
350

351
  if input_fields.groundtruth_confidences in out_tensor_dict:
pkulzc's avatar
pkulzc committed
352
    groundtruth_confidences = out_tensor_dict[
353
        input_fields.groundtruth_confidences]
354
    # Map the confidences to the one-hot encoding of classes
355
    out_tensor_dict[input_fields.groundtruth_confidences] = (
356
        tf.reshape(groundtruth_confidences, [-1, 1]) *
357
        out_tensor_dict[input_fields.groundtruth_classes])
358
359
360
  else:
    groundtruth_confidences = tf.ones_like(
        zero_indexed_groundtruth_classes, dtype=tf.float32)
361
362
    out_tensor_dict[input_fields.groundtruth_confidences] = (
        out_tensor_dict[input_fields.groundtruth_classes])
363

364
  if merge_multiple_boxes:
365
366
    merged_boxes, merged_classes, merged_confidences, _ = (
        util_ops.merge_boxes_with_multiple_labels(
367
            out_tensor_dict[input_fields.groundtruth_boxes],
368
369
370
            zero_indexed_groundtruth_classes,
            groundtruth_confidences,
            num_classes))
371
    merged_classes = tf.cast(merged_classes, tf.float32)
372
373
374
    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] = (
375
        merged_confidences)
376
377
378
  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]
379

pkulzc's avatar
pkulzc committed
380
  return out_tensor_dict
381
382


383
384
385
386
387
def pad_input_data_to_static_shapes(tensor_dict,
                                    max_num_boxes,
                                    num_classes,
                                    spatial_image_shape=None,
                                    max_num_context_features=None,
388
389
                                    context_feature_length=None,
                                    max_dp_points=336):
390
391
  """Pads input tensors to static shapes.

392
393
394
  In case num_additional_channels > 0, we assume that the additional channels
  have already been concatenated to the base image.

395
396
397
398
399
400
401
402
  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.
403
404
405
    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.
406
407
408
409
410
    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.
411
412
413
414
415
416

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

  Raises:
417
    ValueError: If groundtruth classes is neither rank 1 nor rank 2, or if we
418
419
420
      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.
421
422
423
424
425
426
  """
  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

427
  input_fields = fields.InputDataFields
428
  num_additional_channels = 0
429
  if input_fields.image_additional_channels in tensor_dict:
430
    num_additional_channels = shape_utils.get_dim_as_int(tensor_dict[
431
        input_fields.image_additional_channels].shape[2])
432
433
434
435

  # 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
436
  if input_fields.image in tensor_dict:
437
    num_channels = shape_utils.get_dim_as_int(
438
        tensor_dict[input_fields.image].shape[2])
439
440
441
442
443
444

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

445
    if (input_fields.original_image in tensor_dict and
446
        shape_utils.get_dim_as_int(
447
            tensor_dict[input_fields.original_image].shape[2]) ==
448
449
450
451
        num_channels):
      raise ValueError(
          'Image must be already concatenated with additional channels.')

452
  if input_fields.context_features in tensor_dict and (
453
454
455
456
457
      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')

458
  padding_shapes = {
459
460
461
      input_fields.image: [height, width, num_channels],
      input_fields.original_image_spatial_shape: [2],
      input_fields.image_additional_channels: [
462
463
          height, width, num_additional_channels
      ],
464
465
466
467
468
469
470
      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: [
471
472
          max_num_boxes, height, width
      ],
473
474
475
476
477
      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: [
478
479
          max_num_boxes, num_classes
      ],
480
481
482
483
484
485
486
      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],
487
488
  }

489
490
  if input_fields.original_image in tensor_dict:
    padding_shapes[input_fields.original_image] = [
491
        height, width,
492
        shape_utils.get_dim_as_int(tensor_dict[input_fields.
493
                                               original_image].shape[2])
494
    ]
495
  if input_fields.groundtruth_keypoints in tensor_dict:
496
    tensor_shape = (
497
        tensor_dict[input_fields.groundtruth_keypoints].shape)
498
499
500
    padding_shape = [max_num_boxes,
                     shape_utils.get_dim_as_int(tensor_shape[1]),
                     shape_utils.get_dim_as_int(tensor_shape[2])]
501
502
503
    padding_shapes[input_fields.groundtruth_keypoints] = padding_shape
  if input_fields.groundtruth_keypoint_visibilities in tensor_dict:
    tensor_shape = tensor_dict[input_fields.
504
                               groundtruth_keypoint_visibilities].shape
505
    padding_shape = [max_num_boxes, shape_utils.get_dim_as_int(tensor_shape[1])]
506
    padding_shapes[input_fields.
507
508
                   groundtruth_keypoint_visibilities] = padding_shape

509
  if input_fields.groundtruth_keypoint_weights in tensor_dict:
510
    tensor_shape = (
511
        tensor_dict[input_fields.groundtruth_keypoint_weights].shape)
512
    padding_shape = [max_num_boxes, shape_utils.get_dim_as_int(tensor_shape[1])]
513
    padding_shapes[input_fields.
514
                   groundtruth_keypoint_weights] = padding_shape
515
  if input_fields.groundtruth_dp_num_points in tensor_dict:
516
    padding_shapes[
517
        input_fields.groundtruth_dp_num_points] = [max_num_boxes]
518
    padding_shapes[
519
        input_fields.groundtruth_dp_part_ids] = [
520
521
            max_num_boxes, max_dp_points]
    padding_shapes[
522
        input_fields.groundtruth_dp_surface_coords] = [
523
            max_num_boxes, max_dp_points, 4]
524
525
526
527
528
529
530
531
  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:
532
    padding_shapes[
533
        input_fields.groundtruth_not_exhaustive_classes] = [num_classes]
534
535

  # Prepare for ContextRCNN related fields.
536
  if input_fields.context_features in tensor_dict:
537
    padding_shape = [max_num_context_features, context_feature_length]
538
    padding_shapes[input_fields.context_features] = padding_shape
539
540

    tensor_shape = tf.shape(
541
542
543
544
545
546
547
548
        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]
549

550
551
  if input_fields.is_annotated in tensor_dict:
    padding_shapes[input_fields.is_annotated] = []
552

553
554
  padded_tensor_dict = {}
  for tensor_name in tensor_dict:
555
556
    padded_tensor_dict[tensor_name] = shape_utils.pad_or_clip_nd(
        tensor_dict[tensor_name], padding_shapes[tensor_name])
557
558
559

  # Make sure that the number of groundtruth boxes now reflects the
  # padded/clipped tensors.
560
561
  if input_fields.num_groundtruth_boxes in padded_tensor_dict:
    padded_tensor_dict[input_fields.num_groundtruth_boxes] = (
562
        tf.minimum(
563
            padded_tensor_dict[input_fields.num_groundtruth_boxes],
564
            max_num_boxes))
565
566
567
  return padded_tensor_dict


568
569
570
571
572
573
574
575
576
577
578
579
580
581
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(
582
      tf.cast(tensor_dict[fields.InputDataFields.image], dtype=tf.float32), 0)
583
584
585
586
587

  include_instance_masks = (fields.InputDataFields.groundtruth_instance_masks
                            in tensor_dict)
  include_keypoints = (fields.InputDataFields.groundtruth_keypoints
                       in tensor_dict)
588
589
  include_keypoint_visibilities = (
      fields.InputDataFields.groundtruth_keypoint_visibilities in tensor_dict)
590
591
592
593
  include_label_weights = (fields.InputDataFields.groundtruth_weights
                           in tensor_dict)
  include_label_confidences = (fields.InputDataFields.groundtruth_confidences
                               in tensor_dict)
594
595
  include_multiclass_scores = (fields.InputDataFields.multiclass_scores in
                               tensor_dict)
596
597
598
599
  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)
600
601
602
  tensor_dict = preprocessor.preprocess(
      tensor_dict, data_augmentation_options,
      func_arg_map=preprocessor.get_default_func_arg_map(
603
604
          include_label_weights=include_label_weights,
          include_label_confidences=include_label_confidences,
605
          include_multiclass_scores=include_multiclass_scores,
606
          include_instance_masks=include_instance_masks,
607
          include_keypoints=include_keypoints,
608
609
          include_keypoint_visibilities=include_keypoint_visibilities,
          include_dense_pose=include_dense_pose))
610
611
612
613
614
  tensor_dict[fields.InputDataFields.image] = tf.squeeze(
      tensor_dict[fields.InputDataFields.image], axis=0)
  return tensor_dict


615
616
617
618
619
620
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,
621
      fields.InputDataFields.groundtruth_weights,
622
623
624
625
626
627
  ]
  labels_dict = {}
  for key in required_label_keys:
    labels_dict[key] = input_dict[key]

  optional_label_keys = [
628
      fields.InputDataFields.groundtruth_confidences,
629
      fields.InputDataFields.groundtruth_labeled_classes,
630
631
632
633
      fields.InputDataFields.groundtruth_keypoints,
      fields.InputDataFields.groundtruth_instance_masks,
      fields.InputDataFields.groundtruth_area,
      fields.InputDataFields.groundtruth_is_crowd,
634
      fields.InputDataFields.groundtruth_group_of,
635
636
637
      fields.InputDataFields.groundtruth_difficult,
      fields.InputDataFields.groundtruth_keypoint_visibilities,
      fields.InputDataFields.groundtruth_keypoint_weights,
638
639
      fields.InputDataFields.groundtruth_dp_num_points,
      fields.InputDataFields.groundtruth_dp_part_ids,
640
      fields.InputDataFields.groundtruth_dp_surface_coords,
641
642
643
      fields.InputDataFields.groundtruth_track_ids,
      fields.InputDataFields.groundtruth_verified_neg_classes,
      fields.InputDataFields.groundtruth_not_exhaustive_classes
644
645
646
647
648
649
650
651
652
653
654
  ]

  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


655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
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


684
def _get_features_dict(input_dict, include_source_id=False):
685
  """Extracts features dict from input dict."""
686
687
688
689
690

  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)
691
692
693
694
695
  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
696
697
698
          input_dict[fields.InputDataFields.true_image_shape],
      fields.InputDataFields.original_image_spatial_shape:
          input_dict[fields.InputDataFields.original_image_spatial_shape]
699
  }
700
701
  if include_source_id:
    features[fields.InputDataFields.source_id] = source_id
702
703
704
  if fields.InputDataFields.original_image in input_dict:
    features[fields.InputDataFields.original_image] = input_dict[
        fields.InputDataFields.original_image]
705
706
707
  if fields.InputDataFields.image_additional_channels in input_dict:
    features[fields.InputDataFields.image_additional_channels] = input_dict[
        fields.InputDataFields.image_additional_channels]
708
709
710
711
712
713
  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]
714
715
716
  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])
717
718
719
  return features


720
721
def create_train_input_fn(train_config, train_input_config,
                          model_config):
722
723
724
725
726
  """Creates a train `input` function for `Estimator`.

  Args:
    train_config: A train_pb2.TrainConfig.
    train_input_config: An input_reader_pb2.InputReader.
727
    model_config: A model_pb2.DetectionModel.
728
729
730
731
732

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

733
  def _train_input_fn(params=None):
734
735
    return train_input(train_config, train_input_config, model_config,
                       params=params)
736

737
  return _train_input_fn
738

739

740
def train_input(train_config, train_input_config,
741
                model_config, model=None, params=None, input_context=None):
742
743
744
745
746
747
748
749
750
  """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.
751
752
753
    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.
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786

  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.
787
788
789
790
791
792
      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.
793
794
      labels[fields.InputDataFields.groundtruth_labeled_classes] is a
        [batch_size, num_classes] float32 k-hot tensor of classes.
795
796
797
798
799
800
801
802
803
804
805
      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.
806
807
      labels[fields.InputDataFields.groundtruth_track_ids] is a
        [batch_size, num_boxes] int32 tensor with the track ID for each object.
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828

  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

829
830
  num_classes = config_util.get_number_of_classes(model_config)

831
832
833
834
835
836
837
838
839
840
841
842
  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)
843
    keypoint_type_weight = train_input_config.keypoint_type_weight or None
844
845
846
    transform_data_fn = functools.partial(
        transform_input_data, model_preprocess_fn=model_preprocess_fn,
        image_resizer_fn=image_resizer_fn,
847
        num_classes=num_classes,
848
849
850
851
        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,
852
853
        use_bfloat16=train_config.use_bfloat16,
        keypoint_type_weight=keypoint_type_weight)
854
855
856
857

    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,
858
        num_classes=num_classes,
859
        spatial_image_shape=config_util.get_spatial_image_size(
860
861
862
863
864
865
866
867
            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))
868
  reduce_to_frame_fn = get_reduce_to_frame_fn(train_input_config, True)
869
870
871
872

  dataset = INPUT_BUILDER_UTIL_MAP['dataset_build'](
      train_input_config,
      transform_input_data_fn=transform_and_pad_input_data_fn,
873
      batch_size=params['batch_size'] if params else train_config.batch_size,
874
875
      input_context=input_context,
      reduce_to_frame_fn=reduce_to_frame_fn)
876
  return dataset
877
878


879
def create_eval_input_fn(eval_config, eval_input_config, model_config):
880
881
882
883
884
  """Creates an eval `input` function for `Estimator`.

  Args:
    eval_config: An eval_pb2.EvalConfig.
    eval_input_config: An input_reader_pb2.InputReader.
885
    model_config: A model_pb2.DetectionModel.
886
887
888
889
890

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

891
  def _eval_input_fn(params=None):
892
893
    return eval_input(eval_config, eval_input_config, model_config,
                      params=params)
894

895
  return _eval_input_fn
896

897

898
def eval_input(eval_config, eval_input_config, model_config,
899
               model=None, params=None, input_context=None):
900
901
902
903
904
905
906
907
908
  """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.
909
910
911
    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.
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940

  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.
941
942
943
944
945
946
      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.
947
948
949
950
951
      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.
952
953
954
955
956
957
958
959
960
961
962
      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.
963
964
      labels[fields.InputDataFields.groundtruth_track_ids] is a
        [batch_size, num_boxes] int32 tensor with the track ID for each object.
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980

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

981
982
983
984
985
986
987
988
  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)

989
990
991
992
993
994
995
996
997
998
999
1000
  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)
1001
    keypoint_type_weight = eval_input_config.keypoint_type_weight or None
1002
1003
1004
1005
1006
1007

    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,
1008
1009
        retain_original_image=eval_config.retain_original_images,
        retain_original_image_additional_channels=
1010
1011
        eval_config.retain_original_image_additional_channels,
        keypoint_type_weight=keypoint_type_weight)
1012
1013
1014
1015
1016
    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(
1017
1018
1019
1020
1021
1022
1023
1024
            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))
1025
1026
1027

  reduce_to_frame_fn = get_reduce_to_frame_fn(eval_input_config, False)

1028
1029
1030
  dataset = INPUT_BUILDER_UTIL_MAP['dataset_build'](
      eval_input_config,
      batch_size=params['batch_size'] if params else eval_config.batch_size,
1031
      transform_input_data_fn=transform_and_pad_input_data_fn,
1032
      input_context=input_context,
1033
      reduce_to_frame_fn=reduce_to_frame_fn)
1034
  return dataset
1035
1036


1037
def create_predict_input_fn(model_config, predict_input_config):
1038
1039
  """Creates a predict `input` function for `Estimator`.

1040
1041
  Args:
    model_config: A model_pb2.DetectionModel.
1042
    predict_input_config: An input_reader_pb2.InputReader.
1043

1044
1045
1046
1047
  Returns:
    `input_fn` for `Estimator` in PREDICT mode.
  """

1048
  def _predict_input_fn(params=None):
1049
1050
    """Decodes serialized tf.Examples and returns `ServingInputReceiver`.

1051
1052
1053
    Args:
      params: Parameter dictionary passed from the estimator.

1054
1055
1056
    Returns:
      `ServingInputReceiver`.
    """
1057
    del params
1058
    example = tf.placeholder(dtype=tf.string, shape=[], name='tf_example')
1059

1060
    num_classes = config_util.get_number_of_classes(model_config)
1061
1062
1063
    model_preprocess_fn = INPUT_BUILDER_UTIL_MAP['model_build'](
        model_config, is_training=False).preprocess

1064
1065
    image_resizer_config = config_util.get_image_resizer_config(model_config)
    image_resizer_fn = image_resizer_builder.build(image_resizer_config)
1066

1067
    transform_fn = functools.partial(
1068
        transform_input_data, model_preprocess_fn=model_preprocess_fn,
1069
1070
1071
1072
        image_resizer_fn=image_resizer_fn,
        num_classes=num_classes,
        data_augmentation_fn=None)

1073
1074
1075
    decoder = tf_example_decoder.TfExampleDecoder(
        load_instance_masks=False,
        num_additional_channels=predict_input_config.num_additional_channels)
1076
    input_dict = transform_fn(decoder.decode(example))
1077
    images = tf.cast(input_dict[fields.InputDataFields.image], dtype=tf.float32)
1078
    images = tf.expand_dims(images, axis=0)
1079
1080
    true_image_shape = tf.expand_dims(
        input_dict[fields.InputDataFields.true_image_shape], axis=0)
1081
1082

    return tf.estimator.export.ServingInputReceiver(
1083
1084
1085
        features={
            fields.InputDataFields.image: images,
            fields.InputDataFields.true_image_shape: true_image_shape},
1086
1087
1088
        receiver_tensors={SERVING_FED_EXAMPLE_KEY: example})

  return _predict_input_fn
1089
1090
1091
1092
1093
1094
1095
1096
1097
1098
1099
1100
1101
1102
1103
1104
1105
1106
1107
1108
1109


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 != (
1110
1111
      input_reader_pb2.InputType.Value('TF_SEQUENCE_EXAMPLE')):
    return lambda dataset, dataset_map_fn, batch_size, config: dataset
1112
  else:
1113
1114
    def reduce_to_frame(dataset, dataset_map_fn, batch_size,
                        input_reader_config):
1115
1116
1117
1118
      """Returns a function reducing sequence tensors to single frame tensors.

      Args:
        dataset: A tf dataset containing sequence tensors.
1119
1120
1121
1122
1123
1124
        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
1125
1126
1127
1128
1129
1130
1131
1132
1133
1134
1135
1136
1137
1138
1139
1140
1141
1142
1143
1144
1145

      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)
1146
1147
1148
1149
1150
1151
          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)
1152
1153
1154
1155
1156
1157
1158
1159
1160
          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
1161
1162
        dataset = dataset_map_fn(dataset, get_single_frame, batch_size,
                                 input_reader_config)
1163
      else:
1164
1165
        dataset = dataset_map_fn(dataset, util_ops.tile_context_tensors,
                                 batch_size, input_reader_config)
1166
1167
        dataset = dataset.unbatch()
      # Decode frame here as SequenceExample tensors contain encoded images.
1168
1169
      dataset = dataset_map_fn(dataset, util_ops.decode_image, batch_size,
                               input_reader_config)
1170
1171
      return dataset
    return reduce_to_frame