classification_input.py 6.98 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."""
Fan Yang's avatar
Fan Yang committed
16
from typing import 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
29
30
31
32
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)


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

Fan Yang's avatar
Fan Yang committed
33
34
  def __init__(self,
               image_field_key: str = 'image/encoded',
Abdullah Rashwan's avatar
Abdullah Rashwan committed
35
36
               label_field_key: str = 'image/class/label',
               is_multilabel: bool = False):
Abdullah Rashwan's avatar
Abdullah Rashwan committed
37
    self._keys_to_features = {
Fan Yang's avatar
Fan Yang committed
38
        image_field_key: tf.io.FixedLenFeature((), tf.string, default_value=''),
Abdullah Rashwan's avatar
Abdullah Rashwan committed
39
    }
Abdullah Rashwan's avatar
Abdullah Rashwan committed
40
41
42
43
44
45
46
    if is_multilabel:
      self._keys_to_features.update(
          {label_field_key: tf.io.VarLenFeature(dtype=tf.int64)})
    else:
      self._keys_to_features.update({
          label_field_key: tf.io.FixedLenFeature((), tf.int64, default_value=-1)
      })
Abdullah Rashwan's avatar
Abdullah Rashwan committed
47

Fan Yang's avatar
Fan Yang committed
48
49
  def decode(self,
             serialized_example: tf.train.Example) -> Dict[str, tf.Tensor]:
Abdullah Rashwan's avatar
Abdullah Rashwan committed
50
51
52
53
54
55
56
57
    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
58
59
               output_size: List[int],
               num_classes: float,
Fan Yang's avatar
Fan Yang committed
60
61
               image_field_key: str = 'image/encoded',
               label_field_key: str = 'image/class/label',
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
62
               aug_rand_hflip: bool = True,
63
               aug_type: Optional[common.Augmentation] = None,
Abdullah Rashwan's avatar
Abdullah Rashwan committed
64
               is_multilabel: bool = False,
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
65
               dtype: str = 'float32'):
Abdullah Rashwan's avatar
Abdullah Rashwan committed
66
67
68
    """Initializes parameters for parsing annotations in the dataset.

    Args:
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
69
      output_size: `Tensor` or `list` for [height, width] of output image. The
Abdullah Rashwan's avatar
Abdullah Rashwan committed
70
71
        output_size should be divided by the largest feature stride 2^max_level.
      num_classes: `float`, number of classes.
Fan Yang's avatar
Fan Yang committed
72
73
      image_field_key: A `str` of the key name to encoded image in TFExample.
      label_field_key: A `str` of the key name to label in TFExample.
Abdullah Rashwan's avatar
Abdullah Rashwan committed
74
75
      aug_rand_hflip: `bool`, if True, augment training with random
        horizontal flip.
76
77
      aug_type: An optional Augmentation object to choose from AutoAugment and
        RandAugment.
Abdullah Rashwan's avatar
Abdullah Rashwan committed
78
      is_multilabel: A `bool`, whether or not each example has multiple labels.
Abdullah Rashwan's avatar
Abdullah Rashwan committed
79
80
81
82
83
84
      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
85
86
    self._image_field_key = image_field_key
    self._label_field_key = label_field_key
Abdullah Rashwan's avatar
Abdullah Rashwan committed
87
    self._is_multilabel = is_multilabel
Fan Yang's avatar
Fan Yang committed
88

Abdullah Rashwan's avatar
Abdullah Rashwan committed
89
90
91
92
93
94
95
96
    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))
97
98
99
100
101
102
103
    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
104
        self._augmenter = augment.RandAugment(
105
106
107
108
            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)
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
109
      else:
110
111
        raise ValueError('Augmentation policy {} not supported.'.format(
            aug_type.type))
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
112
113
    else:
      self._augmenter = None
Abdullah Rashwan's avatar
Abdullah Rashwan committed
114
115
116

  def _parse_train_data(self, decoded_tensors):
    """Parses data for training."""
Fan Yang's avatar
Fan Yang committed
117
118
    label = tf.cast(decoded_tensors[self._label_field_key], dtype=tf.int32)
    image_bytes = decoded_tensors[self._image_field_key]
Abdullah Rashwan's avatar
Abdullah Rashwan committed
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
    image_shape = tf.image.extract_jpeg_shape(image_bytes)

    # Crops image.
    # TODO(pengchong): support image format other than JPEG.
    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)

    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)

A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
137
138
139
140
    # Apply autoaug or randaug.
    if self._augmenter is not None:
      image = self._augmenter.distort(image)

Abdullah Rashwan's avatar
Abdullah Rashwan committed
141
142
143
144
145
146
147
148
    # 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)

Abdullah Rashwan's avatar
Abdullah Rashwan committed
149
150
151
152
153
    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)

Abdullah Rashwan's avatar
Abdullah Rashwan committed
154
155
156
157
    return image, label

  def _parse_eval_data(self, decoded_tensors):
    """Parses data for evaluation."""
Fan Yang's avatar
Fan Yang committed
158
159
    label = tf.cast(decoded_tensors[self._label_field_key], dtype=tf.int32)
    image_bytes = decoded_tensors[self._image_field_key]
Abdullah Rashwan's avatar
Abdullah Rashwan committed
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
    image_shape = tf.image.extract_jpeg_shape(image_bytes)

    # Center crops and resizes image.
    image = preprocess_ops.center_crop_image_v2(image_bytes, image_shape)

    image = tf.image.resize(
        image, self._output_size, method=tf.image.ResizeMethod.BILINEAR)

    image = tf.reshape(image, [self._output_size[0], self._output_size[1], 3])

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

Abdullah Rashwan's avatar
Abdullah Rashwan committed
178
179
180
181
182
    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)

Abdullah Rashwan's avatar
Abdullah Rashwan committed
183
    return image, label