classification_input.py 8.46 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
"""Classification decoder and parser."""
Rajagopal Ananthanarayanan's avatar
Rajagopal Ananthanarayanan committed
16
from typing import Any, Dict, List, Optional
Abdullah Rashwan's avatar
Abdullah Rashwan committed
17
18
19
# Import libraries
import tensorflow as tf

20
from official.vision.beta.configs import common
Abdullah Rashwan's avatar
Abdullah Rashwan committed
21
22
from official.vision.beta.dataloaders import decoder
from official.vision.beta.dataloaders import parser
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
23
from official.vision.beta.ops import augment
Abdullah Rashwan's avatar
Abdullah Rashwan committed
24
25
26
27
28
from official.vision.beta.ops import preprocess_ops

MEAN_RGB = (0.485 * 255, 0.456 * 255, 0.406 * 255)
STDDEV_RGB = (0.229 * 255, 0.224 * 255, 0.225 * 255)

Rajagopal Ananthanarayanan's avatar
Rajagopal Ananthanarayanan committed
29
30
31
DEFAULT_IMAGE_FIELD_KEY = 'image/encoded'
DEFAULT_LABEL_FIELD_KEY = 'image/class/label'

Abdullah Rashwan's avatar
Abdullah Rashwan committed
32
33
34
35

class Decoder(decoder.Decoder):
  """A tf.Example decoder for classification task."""

Fan Yang's avatar
Fan Yang committed
36
  def __init__(self,
Rajagopal Ananthanarayanan's avatar
Rajagopal Ananthanarayanan committed
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
               image_field_key: str = DEFAULT_IMAGE_FIELD_KEY,
               label_field_key: str = DEFAULT_LABEL_FIELD_KEY,
               is_multilabel: bool = False,
               keys_to_features: Optional[Dict[str, Any]] = None):
    if not keys_to_features:
      keys_to_features = {
          image_field_key:
              tf.io.FixedLenFeature((), tf.string, default_value=''),
      }
      if is_multilabel:
        keys_to_features.update(
            {label_field_key: tf.io.VarLenFeature(dtype=tf.int64)})
      else:
        keys_to_features.update({
            label_field_key:
                tf.io.FixedLenFeature((), tf.int64, default_value=-1)
        })
    self._keys_to_features = keys_to_features
Abdullah Rashwan's avatar
Abdullah Rashwan committed
55

Rajagopal Ananthanarayanan's avatar
Rajagopal Ananthanarayanan committed
56
  def decode(self, serialized_example):
Abdullah Rashwan's avatar
Abdullah Rashwan committed
57
58
59
60
61
62
63
64
    return tf.io.parse_single_example(
        serialized_example, self._keys_to_features)


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

  def __init__(self,
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
65
66
               output_size: List[int],
               num_classes: float,
Rajagopal Ananthanarayanan's avatar
Rajagopal Ananthanarayanan committed
67
68
               image_field_key: str = DEFAULT_IMAGE_FIELD_KEY,
               label_field_key: str = DEFAULT_LABEL_FIELD_KEY,
Abdullah Rashwan's avatar
Abdullah Rashwan committed
69
               decode_jpeg_only: bool = True,
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
70
               aug_rand_hflip: bool = True,
71
               aug_type: Optional[common.Augmentation] = None,
Abdullah Rashwan's avatar
Abdullah Rashwan committed
72
               is_multilabel: bool = False,
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
73
               dtype: str = 'float32'):
Abdullah Rashwan's avatar
Abdullah Rashwan committed
74
75
76
    """Initializes parameters for parsing annotations in the dataset.

    Args:
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
77
      output_size: `Tensor` or `list` for [height, width] of output image. The
Abdullah Rashwan's avatar
Abdullah Rashwan committed
78
79
        output_size should be divided by the largest feature stride 2^max_level.
      num_classes: `float`, number of classes.
Rajagopal Ananthanarayanan's avatar
Rajagopal Ananthanarayanan committed
80
81
      image_field_key: `str`, the key name to encoded image in tf.Example.
      label_field_key: `str`, the key name to label in tf.Example.
Abdullah Rashwan's avatar
Abdullah Rashwan committed
82
83
      decode_jpeg_only: `bool`, if True, only JPEG format is decoded, this is
        faster than decoding other types. Default is True.
Abdullah Rashwan's avatar
Abdullah Rashwan committed
84
85
      aug_rand_hflip: `bool`, if True, augment training with random
        horizontal flip.
86
87
      aug_type: An optional Augmentation object to choose from AutoAugment and
        RandAugment.
Abdullah Rashwan's avatar
Abdullah Rashwan committed
88
      is_multilabel: A `bool`, whether or not each example has multiple labels.
Abdullah Rashwan's avatar
Abdullah Rashwan committed
89
90
91
92
93
94
      dtype: `str`, cast output image in dtype. It can be 'float32', 'float16',
        or 'bfloat16'.
    """
    self._output_size = output_size
    self._aug_rand_hflip = aug_rand_hflip
    self._num_classes = num_classes
Fan Yang's avatar
Fan Yang committed
95
    self._image_field_key = image_field_key
Abdullah Rashwan's avatar
Abdullah Rashwan committed
96
97
98
99
100
101
102
103
    if dtype == 'float32':
      self._dtype = tf.float32
    elif dtype == 'float16':
      self._dtype = tf.float16
    elif dtype == 'bfloat16':
      self._dtype = tf.bfloat16
    else:
      raise ValueError('dtype {!r} is not supported!'.format(dtype))
104
105
106
107
108
109
110
    if aug_type:
      if aug_type.type == 'autoaug':
        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)
      elif aug_type.type == 'randaug':
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
111
        self._augmenter = augment.RandAugment(
112
113
114
            num_layers=aug_type.randaug.num_layers,
            magnitude=aug_type.randaug.magnitude,
            cutout_const=aug_type.randaug.cutout_const,
Fan Yang's avatar
Fan Yang committed
115
116
            translate_const=aug_type.randaug.translate_const,
            prob_to_apply=aug_type.randaug.prob_to_apply)
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
117
      else:
118
119
        raise ValueError('Augmentation policy {} not supported.'.format(
            aug_type.type))
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
120
121
    else:
      self._augmenter = None
Rajagopal Ananthanarayanan's avatar
Rajagopal Ananthanarayanan committed
122
123
    self._label_field_key = label_field_key
    self._is_multilabel = is_multilabel
Abdullah Rashwan's avatar
Abdullah Rashwan committed
124
    self._decode_jpeg_only = decode_jpeg_only
Abdullah Rashwan's avatar
Abdullah Rashwan committed
125
126
127

  def _parse_train_data(self, decoded_tensors):
    """Parses data for training."""
Rajagopal Ananthanarayanan's avatar
Rajagopal Ananthanarayanan committed
128
129
130
131
132
133
134
135
136
137
138
    image = self._parse_train_image(decoded_tensors)
    label = tf.cast(decoded_tensors[self._label_field_key], dtype=tf.int32)
    if self._is_multilabel:
      if isinstance(label, tf.sparse.SparseTensor):
        label = tf.sparse.to_dense(label)
      label = tf.reduce_sum(tf.one_hot(label, self._num_classes), axis=0)
    return image, label

  def _parse_eval_data(self, decoded_tensors):
    """Parses data for evaluation."""
    image = self._parse_eval_image(decoded_tensors)
Fan Yang's avatar
Fan Yang committed
139
    label = tf.cast(decoded_tensors[self._label_field_key], dtype=tf.int32)
Rajagopal Ananthanarayanan's avatar
Rajagopal Ananthanarayanan committed
140
141
142
143
144
145
146
147
    if self._is_multilabel:
      if isinstance(label, tf.sparse.SparseTensor):
        label = tf.sparse.to_dense(label)
      label = tf.reduce_sum(tf.one_hot(label, self._num_classes), axis=0)
    return image, label

  def _parse_train_image(self, decoded_tensors):
    """Parses image data for training."""
Fan Yang's avatar
Fan Yang committed
148
    image_bytes = decoded_tensors[self._image_field_key]
Abdullah Rashwan's avatar
Abdullah Rashwan committed
149

Abdullah Rashwan's avatar
Abdullah Rashwan committed
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
    if self._decode_jpeg_only:
      image_shape = tf.image.extract_jpeg_shape(image_bytes)

      # Crops image.
      cropped_image = preprocess_ops.random_crop_image_v2(
          image_bytes, image_shape)
      image = tf.cond(
          tf.reduce_all(tf.equal(tf.shape(cropped_image), image_shape)),
          lambda: preprocess_ops.center_crop_image_v2(image_bytes, image_shape),
          lambda: cropped_image)
    else:
      # Decodes image.
      image = tf.io.decode_image(image_bytes, channels=3)
      image.set_shape([None, None, 3])

      # Crops image.
      cropped_image = preprocess_ops.random_crop_image(image)

      image = tf.cond(
          tf.reduce_all(tf.equal(tf.shape(cropped_image), tf.shape(image))),
          lambda: preprocess_ops.center_crop_image(image),
          lambda: cropped_image)
Abdullah Rashwan's avatar
Abdullah Rashwan committed
172
173
174
175
176
177
178

    if self._aug_rand_hflip:
      image = tf.image.random_flip_left_right(image)

    # Resizes image.
    image = tf.image.resize(
        image, self._output_size, method=tf.image.ResizeMethod.BILINEAR)
Abdullah Rashwan's avatar
Abdullah Rashwan committed
179
    image.set_shape([self._output_size[0], self._output_size[1], 3])
Abdullah Rashwan's avatar
Abdullah Rashwan committed
180

A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
181
182
183
184
    # Apply autoaug or randaug.
    if self._augmenter is not None:
      image = self._augmenter.distort(image)

Abdullah Rashwan's avatar
Abdullah Rashwan committed
185
186
187
188
189
190
191
192
    # Normalizes image with mean and std pixel values.
    image = preprocess_ops.normalize_image(image,
                                           offset=MEAN_RGB,
                                           scale=STDDEV_RGB)

    # Convert image to self._dtype.
    image = tf.image.convert_image_dtype(image, self._dtype)

Rajagopal Ananthanarayanan's avatar
Rajagopal Ananthanarayanan committed
193
    return image
Abdullah Rashwan's avatar
Abdullah Rashwan committed
194

Rajagopal Ananthanarayanan's avatar
Rajagopal Ananthanarayanan committed
195
196
  def _parse_eval_image(self, decoded_tensors):
    """Parses image data for evaluation."""
Fan Yang's avatar
Fan Yang committed
197
    image_bytes = decoded_tensors[self._image_field_key]
Abdullah Rashwan's avatar
Abdullah Rashwan committed
198

Abdullah Rashwan's avatar
Abdullah Rashwan committed
199
200
201
202
203
204
205
206
207
208
209
210
    if self._decode_jpeg_only:
      image_shape = tf.image.extract_jpeg_shape(image_bytes)

      # Center crops.
      image = preprocess_ops.center_crop_image_v2(image_bytes, image_shape)
    else:
      # Decodes image.
      image = tf.io.decode_image(image_bytes, channels=3)
      image.set_shape([None, None, 3])

      # Center crops.
      image = preprocess_ops.center_crop_image(image)
Abdullah Rashwan's avatar
Abdullah Rashwan committed
211
212
213

    image = tf.image.resize(
        image, self._output_size, method=tf.image.ResizeMethod.BILINEAR)
Abdullah Rashwan's avatar
Abdullah Rashwan committed
214
    image.set_shape([self._output_size[0], self._output_size[1], 3])
Abdullah Rashwan's avatar
Abdullah Rashwan committed
215
216
217
218
219
220
221
222
223

    # Normalizes image with mean and std pixel values.
    image = preprocess_ops.normalize_image(image,
                                           offset=MEAN_RGB,
                                           scale=STDDEV_RGB)

    # Convert image to self._dtype.
    image = tf.image.convert_image_dtype(image, self._dtype)

Rajagopal Ananthanarayanan's avatar
Rajagopal Ananthanarayanan committed
224
    return image