video_input.py 11.8 KB
Newer Older
Abdullah Rashwan's avatar
Abdullah Rashwan committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
# Lint as: python3
# Copyright 2020 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.
# ==============================================================================
"""Parser for video and label datasets."""

Yeqing Li's avatar
Yeqing Li committed
18
from typing import Dict, Optional, Tuple, Union
Abdullah Rashwan's avatar
Abdullah Rashwan committed
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36

from absl import logging
import tensorflow as tf

from official.vision.beta.configs import video_classification as exp_cfg
from official.vision.beta.dataloaders import decoder
from official.vision.beta.dataloaders import parser
from official.vision.beta.ops import preprocess_ops_3d

IMAGE_KEY = 'image/encoded'
LABEL_KEY = 'clip/label/index'


def _process_image(image: tf.Tensor,
                   is_training: bool = True,
                   num_frames: int = 32,
                   stride: int = 1,
                   num_test_clips: int = 1,
Yin Cui's avatar
Yin Cui committed
37
38
39
                   min_resize: int = 256,
                   crop_size: int = 224,
                   num_crops: int = 1,
Abdullah Rashwan's avatar
Abdullah Rashwan committed
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
                   zero_centering_image: bool = False,
                   seed: Optional[int] = None) -> tf.Tensor:
  """Processes a serialized image tensor.

  Args:
    image: Input Tensor of shape [timesteps] and type tf.string of serialized
      frames.
    is_training: Whether or not in training mode. If True, random sample, crop
      and left right flip is used.
    num_frames: Number of frames per subclip.
    stride: Temporal stride to sample frames.
    num_test_clips: Number of test clips (1 by default). If more than 1, this
      will sample multiple linearly spaced clips within each video at test time.
      If 1, then a single clip in the middle of the video is sampled. The clips
      are aggreagated in the batch dimension.
    min_resize: Frames are resized so that min(height, width) is min_resize.
    crop_size: Final size of the frame after cropping the resized frames. Both
      height and width are the same.
Yin Cui's avatar
Yin Cui committed
58
    num_crops: Number of crops to perform on the resized frames.
Abdullah Rashwan's avatar
Abdullah Rashwan committed
59
60
61
62
63
64
65
66
67
68
69
70
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
    zero_centering_image: If True, frames are normalized to values in [-1, 1].
      If False, values in [0, 1].
    seed: A deterministic seed to use when sampling.

  Returns:
    Processed frames. Tensor of shape
      [num_frames * num_test_clips, crop_size, crop_size, 3].
  """
  # Validate parameters.
  if is_training and num_test_clips != 1:
    logging.warning(
        '`num_test_clips` %d is ignored since `is_training` is `True`.',
        num_test_clips)

  # Temporal sampler.
  if is_training:
    # Sample random clip.
    image = preprocess_ops_3d.sample_sequence(image, num_frames, True, stride,
                                              seed)
  elif num_test_clips > 1:
    # Sample linspace clips.
    image = preprocess_ops_3d.sample_linspace_sequence(image, num_test_clips,
                                                       num_frames, stride)
  else:
    # Sample middle clip.
    image = preprocess_ops_3d.sample_sequence(image, num_frames, False, stride)

  # Decode JPEG string to tf.uint8.
  image = preprocess_ops_3d.decode_jpeg(image, 3)

  # Resize images (resize happens only if necessary to save compute).
  image = preprocess_ops_3d.resize_smallest(image, min_resize)

  if is_training:
    # Standard image data augmentation: random crop and random flip.
    image = preprocess_ops_3d.crop_image(image, crop_size, crop_size, True,
                                         seed)
    image = preprocess_ops_3d.random_flip_left_right(image, seed)
  else:
Yin Cui's avatar
Yin Cui committed
98
99
100
    # Crop of the frames.
    image = preprocess_ops_3d.crop_image(image, crop_size, crop_size, False,
                                         num_crops)
Abdullah Rashwan's avatar
Abdullah Rashwan committed
101
102
103
104
105
106
107
108

  # Cast the frames in float32, normalizing according to zero_centering_image.
  return preprocess_ops_3d.normalize_image(image, zero_centering_image)


def _postprocess_image(image: tf.Tensor,
                       is_training: bool = True,
                       num_frames: int = 32,
Yin Cui's avatar
Yin Cui committed
109
110
                       num_test_clips: int = 1,
                       num_test_crops: int = 1) -> tf.Tensor:
Abdullah Rashwan's avatar
Abdullah Rashwan committed
111
112
113
114
115
116
117
118
119
120
121
122
123
  """Processes a batched Tensor of frames.

  The same parameters used in process should be used here.

  Args:
    image: Input Tensor of shape [batch, timesteps, height, width, 3].
    is_training: Whether or not in training mode. If True, random sample, crop
      and left right flip is used.
    num_frames: Number of frames per subclip.
    num_test_clips: Number of test clips (1 by default). If more than 1, this
      will sample multiple linearly spaced clips within each video at test time.
      If 1, then a single clip in the middle of the video is sampled. The clips
      are aggreagated in the batch dimension.
Yin Cui's avatar
Yin Cui committed
124
125
126
    num_test_crops: Number of test crops (1 by default). If more than 1, there
      are multiple crops for each clip at test time. If 1, there is a single
      central crop. The crops are aggreagated in the batch dimension.
Abdullah Rashwan's avatar
Abdullah Rashwan committed
127
128
129

  Returns:
    Processed frames. Tensor of shape
Yin Cui's avatar
Yin Cui committed
130
      [batch * num_test_clips * num_test_crops, num_frames, height, width, 3].
Abdullah Rashwan's avatar
Abdullah Rashwan committed
131
  """
Yin Cui's avatar
Yin Cui committed
132
133
134
135
136
  num_views = num_test_clips * num_test_crops
  if num_views > 1 and not is_training:
    # In this case, multiple views are merged together in batch dimenstion which
    # will be batch * num_views.
    image = tf.reshape(image, [-1, num_frames] + image.shape[2:].as_list())
Abdullah Rashwan's avatar
Abdullah Rashwan committed
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155

  return image


def _process_label(label: tf.Tensor,
                   one_hot_label: bool = True,
                   num_classes: Optional[int] = None) -> tf.Tensor:
  """Processes label Tensor."""
  # Validate parameters.
  if one_hot_label and not num_classes:
    raise ValueError(
        '`num_classes` should be given when requesting one hot label.')

  # Cast to tf.int32.
  label = tf.cast(label, dtype=tf.int32)

  if one_hot_label:
    # Replace label index by one hot representation.
    label = tf.one_hot(label, num_classes)
Yeqing Li's avatar
Yeqing Li committed
156
157
158
159
160
    if len(label.shape.as_list()) > 1:
      label = tf.reduce_sum(label, axis=0)
    if num_classes == 1:
      # The trick for single label.
      label = 1 - label
Abdullah Rashwan's avatar
Abdullah Rashwan committed
161
162
163
164
165
166
167
168

  return label


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

  def __init__(self, image_key: str = IMAGE_KEY, label_key: str = LABEL_KEY):
Yeqing Li's avatar
Yeqing Li committed
169
170
    self._image_key = image_key
    self._label_key = label_key
Abdullah Rashwan's avatar
Abdullah Rashwan committed
171
172
    self._context_description = {
        # One integer stored in context.
Yeqing Li's avatar
Yeqing Li committed
173
        self._label_key: tf.io.VarLenFeature(tf.int64),
Abdullah Rashwan's avatar
Abdullah Rashwan committed
174
175
176
177
178
179
    }
    self._sequence_description = {
        # Each image is a string encoding JPEG.
        self._image_key: tf.io.FixedLenSequenceFeature((), tf.string),
    }

Yeqing Li's avatar
Yeqing Li committed
180
181
182
183
184
185
186
187
188
189
190
191
  def add_feature(self, feature_name: str,
                  feature_type: Union[tf.io.VarLenFeature,
                                      tf.io.FixedLenFeature,
                                      tf.io.FixedLenSequenceFeature]):
    self._sequence_description[feature_name] = feature_type

  def add_context(self, feature_name: str,
                  feature_type: Union[tf.io.VarLenFeature,
                                      tf.io.FixedLenFeature,
                                      tf.io.FixedLenSequenceFeature]):
    self._context_description[feature_name] = feature_type

Abdullah Rashwan's avatar
Abdullah Rashwan committed
192
193
  def decode(self, serialized_example):
    """Parses a single tf.Example into image and label tensors."""
Yeqing Li's avatar
Yeqing Li committed
194
    result = {}
Abdullah Rashwan's avatar
Abdullah Rashwan committed
195
196
197
    context, sequences = tf.io.parse_single_sequence_example(
        serialized_example, self._context_description,
        self._sequence_description)
Yeqing Li's avatar
Yeqing Li committed
198
199
200
201
202
203
    result.update(context)
    result.update(sequences)
    for key, value in result.items():
      if isinstance(value, tf.SparseTensor):
        result[key] = tf.sparse.to_dense(value)
    return result
Abdullah Rashwan's avatar
Abdullah Rashwan committed
204
205
206
207
208
209
210
211
212
213
214
215
216
217


class Parser(parser.Parser):
  """Parses a video and label dataset."""

  def __init__(self,
               input_params: exp_cfg.DataConfig,
               image_key: str = IMAGE_KEY,
               label_key: str = LABEL_KEY):
    self._num_frames = input_params.feature_shape[0]
    self._stride = input_params.temporal_stride
    self._num_test_clips = input_params.num_test_clips
    self._min_resize = input_params.min_image_size
    self._crop_size = input_params.feature_shape[1]
Yin Cui's avatar
Yin Cui committed
218
    self._num_crops = input_params.num_test_crops
Abdullah Rashwan's avatar
Abdullah Rashwan committed
219
220
221
222
    self._one_hot_label = input_params.one_hot
    self._num_classes = input_params.num_classes
    self._image_key = image_key
    self._label_key = label_key
223
    self._dtype = tf.dtypes.as_dtype(input_params.dtype)
Yeqing Li's avatar
Yeqing Li committed
224
225
226
227
    self._output_audio = input_params.output_audio
    if self._output_audio:
      self._audio_feature = input_params.audio_feature
      self._audio_shape = input_params.audio_feature_shape
Abdullah Rashwan's avatar
Abdullah Rashwan committed
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242

  def _parse_train_data(
      self, decoded_tensors: Dict[str, tf.Tensor]
  ) -> Tuple[Dict[str, tf.Tensor], tf.Tensor]:
    """Parses data for training."""
    # Process image and label.
    image = decoded_tensors[self._image_key]
    image = _process_image(
        image=image,
        is_training=True,
        num_frames=self._num_frames,
        stride=self._stride,
        num_test_clips=self._num_test_clips,
        min_resize=self._min_resize,
        crop_size=self._crop_size)
243
    image = tf.cast(image, dtype=self._dtype)
Yeqing Li's avatar
Yeqing Li committed
244
    features = {'image': image}
Yeqing Li's avatar
Yeqing Li committed
245
246

    label = decoded_tensors[self._label_key]
Abdullah Rashwan's avatar
Abdullah Rashwan committed
247
248
    label = _process_label(label, self._one_hot_label, self._num_classes)

Yeqing Li's avatar
Yeqing Li committed
249
250
251
252
253
254
255
256
257
258
    if self._output_audio:
      audio = decoded_tensors[self._audio_feature]
      audio = tf.cast(audio, dtype=self._dtype)
      # TODO(yeqing): synchronize audio/video sampling. Especially randomness.
      audio = preprocess_ops_3d.sample_sequence(
          audio, self._audio_shape[0], random=False, stride=1)
      audio = tf.ensure_shape(audio, self._audio_shape)
      features['audio'] = audio

    return features, label
Abdullah Rashwan's avatar
Abdullah Rashwan committed
259
260
261
262
263
264
265
266
267
268
269
270
271

  def _parse_eval_data(
      self, decoded_tensors: Dict[str, tf.Tensor]
  ) -> Tuple[Dict[str, tf.Tensor], tf.Tensor]:
    """Parses data for evaluation."""
    image = decoded_tensors[self._image_key]
    image = _process_image(
        image=image,
        is_training=False,
        num_frames=self._num_frames,
        stride=self._stride,
        num_test_clips=self._num_test_clips,
        min_resize=self._min_resize,
Yin Cui's avatar
Yin Cui committed
272
273
        crop_size=self._crop_size,
        num_crops=self._num_crops)
274
    image = tf.cast(image, dtype=self._dtype)
Yeqing Li's avatar
Yeqing Li committed
275
    features = {'image': image}
Yeqing Li's avatar
Yeqing Li committed
276
277

    label = decoded_tensors[self._label_key]
Abdullah Rashwan's avatar
Abdullah Rashwan committed
278
279
    label = _process_label(label, self._one_hot_label, self._num_classes)

Yeqing Li's avatar
Yeqing Li committed
280
281
282
283
284
285
286
287
288
    if self._output_audio:
      audio = decoded_tensors[self._audio_feature]
      audio = tf.cast(audio, dtype=self._dtype)
      audio = preprocess_ops_3d.sample_sequence(
          audio, 20, random=False, stride=1)
      audio = tf.ensure_shape(audio, [20, 2048])
      features['audio'] = audio

    return features, label
Abdullah Rashwan's avatar
Abdullah Rashwan committed
289
290
291
292
293
294
295
296
297
298


class PostBatchProcessor(object):
  """Processes a video and label dataset which is batched."""

  def __init__(self, input_params: exp_cfg.DataConfig):
    self._is_training = input_params.is_training

    self._num_frames = input_params.feature_shape[0]
    self._num_test_clips = input_params.num_test_clips
Yin Cui's avatar
Yin Cui committed
299
    self._num_test_crops = input_params.num_test_crops
Abdullah Rashwan's avatar
Abdullah Rashwan committed
300

Yeqing Li's avatar
Yeqing Li committed
301
302
  def __call__(self, features: Dict[str, tf.Tensor],
               label: tf.Tensor) -> Tuple[Dict[str, tf.Tensor], tf.Tensor]:
Abdullah Rashwan's avatar
Abdullah Rashwan committed
303
    """Parses a single tf.Example into image and label tensors."""
Yeqing Li's avatar
Yeqing Li committed
304
305
306
307
308
309
    for key in ['image', 'audio']:
      if key in features:
        features[key] = _postprocess_image(
            image=features[key],
            is_training=self._is_training,
            num_frames=self._num_frames,
Yin Cui's avatar
Yin Cui committed
310
311
            num_test_clips=self._num_test_clips,
            num_test_crops=self._num_test_crops)
Yeqing Li's avatar
Yeqing Li committed
312
313

    return features, label