karolmajek/Axial-DeepLab-SWideRNet

1# coding=utf-82# Copyright 2021 The Deeplab2 Authors.3#4# Licensed under the Apache License, Version 2.0 (the "License");5# you may not use this file except in compliance with the License.6# You may obtain a copy of the License at7#8#     http://www.apache.org/licenses/LICENSE-2.09#10# Unless required by applicable law or agreed to in writing, software11# distributed under the License is distributed on an "AS IS" BASIS,12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.13# See the License for the specific language governing permissions and14# limitations under the License.15 16"""Input reader to load segmentation dataset."""17 18import tensorflow as tf19 20_NUM_INPUTS_PROCESSED_CONCURRENTLY = 3221_SHUFFLE_BUFFER_SIZE = 100022 23 24class InputReader(object):25  """Input function that creates a dataset from files."""26 27  def __init__(self,28               file_pattern,29               decoder_fn,30               generator_fn=None,31               is_training=False):32    """Initializes the input reader.33 34    Args:35      file_pattern: The file pattern for the data example, in TFRecord format36      decoder_fn: A callable that takes a serialized tf.Example and produces37        parsed (and potentially processed / augmented) tensors.38      generator_fn: An optional `callable` that takes the decoded raw tensors39        dict and generates a ground-truth dictionary that can be consumed by40        the model. It will be executed after decoder_fn (default: None).41      is_training: If this dataset is used for training or not (default: False).42    """43    self._file_pattern = file_pattern44    self._is_training = is_training45    self._decoder_fn = decoder_fn46    self._generator_fn = generator_fn47 48  def __call__(self, batch_size=1, max_num_examples=-1):49    """Provides tf.data.Dataset object.50 51    Args:52      batch_size: Expected batch size input data.53      max_num_examples: Positive integer or -1. If positive, the returned54        dataset will only take (at most) this number of examples and raise55        tf.errors.OutOfRangeError after that (default: -1).56 57    Returns:58      tf.data.Dataset object.59    """60    dataset = tf.data.Dataset.list_files(self._file_pattern)61 62    if self._is_training:63      # File level shuffle.64      dataset = dataset.shuffle(dataset.cardinality(),65                                reshuffle_each_iteration=True)66      dataset = dataset.repeat()67 68    # During training, interleave TFRecord conversion for maximum efficiency.69    # During evaluation, read input in consecutive order for tasks requiring70    # such behavior.71    dataset = dataset.interleave(72        map_func=tf.data.TFRecordDataset,73        cycle_length=(_NUM_INPUTS_PROCESSED_CONCURRENTLY74                      if self._is_training else 1),75        num_parallel_calls=tf.data.experimental.AUTOTUNE,76        deterministic=not self._is_training)77 78    if self._is_training:79      dataset = dataset.shuffle(_SHUFFLE_BUFFER_SIZE)80    if max_num_examples > 0:81      dataset = dataset.take(max_num_examples)82 83    # Parses the fetched records to input tensors for model function.84    dataset = dataset.map(85        self._decoder_fn, num_parallel_calls=tf.data.experimental.AUTOTUNE)86    if self._generator_fn is not None:87      dataset = dataset.map(88          self._generator_fn, num_parallel_calls=tf.data.experimental.AUTOTUNE)89    dataset = dataset.batch(batch_size, drop_remainder=True)90    dataset = dataset.prefetch(tf.data.experimental.AUTOTUNE)91    return dataset92