retinanet_input.py 13.4 KB
Newer Older
Yeqing Li's avatar
Yeqing Li committed
1
# Copyright 2021 The TensorFlow Authors. All Rights Reserved.
Abdullah Rashwan's avatar
Abdullah Rashwan committed
2
3
4
5
6
7
8
9
10
11
12
13
#
# 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.
Yeqing Li's avatar
Yeqing Li committed
14

Abdullah Rashwan's avatar
Abdullah Rashwan committed
15
16
17
18
19
20
21
"""Data parser and processing for RetinaNet.

Parse image and ground truths in a dataset to training targets and package them
into (image, labels) tuple for RetinaNet.
"""

# Import libraries
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
22
from absl import logging
Abdullah Rashwan's avatar
Abdullah Rashwan committed
23
24
25
26
27
import tensorflow as tf

from official.vision.beta.dataloaders import parser
from official.vision.beta.dataloaders import utils
from official.vision.beta.ops import anchor
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
28
from official.vision.beta.ops import augment
Abdullah Rashwan's avatar
Abdullah Rashwan committed
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
from official.vision.beta.ops import box_ops
from official.vision.beta.ops import preprocess_ops


class Parser(parser.Parser):
  """Parser to parse an image and its annotations into a dictionary of tensors."""

  def __init__(self,
               output_size,
               min_level,
               max_level,
               num_scales,
               aspect_ratios,
               anchor_size,
               match_threshold=0.5,
               unmatched_threshold=0.5,
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
45
               aug_type=None,
Abdullah Rashwan's avatar
Abdullah Rashwan committed
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
               aug_rand_hflip=False,
               aug_scale_min=1.0,
               aug_scale_max=1.0,
               use_autoaugment=False,
               autoaugment_policy_name='v0',
               skip_crowd_during_training=True,
               max_num_instances=100,
               dtype='bfloat16',
               mode=None):
    """Initializes parameters for parsing annotations in the dataset.

    Args:
      output_size: `Tensor` or `list` for [height, width] of output image. The
        output_size should be divided by the largest feature stride 2^max_level.
      min_level: `int` number of minimum level of the output feature pyramid.
      max_level: `int` number of maximum level of the output feature pyramid.
      num_scales: `int` number representing intermediate scales added on each
        level. For instances, num_scales=2 adds one additional intermediate
        anchor scales [2^0, 2^0.5] on each level.
      aspect_ratios: `list` of float numbers representing the aspect raito
        anchors added on each level. The number indicates the ratio of width to
        height. For instances, aspect_ratios=[1.0, 2.0, 0.5] adds three anchors
        on each scale level.
      anchor_size: `float` number representing the scale of size of the base
        anchor to the feature stride 2^level.
      match_threshold: `float` number between 0 and 1 representing the
        lower-bound threshold to assign positive labels for anchors. An anchor
        with a score over the threshold is labeled positive.
      unmatched_threshold: `float` number between 0 and 1 representing the
        upper-bound threshold to assign negative labels for anchors. An anchor
        with a score below the threshold is labeled negative.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
77
      aug_type: An optional Augmentation object to choose from AutoAugment and
78
        RandAugment.
Abdullah Rashwan's avatar
Abdullah Rashwan committed
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
      aug_rand_hflip: `bool`, if True, augment training with random horizontal
        flip.
      aug_scale_min: `float`, the minimum scale applied to `output_size` for
        data augmentation during training.
      aug_scale_max: `float`, the maximum scale applied to `output_size` for
        data augmentation during training.
      use_autoaugment: `bool`, if True, use the AutoAugment augmentation policy
        during training.
      autoaugment_policy_name: `string` that specifies the name of the
        AutoAugment policy that will be used during training.
      skip_crowd_during_training: `bool`, if True, skip annotations labeled with
        `is_crowd` equals to 1.
      max_num_instances: `int` number of maximum number of instances in an
        image. The groundtruth data will be padded to `max_num_instances`.
      dtype: `str`, data type. One of {`bfloat16`, `float32`, `float16`}.
      mode: a ModeKeys. Specifies if this is training, evaluation, prediction or
        prediction with groundtruths in the outputs.
    """
    self._mode = mode
    self._max_num_instances = max_num_instances
    self._skip_crowd_during_training = skip_crowd_during_training

    # Anchor.
    self._output_size = output_size
    self._min_level = min_level
    self._max_level = max_level
    self._num_scales = num_scales
    self._aspect_ratios = aspect_ratios
    self._anchor_size = anchor_size
    self._match_threshold = match_threshold
    self._unmatched_threshold = unmatched_threshold

    # Data augmentation.
    self._aug_rand_hflip = aug_rand_hflip
    self._aug_scale_min = aug_scale_min
    self._aug_scale_max = aug_scale_max

A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
116
117
118
119
120
121
122
123
124
    # Data augmentation with AutoAugment or RandAugment.
    self._augmenter = None
    if aug_type is not None:
      if aug_type.type == 'autoaug':
        logging.info('Using AutoAugment.')
        self._augmenter = augment.AutoAugment(
            augmentation_name=aug_type.autoaug.augmentation_name,
            cutout_const=aug_type.autoaug.cutout_const,
            translate_const=aug_type.autoaug.translate_const)
125
126
127
128
129
130
131
132
133
      elif aug_type.type == 'randaug':
        logging.info('Using RandAugment.')
        self._augmenter = augment.RandAugment.build_for_detection(
            num_layers=aug_type.randaug.num_layers,
            magnitude=aug_type.randaug.magnitude,
            cutout_const=aug_type.randaug.cutout_const,
            translate_const=aug_type.randaug.translate_const,
            prob_to_apply=aug_type.randaug.prob_to_apply,
            exclude_ops=aug_type.randaug.exclude_ops)
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
134
135
136
137
      else:
        raise ValueError(f'Augmentation policy {aug_type.type} not supported.')

    # Deprecated. Data Augmentation with AutoAugment.
Abdullah Rashwan's avatar
Abdullah Rashwan committed
138
139
140
    self._use_autoaugment = use_autoaugment
    self._autoaugment_policy_name = autoaugment_policy_name

A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
141
142
    # Data type.
    self._dtype = dtype
Abdullah Rashwan's avatar
Abdullah Rashwan committed
143
144
145
146
147

  def _parse_train_data(self, data):
    """Parses data for training and evaluation."""
    classes = data['groundtruth_classes']
    boxes = data['groundtruth_boxes']
Xianzhi Du's avatar
Xianzhi Du committed
148
149
150
151
    # If not empty, `attributes` is a dict of (name, ground_truth) pairs.
    # `ground_gruth` of attributes is assumed in shape [N, attribute_size].
    # TODO(xianzhi): support parsing attributes weights.
    attributes = data.get('groundtruth_attributes', {})
Abdullah Rashwan's avatar
Abdullah Rashwan committed
152
    is_crowds = data['groundtruth_is_crowd']
Xianzhi Du's avatar
Xianzhi Du committed
153

Abdullah Rashwan's avatar
Abdullah Rashwan committed
154
155
156
157
158
159
160
161
162
163
    # Skips annotations with `is_crowd` = True.
    if self._skip_crowd_during_training:
      num_groundtrtuhs = tf.shape(input=classes)[0]
      with tf.control_dependencies([num_groundtrtuhs, is_crowds]):
        indices = tf.cond(
            pred=tf.greater(tf.size(input=is_crowds), 0),
            true_fn=lambda: tf.where(tf.logical_not(is_crowds))[:, 0],
            false_fn=lambda: tf.cast(tf.range(num_groundtrtuhs), tf.int64))
      classes = tf.gather(classes, indices)
      boxes = tf.gather(boxes, indices)
Xianzhi Du's avatar
Xianzhi Du committed
164
165
      for k, v in attributes.items():
        attributes[k] = tf.gather(v, indices)
Abdullah Rashwan's avatar
Abdullah Rashwan committed
166

A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
167
    # Gets original image.
Abdullah Rashwan's avatar
Abdullah Rashwan committed
168
169
    image = data['image']

A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
170
171
172
    # Apply autoaug or randaug.
    if self._augmenter is not None:
      image, boxes = self._augmenter.distort_with_boxes(image, boxes)
Abdullah Rashwan's avatar
Abdullah Rashwan committed
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
    image_shape = tf.shape(input=image)[0:2]

    # Normalizes image with mean and std pixel values.
    image = preprocess_ops.normalize_image(image)

    # Flips image randomly during training.
    if self._aug_rand_hflip:
      image, boxes, _ = preprocess_ops.random_horizontal_flip(image, boxes)

    # Converts boxes from normalized coordinates to pixel coordinates.
    boxes = box_ops.denormalize_boxes(boxes, image_shape)

    # Resizes and crops image.
    image, image_info = preprocess_ops.resize_and_crop_image(
        image,
        self._output_size,
        padded_size=preprocess_ops.compute_padded_size(self._output_size,
                                                       2**self._max_level),
        aug_scale_min=self._aug_scale_min,
        aug_scale_max=self._aug_scale_max)
    image_height, image_width, _ = image.get_shape().as_list()

    # Resizes and crops boxes.
    image_scale = image_info[2, :]
    offset = image_info[3, :]
    boxes = preprocess_ops.resize_and_crop_boxes(boxes, image_scale,
                                                 image_info[1, :], offset)
    # Filters out ground truth boxes that are all zeros.
    indices = box_ops.get_non_empty_box_indices(boxes)
    boxes = tf.gather(boxes, indices)
    classes = tf.gather(classes, indices)
Xianzhi Du's avatar
Xianzhi Du committed
204
205
    for k, v in attributes.items():
      attributes[k] = tf.gather(v, indices)
Abdullah Rashwan's avatar
Abdullah Rashwan committed
206
207
208
209
210
211
212
213
214
215
216

    # Assigns anchors.
    input_anchor = anchor.build_anchor_generator(
        min_level=self._min_level,
        max_level=self._max_level,
        num_scales=self._num_scales,
        aspect_ratios=self._aspect_ratios,
        anchor_size=self._anchor_size)
    anchor_boxes = input_anchor(image_size=(image_height, image_width))
    anchor_labeler = anchor.AnchorLabeler(self._match_threshold,
                                          self._unmatched_threshold)
Xianzhi Du's avatar
Xianzhi Du committed
217
    (cls_targets, box_targets, att_targets, cls_weights,
Abdullah Rashwan's avatar
Abdullah Rashwan committed
218
     box_weights) = anchor_labeler.label_anchors(
Xianzhi Du's avatar
Xianzhi Du committed
219
         anchor_boxes, boxes, tf.expand_dims(classes, axis=1), attributes)
Abdullah Rashwan's avatar
Abdullah Rashwan committed
220

A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
221
222
    # Casts input image to desired data type.
    image = tf.cast(image, dtype=self._dtype)
Abdullah Rashwan's avatar
Abdullah Rashwan committed
223
224
225
226
227
228
229
230
231
232

    # Packs labels for model_fn outputs.
    labels = {
        'cls_targets': cls_targets,
        'box_targets': box_targets,
        'anchor_boxes': anchor_boxes,
        'cls_weights': cls_weights,
        'box_weights': box_weights,
        'image_info': image_info,
    }
Xianzhi Du's avatar
Xianzhi Du committed
233
234
    if att_targets:
      labels['attribute_targets'] = att_targets
Abdullah Rashwan's avatar
Abdullah Rashwan committed
235
236
237
238
239
240
241
    return image, labels

  def _parse_eval_data(self, data):
    """Parses data for training and evaluation."""
    groundtruths = {}
    classes = data['groundtruth_classes']
    boxes = data['groundtruth_boxes']
Xianzhi Du's avatar
Xianzhi Du committed
242
243
244
245
    # If not empty, `attributes` is a dict of (name, ground_truth) pairs.
    # `ground_gruth` of attributes is assumed in shape [N, attribute_size].
    # TODO(xianzhi): support parsing attributes weights.
    attributes = data.get('groundtruth_attributes', {})
Abdullah Rashwan's avatar
Abdullah Rashwan committed
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275

    # Gets original image and its size.
    image = data['image']
    image_shape = tf.shape(input=image)[0:2]

    # Normalizes image with mean and std pixel values.
    image = preprocess_ops.normalize_image(image)

    # Converts boxes from normalized coordinates to pixel coordinates.
    boxes = box_ops.denormalize_boxes(boxes, image_shape)

    # Resizes and crops image.
    image, image_info = preprocess_ops.resize_and_crop_image(
        image,
        self._output_size,
        padded_size=preprocess_ops.compute_padded_size(self._output_size,
                                                       2**self._max_level),
        aug_scale_min=1.0,
        aug_scale_max=1.0)
    image_height, image_width, _ = image.get_shape().as_list()

    # Resizes and crops boxes.
    image_scale = image_info[2, :]
    offset = image_info[3, :]
    boxes = preprocess_ops.resize_and_crop_boxes(boxes, image_scale,
                                                 image_info[1, :], offset)
    # Filters out ground truth boxes that are all zeros.
    indices = box_ops.get_non_empty_box_indices(boxes)
    boxes = tf.gather(boxes, indices)
    classes = tf.gather(classes, indices)
Xianzhi Du's avatar
Xianzhi Du committed
276
277
    for k, v in attributes.items():
      attributes[k] = tf.gather(v, indices)
Abdullah Rashwan's avatar
Abdullah Rashwan committed
278
279
280
281
282
283
284
285
286
287
288

    # Assigns anchors.
    input_anchor = anchor.build_anchor_generator(
        min_level=self._min_level,
        max_level=self._max_level,
        num_scales=self._num_scales,
        aspect_ratios=self._aspect_ratios,
        anchor_size=self._anchor_size)
    anchor_boxes = input_anchor(image_size=(image_height, image_width))
    anchor_labeler = anchor.AnchorLabeler(self._match_threshold,
                                          self._unmatched_threshold)
Xianzhi Du's avatar
Xianzhi Du committed
289
    (cls_targets, box_targets, att_targets, cls_weights,
Abdullah Rashwan's avatar
Abdullah Rashwan committed
290
     box_weights) = anchor_labeler.label_anchors(
Xianzhi Du's avatar
Xianzhi Du committed
291
         anchor_boxes, boxes, tf.expand_dims(classes, axis=1), attributes)
Abdullah Rashwan's avatar
Abdullah Rashwan committed
292

A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
293
294
    # Casts input image to desired data type.
    image = tf.cast(image, dtype=self._dtype)
Abdullah Rashwan's avatar
Abdullah Rashwan committed
295
296
297
298
299
300
301
302
303
304
305
306
307
308

    # Sets up groundtruth data for evaluation.
    groundtruths = {
        'source_id': data['source_id'],
        'height': data['height'],
        'width': data['width'],
        'num_detections': tf.shape(data['groundtruth_classes']),
        'image_info': image_info,
        'boxes': box_ops.denormalize_boxes(
            data['groundtruth_boxes'], image_shape),
        'classes': data['groundtruth_classes'],
        'areas': data['groundtruth_area'],
        'is_crowds': tf.cast(data['groundtruth_is_crowd'], tf.int32),
    }
Xianzhi Du's avatar
Xianzhi Du committed
309
310
    if 'groundtruth_attributes' in data:
      groundtruths['attributes'] = data['groundtruth_attributes']
Abdullah Rashwan's avatar
Abdullah Rashwan committed
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
    groundtruths['source_id'] = utils.process_source_id(
        groundtruths['source_id'])
    groundtruths = utils.pad_groundtruths_to_fixed_size(
        groundtruths, self._max_num_instances)

    # Packs labels for model_fn outputs.
    labels = {
        'cls_targets': cls_targets,
        'box_targets': box_targets,
        'anchor_boxes': anchor_boxes,
        'cls_weights': cls_weights,
        'box_weights': box_weights,
        'image_info': image_info,
        'groundtruths': groundtruths,
    }
Xianzhi Du's avatar
Xianzhi Du committed
326
327
    if att_targets:
      labels['attribute_targets'] = att_targets
Abdullah Rashwan's avatar
Abdullah Rashwan committed
328
    return image, labels