inputs.py 41.6 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Model input function for tf-learn object detection model."""

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

import functools

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

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

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

54

pkulzc's avatar
pkulzc committed
55
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
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
def _convert_labeled_classes_to_k_hot(groundtruth_labeled_classes, num_classes):
  """Returns k-hot encoding of the labeled classes."""

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

  return tf.cond(tf.size(groundtruth_labeled_classes) > 0, true_fn, false_fn)


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

  recognized_indices = tf.where(tf.greater(class_ids, unrecognized_label))
  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()


129
130
131
132
133
134
def transform_input_data(tensor_dict,
                         model_preprocess_fn,
                         image_resizer_fn,
                         num_classes,
                         data_augmentation_fn=None,
                         merge_multiple_boxes=False,
135
                         retain_original_image=False,
136
                         use_multiclass_scores=False,
137
                         use_bfloat16=False,
138
139
                         retain_original_image_additional_channels=False,
                         keypoint_type_weight=None):
140
141
142
  """A single function that is responsible for all input data transformations.

  Data transformation functions are applied in the following order.
143
144
145
146
147
  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.
148
149
150
151
152
153
  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
154
     tensor_dict.
155
156
  6. one_hot_encoding: applied to classes tensor in tensor_dict.
  7. merge_multiple_boxes (optional): when groundtruth boxes are exactly the
157
158
159
160
161
162
163
164
165
     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.
166
167
168
169
    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.
170
171
172
173
174
175
176
177
    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
178
179
180
181
    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.
182
    use_bfloat16: (optional) a bool, whether to use bfloat16 in training.
183
184
    retain_original_image_additional_channels: (optional) Whether to retain
      original image additional channels in the output dictionary.
185
186
187
    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.
188
189
190
191
192

  Returns:
    A dictionary keyed by fields.InputDataFields containing the tensors obtained
    after applying all the transformations.
  """
pkulzc's avatar
pkulzc committed
193
  out_tensor_dict = tensor_dict.copy()
194
195
196
197
198
199
200
201
202
203

  labeled_classes_field = fields.InputDataFields.groundtruth_labeled_classes
  if labeled_classes_field in out_tensor_dict:
    # tf_example_decoder casts unrecognized labels to -1. Remove these
    # unrecognized labels before converting labeled_classes to k-hot vector.
    out_tensor_dict[labeled_classes_field] = _remove_unrecognized_classes(
        out_tensor_dict[labeled_classes_field], unrecognized_label=-1)
    out_tensor_dict[labeled_classes_field] = _convert_labeled_classes_to_k_hot(
        out_tensor_dict[labeled_classes_field], num_classes)

pkulzc's avatar
pkulzc committed
204
205
206
207
208
209
210
211
212
213
214
215
216
  if fields.InputDataFields.multiclass_scores in out_tensor_dict:
    out_tensor_dict[
        fields.InputDataFields
        .multiclass_scores] = _multiclass_scores_or_one_hot_labels(
            out_tensor_dict[fields.InputDataFields.multiclass_scores],
            out_tensor_dict[fields.InputDataFields.groundtruth_boxes],
            out_tensor_dict[fields.InputDataFields.groundtruth_classes],
            num_classes)

  if fields.InputDataFields.groundtruth_boxes in out_tensor_dict:
    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)
217

218
  if retain_original_image:
pkulzc's avatar
pkulzc committed
219
220
221
    out_tensor_dict[fields.InputDataFields.original_image] = tf.cast(
        image_resizer_fn(out_tensor_dict[fields.InputDataFields.image],
                         None)[0], tf.uint8)
222

pkulzc's avatar
pkulzc committed
223
224
225
226
  if fields.InputDataFields.image_additional_channels in out_tensor_dict:
    channels = out_tensor_dict[fields.InputDataFields.image_additional_channels]
    out_tensor_dict[fields.InputDataFields.image] = tf.concat(
        [out_tensor_dict[fields.InputDataFields.image], channels], axis=2)
227
228
229
230
    if retain_original_image_additional_channels:
      out_tensor_dict[
          fields.InputDataFields.image_additional_channels] = tf.cast(
              image_resizer_fn(channels, None)[0], tf.uint8)
231

232
233
  # Apply data augmentation ops.
  if data_augmentation_fn is not None:
pkulzc's avatar
pkulzc committed
234
    out_tensor_dict = data_augmentation_fn(out_tensor_dict)
235
236

  # Apply model preprocessing ops and resize instance masks.
pkulzc's avatar
pkulzc committed
237
  image = out_tensor_dict[fields.InputDataFields.image]
238
  preprocessed_resized_image, true_image_shape = model_preprocess_fn(
239
      tf.expand_dims(tf.cast(image, dtype=tf.float32), axis=0))
240
241
242
243
244
245
246
247
248
249
250
251
252
253

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

  if fields.InputDataFields.groundtruth_boxes in tensor_dict:
    bboxes = out_tensor_dict[fields.InputDataFields.groundtruth_boxes]
    boxlist = box_list.BoxList(bboxes)
    realigned_bboxes = box_list_ops.change_coordinate_frame(boxlist, im_box)
254
255
256

    realigned_boxes_tensor = realigned_bboxes.get()
    valid_boxes_tensor = assert_or_prune_invalid_boxes(realigned_boxes_tensor)
257
    out_tensor_dict[
258
        fields.InputDataFields.groundtruth_boxes] = valid_boxes_tensor
259
260
261
262
263
264
265

  if fields.InputDataFields.groundtruth_keypoints in tensor_dict:
    keypoints = out_tensor_dict[fields.InputDataFields.groundtruth_keypoints]
    realigned_keypoints = keypoint_ops.change_coordinate_frame(keypoints,
                                                               im_box)
    out_tensor_dict[
        fields.InputDataFields.groundtruth_keypoints] = realigned_keypoints
266
267
268
269
270
271
272
273
274
275
276
    flds_gt_kpt = fields.InputDataFields.groundtruth_keypoints
    flds_gt_kpt_vis = fields.InputDataFields.groundtruth_keypoint_visibilities
    flds_gt_kpt_weights = fields.InputDataFields.groundtruth_keypoint_weights
    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))
277

278
279
280
  if use_bfloat16:
    preprocessed_resized_image = tf.cast(
        preprocessed_resized_image, tf.bfloat16)
281
282
283
    if fields.InputDataFields.context_features in out_tensor_dict:
      out_tensor_dict[fields.InputDataFields.context_features] = tf.cast(
          out_tensor_dict[fields.InputDataFields.context_features], tf.bfloat16)
pkulzc's avatar
pkulzc committed
284
  out_tensor_dict[fields.InputDataFields.image] = tf.squeeze(
285
      preprocessed_resized_image, axis=0)
pkulzc's avatar
pkulzc committed
286
  out_tensor_dict[fields.InputDataFields.true_image_shape] = tf.squeeze(
287
      true_image_shape, axis=0)
pkulzc's avatar
pkulzc committed
288
289
  if fields.InputDataFields.groundtruth_instance_masks in out_tensor_dict:
    masks = out_tensor_dict[fields.InputDataFields.groundtruth_instance_masks]
290
    _, resized_masks, _ = image_resizer_fn(image, masks)
291
292
    if use_bfloat16:
      resized_masks = tf.cast(resized_masks, tf.bfloat16)
pkulzc's avatar
pkulzc committed
293
294
    out_tensor_dict[
        fields.InputDataFields.groundtruth_instance_masks] = resized_masks
295

pkulzc's avatar
pkulzc committed
296
  zero_indexed_groundtruth_classes = out_tensor_dict[
297
      fields.InputDataFields.groundtruth_classes] - _LABEL_OFFSET
298
  if use_multiclass_scores:
pkulzc's avatar
pkulzc committed
299
300
301
302
303
304
305
    out_tensor_dict[
        fields.InputDataFields.groundtruth_classes] = out_tensor_dict[
            fields.InputDataFields.multiclass_scores]
  else:
    out_tensor_dict[fields.InputDataFields.groundtruth_classes] = tf.one_hot(
        zero_indexed_groundtruth_classes, num_classes)
  out_tensor_dict.pop(fields.InputDataFields.multiclass_scores, None)
306

pkulzc's avatar
pkulzc committed
307
308
  if fields.InputDataFields.groundtruth_confidences in out_tensor_dict:
    groundtruth_confidences = out_tensor_dict[
309
        fields.InputDataFields.groundtruth_confidences]
310
    # Map the confidences to the one-hot encoding of classes
pkulzc's avatar
pkulzc committed
311
    out_tensor_dict[fields.InputDataFields.groundtruth_confidences] = (
312
        tf.reshape(groundtruth_confidences, [-1, 1]) *
pkulzc's avatar
pkulzc committed
313
        out_tensor_dict[fields.InputDataFields.groundtruth_classes])
314
315
316
  else:
    groundtruth_confidences = tf.ones_like(
        zero_indexed_groundtruth_classes, dtype=tf.float32)
pkulzc's avatar
pkulzc committed
317
318
    out_tensor_dict[fields.InputDataFields.groundtruth_confidences] = (
        out_tensor_dict[fields.InputDataFields.groundtruth_classes])
319

320
  if merge_multiple_boxes:
321
322
    merged_boxes, merged_classes, merged_confidences, _ = (
        util_ops.merge_boxes_with_multiple_labels(
pkulzc's avatar
pkulzc committed
323
            out_tensor_dict[fields.InputDataFields.groundtruth_boxes],
324
325
326
            zero_indexed_groundtruth_classes,
            groundtruth_confidences,
            num_classes))
327
    merged_classes = tf.cast(merged_classes, tf.float32)
pkulzc's avatar
pkulzc committed
328
329
330
    out_tensor_dict[fields.InputDataFields.groundtruth_boxes] = merged_boxes
    out_tensor_dict[fields.InputDataFields.groundtruth_classes] = merged_classes
    out_tensor_dict[fields.InputDataFields.groundtruth_confidences] = (
331
        merged_confidences)
pkulzc's avatar
pkulzc committed
332
333
334
  if fields.InputDataFields.groundtruth_boxes in out_tensor_dict:
    out_tensor_dict[fields.InputDataFields.num_groundtruth_boxes] = tf.shape(
        out_tensor_dict[fields.InputDataFields.groundtruth_boxes])[0]
335

pkulzc's avatar
pkulzc committed
336
  return out_tensor_dict
337
338


339
340
341
342
343
344
def pad_input_data_to_static_shapes(tensor_dict,
                                    max_num_boxes,
                                    num_classes,
                                    spatial_image_shape=None,
                                    max_num_context_features=None,
                                    context_feature_length=None):
345
346
  """Pads input tensors to static shapes.

347
348
349
  In case num_additional_channels > 0, we assume that the additional channels
  have already been concatenated to the base image.

350
351
352
353
354
355
356
357
  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.
358
359
360
    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.
361
362
363
364
365
366

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

  Raises:
367
    ValueError: If groundtruth classes is neither rank 1 nor rank 2, or if we
368
369
370
      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.
371
372
373
374
375
376
377
378
379
  """

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

  num_additional_channels = 0
  if fields.InputDataFields.image_additional_channels in tensor_dict:
380
381
    num_additional_channels = shape_utils.get_dim_as_int(tensor_dict[
        fields.InputDataFields.image_additional_channels].shape[2])
382
383
384
385

  # 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
386
  if fields.InputDataFields.image in tensor_dict:
387
388
    num_channels = shape_utils.get_dim_as_int(
        tensor_dict[fields.InputDataFields.image].shape[2])
389
390
391
392
393
394
395

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

    if (fields.InputDataFields.original_image in tensor_dict and
396
397
        shape_utils.get_dim_as_int(
            tensor_dict[fields.InputDataFields.original_image].shape[2]) ==
398
399
400
401
        num_channels):
      raise ValueError(
          'Image must be already concatenated with additional channels.')

402
403
404
405
406
407
  if fields.InputDataFields.context_features in tensor_dict and (
      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')

408
  padding_shapes = {
409
      fields.InputDataFields.image: [height, width, num_channels],
pkulzc's avatar
pkulzc committed
410
      fields.InputDataFields.original_image_spatial_shape: [2],
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
      fields.InputDataFields.image_additional_channels: [
          height, width, num_additional_channels
      ],
      fields.InputDataFields.source_id: [],
      fields.InputDataFields.filename: [],
      fields.InputDataFields.key: [],
      fields.InputDataFields.groundtruth_difficult: [max_num_boxes],
      fields.InputDataFields.groundtruth_boxes: [max_num_boxes, 4],
      fields.InputDataFields.groundtruth_classes: [max_num_boxes, num_classes],
      fields.InputDataFields.groundtruth_instance_masks: [
          max_num_boxes, height, width
      ],
      fields.InputDataFields.groundtruth_is_crowd: [max_num_boxes],
      fields.InputDataFields.groundtruth_group_of: [max_num_boxes],
      fields.InputDataFields.groundtruth_area: [max_num_boxes],
      fields.InputDataFields.groundtruth_weights: [max_num_boxes],
427
428
429
      fields.InputDataFields.groundtruth_confidences: [
          max_num_boxes, num_classes
      ],
430
431
      fields.InputDataFields.num_groundtruth_boxes: [],
      fields.InputDataFields.groundtruth_label_types: [max_num_boxes],
432
      fields.InputDataFields.groundtruth_label_weights: [max_num_boxes],
433
434
      fields.InputDataFields.true_image_shape: [3],
      fields.InputDataFields.groundtruth_image_classes: [num_classes],
435
      fields.InputDataFields.groundtruth_image_confidences: [num_classes],
436
      fields.InputDataFields.groundtruth_labeled_classes: [num_classes],
437
438
439
440
  }

  if fields.InputDataFields.original_image in tensor_dict:
    padding_shapes[fields.InputDataFields.original_image] = [
441
442
443
        height, width,
        shape_utils.get_dim_as_int(tensor_dict[fields.InputDataFields.
                                               original_image].shape[2])
444
445
446
447
    ]
  if fields.InputDataFields.groundtruth_keypoints in tensor_dict:
    tensor_shape = (
        tensor_dict[fields.InputDataFields.groundtruth_keypoints].shape)
448
449
450
    padding_shape = [max_num_boxes,
                     shape_utils.get_dim_as_int(tensor_shape[1]),
                     shape_utils.get_dim_as_int(tensor_shape[2])]
451
452
453
454
    padding_shapes[fields.InputDataFields.groundtruth_keypoints] = padding_shape
  if fields.InputDataFields.groundtruth_keypoint_visibilities in tensor_dict:
    tensor_shape = tensor_dict[fields.InputDataFields.
                               groundtruth_keypoint_visibilities].shape
455
    padding_shape = [max_num_boxes, shape_utils.get_dim_as_int(tensor_shape[1])]
456
457
458
    padding_shapes[fields.InputDataFields.
                   groundtruth_keypoint_visibilities] = padding_shape

459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
  if fields.InputDataFields.groundtruth_keypoint_weights in tensor_dict:
    tensor_shape = (
        tensor_dict[fields.InputDataFields.groundtruth_keypoint_weights].shape)
    padding_shape = [max_num_boxes, shape_utils.get_dim_as_int(tensor_shape[1])]
    padding_shapes[fields.InputDataFields.
                   groundtruth_keypoint_weights] = padding_shape

  # Prepare for ContextRCNN related fields.
  if fields.InputDataFields.context_features in tensor_dict:
    padding_shape = [max_num_context_features, context_feature_length]
    padding_shapes[fields.InputDataFields.context_features] = padding_shape

    tensor_shape = tf.shape(
        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] = []

478
479
  padded_tensor_dict = {}
  for tensor_name in tensor_dict:
480
481
    padded_tensor_dict[tensor_name] = shape_utils.pad_or_clip_nd(
        tensor_dict[tensor_name], padding_shapes[tensor_name])
482
483
484
485
486
487
488
489

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


493
494
495
496
497
498
499
500
501
502
503
504
505
506
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(
507
      tf.cast(tensor_dict[fields.InputDataFields.image], dtype=tf.float32), 0)
508
509
510
511
512

  include_instance_masks = (fields.InputDataFields.groundtruth_instance_masks
                            in tensor_dict)
  include_keypoints = (fields.InputDataFields.groundtruth_keypoints
                       in tensor_dict)
513
514
  include_keypoint_visibilities = (
      fields.InputDataFields.groundtruth_keypoint_visibilities in tensor_dict)
515
516
517
518
  include_label_weights = (fields.InputDataFields.groundtruth_weights
                           in tensor_dict)
  include_label_confidences = (fields.InputDataFields.groundtruth_confidences
                               in tensor_dict)
519
520
  include_multiclass_scores = (fields.InputDataFields.multiclass_scores in
                               tensor_dict)
521
522
523
  tensor_dict = preprocessor.preprocess(
      tensor_dict, data_augmentation_options,
      func_arg_map=preprocessor.get_default_func_arg_map(
524
525
          include_label_weights=include_label_weights,
          include_label_confidences=include_label_confidences,
526
          include_multiclass_scores=include_multiclass_scores,
527
          include_instance_masks=include_instance_masks,
528
529
          include_keypoints=include_keypoints,
          include_keypoint_visibilities=include_keypoint_visibilities))
530
531
532
533
534
  tensor_dict[fields.InputDataFields.image] = tf.squeeze(
      tensor_dict[fields.InputDataFields.image], axis=0)
  return tensor_dict


535
536
537
538
539
540
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,
541
      fields.InputDataFields.groundtruth_weights,
542
543
544
545
546
547
  ]
  labels_dict = {}
  for key in required_label_keys:
    labels_dict[key] = input_dict[key]

  optional_label_keys = [
548
      fields.InputDataFields.groundtruth_confidences,
549
      fields.InputDataFields.groundtruth_labeled_classes,
550
551
552
553
      fields.InputDataFields.groundtruth_keypoints,
      fields.InputDataFields.groundtruth_instance_masks,
      fields.InputDataFields.groundtruth_area,
      fields.InputDataFields.groundtruth_is_crowd,
554
555
556
      fields.InputDataFields.groundtruth_difficult,
      fields.InputDataFields.groundtruth_keypoint_visibilities,
      fields.InputDataFields.groundtruth_keypoint_weights,
557
558
559
560
561
562
563
564
565
566
567
  ]

  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


568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
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


597
def _get_features_dict(input_dict, include_source_id=False):
598
  """Extracts features dict from input dict."""
599
600
601
602
603

  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)
604
605
606
607
608
  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
609
610
611
          input_dict[fields.InputDataFields.true_image_shape],
      fields.InputDataFields.original_image_spatial_shape:
          input_dict[fields.InputDataFields.original_image_spatial_shape]
612
  }
613
614
  if include_source_id:
    features[fields.InputDataFields.source_id] = source_id
615
616
617
  if fields.InputDataFields.original_image in input_dict:
    features[fields.InputDataFields.original_image] = input_dict[
        fields.InputDataFields.original_image]
618
619
620
  if fields.InputDataFields.image_additional_channels in input_dict:
    features[fields.InputDataFields.image_additional_channels] = input_dict[
        fields.InputDataFields.image_additional_channels]
621
622
623
624
625
626
  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]
627
628
629
  return features


630
631
def create_train_input_fn(train_config, train_input_config,
                          model_config):
632
633
634
635
636
  """Creates a train `input` function for `Estimator`.

  Args:
    train_config: A train_pb2.TrainConfig.
    train_input_config: An input_reader_pb2.InputReader.
637
    model_config: A model_pb2.DetectionModel.
638
639
640
641
642

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

643
  def _train_input_fn(params=None):
644
645
    return train_input(train_config, train_input_config, model_config,
                       params=params)
646

647
  return _train_input_fn
648

649

650
def train_input(train_config, train_input_config,
651
                model_config, model=None, params=None, input_context=None):
652
653
654
655
656
657
658
659
660
  """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.
661
662
663
    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.
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696

  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.
697
698
699
700
701
702
      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.
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723

  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

724
725
  num_classes = config_util.get_number_of_classes(model_config)

726
727
728
729
730
731
732
733
734
735
736
737
  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)
738
    keypoint_type_weight = train_input_config.keypoint_type_weight or None
739
740
741
    transform_data_fn = functools.partial(
        transform_input_data, model_preprocess_fn=model_preprocess_fn,
        image_resizer_fn=image_resizer_fn,
742
        num_classes=num_classes,
743
744
745
746
        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,
747
748
        use_bfloat16=train_config.use_bfloat16,
        keypoint_type_weight=keypoint_type_weight)
749
750
751
752

    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,
753
        num_classes=num_classes,
754
        spatial_image_shape=config_util.get_spatial_image_size(
755
756
757
758
759
760
761
762
            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))
763
764
765
766

  dataset = INPUT_BUILDER_UTIL_MAP['dataset_build'](
      train_input_config,
      transform_input_data_fn=transform_and_pad_input_data_fn,
767
768
      batch_size=params['batch_size'] if params else train_config.batch_size,
      input_context=input_context)
769
  return dataset
770
771


772
def create_eval_input_fn(eval_config, eval_input_config, model_config):
773
774
775
776
777
  """Creates an eval `input` function for `Estimator`.

  Args:
    eval_config: An eval_pb2.EvalConfig.
    eval_input_config: An input_reader_pb2.InputReader.
778
    model_config: A model_pb2.DetectionModel.
779
780
781
782
783

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

784
  def _eval_input_fn(params=None):
785
786
    return eval_input(eval_config, eval_input_config, model_config,
                      params=params)
787

788
  return _eval_input_fn
789

790

791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
def eval_input(eval_config, eval_input_config, model_config,
               model=None, params=None):
  """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.

  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.
831
832
833
834
835
836
      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.
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852

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

853
854
855
856
857
858
859
860
  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)

861
862
863
864
865
866
867
868
869
870
871
872
  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)
873
    keypoint_type_weight = eval_input_config.keypoint_type_weight or None
874
875
876
877
878
879

    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,
880
881
        retain_original_image=eval_config.retain_original_images,
        retain_original_image_additional_channels=
882
883
        eval_config.retain_original_image_additional_channels,
        keypoint_type_weight=keypoint_type_weight)
884
885
886
887
888
    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(
889
890
891
892
893
894
895
896
            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))
897
898
899
900
901
  dataset = INPUT_BUILDER_UTIL_MAP['dataset_build'](
      eval_input_config,
      batch_size=params['batch_size'] if params else eval_config.batch_size,
      transform_input_data_fn=transform_and_pad_input_data_fn)
  return dataset
902
903


904
def create_predict_input_fn(model_config, predict_input_config):
905
906
  """Creates a predict `input` function for `Estimator`.

907
908
  Args:
    model_config: A model_pb2.DetectionModel.
909
    predict_input_config: An input_reader_pb2.InputReader.
910

911
912
913
914
  Returns:
    `input_fn` for `Estimator` in PREDICT mode.
  """

915
  def _predict_input_fn(params=None):
916
917
    """Decodes serialized tf.Examples and returns `ServingInputReceiver`.

918
919
920
    Args:
      params: Parameter dictionary passed from the estimator.

921
922
923
    Returns:
      `ServingInputReceiver`.
    """
924
    del params
925
    example = tf.placeholder(dtype=tf.string, shape=[], name='tf_example')
926

927
    num_classes = config_util.get_number_of_classes(model_config)
928
929
930
    model_preprocess_fn = INPUT_BUILDER_UTIL_MAP['model_build'](
        model_config, is_training=False).preprocess

931
932
    image_resizer_config = config_util.get_image_resizer_config(model_config)
    image_resizer_fn = image_resizer_builder.build(image_resizer_config)
933

934
    transform_fn = functools.partial(
935
        transform_input_data, model_preprocess_fn=model_preprocess_fn,
936
937
938
939
        image_resizer_fn=image_resizer_fn,
        num_classes=num_classes,
        data_augmentation_fn=None)

940
941
942
    decoder = tf_example_decoder.TfExampleDecoder(
        load_instance_masks=False,
        num_additional_channels=predict_input_config.num_additional_channels)
943
    input_dict = transform_fn(decoder.decode(example))
944
    images = tf.cast(input_dict[fields.InputDataFields.image], dtype=tf.float32)
945
    images = tf.expand_dims(images, axis=0)
946
947
    true_image_shape = tf.expand_dims(
        input_dict[fields.InputDataFields.true_image_shape], axis=0)
948
949

    return tf.estimator.export.ServingInputReceiver(
950
951
952
        features={
            fields.InputDataFields.image: images,
            fields.InputDataFields.true_image_shape: true_image_shape},
953
954
955
        receiver_tensors={SERVING_FED_EXAMPLE_KEY: example})

  return _predict_input_fn