text: # Copyright 2021 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.
  
  # Lint as: python3
  """Dataset utilities for vision tasks using TFDS and tf.data.Dataset."""
  from __future__ import absolute_import
  from __future__ import division
  # from __future__ import google_type_annotations
  from __future__ import print_function
  import dataclasses
  import os
  from typing import Any, List, Mapping, Optional, Tuple, Union
  
  from absl import logging
  import tensorflow as tf
  import tensorflow_datasets as tfds
  from official.legacy.image_classification import augment
  from official.legacy.image_classification import preprocessing
  from official.modeling.hyperparams import base_config
  
  AUGMENTERS = {
      'autoaugment': augment.AutoAugment,
      'randaugment': augment.RandAugment,
  }
  
  
  @dataclasses.dataclass
  class AugmentConfig(base_config.Config):
    """Configuration for image augmenters.
  
    Attributes:
      name: The name of the image augmentation to use. Possible options are None
        (default), 'autoaugment', or 'randaugment'.
      params: Any parameters used to initialize the augmenter.
    """
    name: Optional[str] = None
    params: Optional[Mapping[str, Any]] = None
  
    def build(self) -> augment.ImageAugment:
      """Build the augmenter using this config."""
      params = self.params or {}
      augmenter = AUGMENTERS.get(self.name, None)
      return augmenter(**params) if augmenter is not None else None
  
  
  @dataclasses.dataclass
  class DatasetConfig(base_config.Config):
    """The base configuration for building datasets.
  
    Attributes:
      name: The name of the Dataset. Usually should correspond to a TFDS dataset.
      data_dir: The path where the dataset files are stored, if available.
      filenames: Optional list of strings representing the TFRecord names.
      builder: The builder type used to load the dataset. Value should be one of
        'tfds' (load using TFDS), 'records' (load from TFRecords), or 'synthetic'
        (generate dummy synthetic data without reading from files).
      split: The split of the dataset. Usually 'train', 'validation', or 'test'.
      image_size: The size of the image in the dataset. This assumes that `width`
        == `height`. Set to 'infer' to infer the image size from TFDS info. This
        requires `name` to be a registered dataset in TFDS.
      num_classes: The number of classes given by the dataset. Set to 'infer' to
        infer the image size from TFDS info. This requires `name` to be a
        registered dataset in TFDS.
      num_channels: The number of channels given by the dataset. Set to 'infer' to
        infer the image size from TFDS info. This requires `name` to be a
        registered dataset in TFDS.
      num_examples: The number of examples given by the dataset. Set to 'infer' to
        infer the image size from TFDS info. This requires `name` to be a
        registered dataset in TFDS.
      batch_size: The base batch size for the dataset.
      use_per_replica_batch_size: Whether to scale the batch size based on
        available resources. If set to `True`, the dataset builder will return
        batch_size multiplied by `num_devices`, the number of device replicas
        (e.g., the number of GPUs or TPU cores). This setting should be `True` if
        the strategy argument is passed to `build()` and `num_devices > 1`.
      num_devices: The number of replica devices to use. This should be set by
        `strategy.num_replicas_in_sync` when using a distribution strategy.
      dtype: The desired dtype of the dataset. This will be set during
        preprocessing.
      one_hot: Whether to apply one hot encoding. Set to `True` to be able to use
        label smoothing.
      augmenter: The augmenter config to use. No augmentation is used by default.
      download: Whether to download data using TFDS.
      shuffle_buffer_size: The buffer size used for shuffling training data.
      file_shuffle_buffer_size: The buffer size used for shuffling raw training
        files.
      skip_decoding: Whether to skip image decoding when loading from TFDS.
      cache: whether to cache to dataset examples. Can be used to avoid re-reading
        from disk on the second epoch. Requires significant memory overhead.
      tf_data_service: The URI of a tf.data service to offload preprocessing onto
        during training. The URI should be in the format "protocol://address",
        e.g. "grpc://tf-data-service:5050".
      mean_subtract: whether or not to apply mean subtraction to the dataset.
      standardize: whether or not to apply standardization to the dataset.
    """
    name: Optional[str] = None
    data_dir: Optional[str] = None
    filenames: Optional[List[str]] = None
    builder: str = 'tfds'
    split: str = 'train'
    image_size: Union[int, str] = 'infer'
    num_classes: Union[int, str] = 'infer'
    num_channels: Union[int, str] = 'infer'
    num_examples: Union[int, str] = 'infer'
    batch_size: int = 128
    use_per_replica_batch_size: bool = True
    num_devices: int = 1
    dtype: str = 'float32'
    one_hot: bool = True
    augmenter: AugmentConfig = AugmentConfig()
    download: bool = False
    shuffle_buffer_size: int = 10000
    file_shuffle_buffer_size: int = 1024
    skip_decoding: bool = True
    cache: bool = False
    tf_data_service: Optional[str] = None
    mean_subtract: bool = False
    standardize: bool = False
  
    @property
    def has_data(self):
      """Whether this dataset is has any data associated with it."""
      return self.name or self.data_dir or self.filenames
  
  
  @dataclasses.dataclass
  class ImageNetConfig(DatasetConfig):
    """The base ImageNet dataset config."""
    name: str = 'imagenet2012'
    # Note: for large datasets like ImageNet, using records is faster than tfds
    builder: str = 'records'
    image_size: int = 224
    num_channels: int = 3
    num_examples: int = 1281167
    num_classes: int = 1000
    batch_size: int = 128
  
  
  @dataclasses.dataclass
  class Cifar10Config(DatasetConfig):
    """The base CIFAR-10 dataset config."""
    name: str = 'cifar10'
    image_size: int = 224
    batch_size: int = 128
    download: bool = True
    cache: bool = True
  
  
  class DatasetBuilder:
    """An object for building datasets.
  
    Allows building various pipelines fetching examples, preprocessing, etc.
    Maintains additional state information calculated from the dataset, i.e.,
    training set split, batch size, and number of steps (batches).
    """
  
    def __init__(self, config: DatasetConfig, **overrides: Any):
      """Initialize the builder from the config."""
      self.config = config.replace(**overrides)
      self.builder_info = None
  
      if self.config.augmenter is not None:
        logging.info('Using augmentation: %s', self.config.augmenter.name)
        self.augmenter = self.config.augmenter.build()
      else:
        self.augmenter = None
  
    @property
    def is_training(self) -> bool:
      """Whether this is the training set."""
      return self.config.split == 'train'
  
    @property
    def batch_size(self) -> int:
      """The batch size, multiplied by the number of replicas (if configured)."""
      if self.config.use_per_replica_batch_size:
        return self.config.batch_size * self.config.num_devices
      else:
        return self.config.batch_size
  
    @property
    def global_batch_size(self):
      """The global batch size across all replicas."""
      return self.batch_size
  
    @property
    def local_batch_size(self):
      """The base unscaled batch size."""
      if self.config.use_per_replica_batch_size:
        return self.config.batch_size
      else:
        return self.config.batch_size // self.config.num_devices
  
    @property
    def num_steps(self) -> int:
      """The number of steps (batches) to exhaust this dataset."""
      # Always divide by the global batch size to get the correct # of steps
      return self.num_examples // self.global_batch_size
  
    @property
    def dtype(self) -> tf.dtypes.DType:
      """Converts the config's dtype string to a tf dtype.
  
      Returns:
        A mapping from string representation of a dtype to the `tf.dtypes.DType`.
  
      Raises:
        ValueError if the config's dtype is not supported.
  
      """
      dtype_map = {
          'float32': tf.float32,
          'bfloat16': tf.bfloat16,
          'float16': tf.float16,
          'fp32': tf.float32,
          'bf16': tf.bfloat16,
      }
      try:
        return dtype_map[self.config.dtype]
      except:
        raise ValueError('Invalid DType provided. Supported types: {}'.format(
            dtype_map.keys()))
  
    @property
    def image_size(self) -> int:
      """The size of each image (can be inferred from the dataset)."""
  
      if self.config.image_size == 'infer':
        return self.info.features['image'].shape[0]
      else:
        return int(self.config.image_size)
  
    @property
    def num_channels(self) -> int:
      """The number of image channels (can be inferred from the dataset)."""
      if self.config.num_channels == 'infer':
        return self.info.features['image'].shape[-1]
      else:
        return int(self.config.num_channels)
  
    @property
    def num_examples(self) -> int:
      """The number of examples (can be inferred from the dataset)."""
      if self.config.num_examples == 'infer':
        return self.info.splits[self.config.split].num_examples
      else:
        return int(self.config.num_examples)
  
    @property
    def num_classes(self) -> int:
      """The number of classes (can be inferred from the dataset)."""
      if self.config.num_classes == 'infer':
        return self.info.features['label'].num_classes
      else:
        return int(self.config.num_classes)
  
    @property
    def info(self) -> tfds.core.DatasetInfo:
      """The TFDS dataset info, if available."""
      try:
        if self.builder_info is None:
          self.builder_info = tfds.builder(self.config.name).info
      except ConnectionError as e:
        logging.error('Failed to use TFDS to load info. Please set dataset info '
                      '(image_size, num_channels, num_examples, num_classes) in '
                      'the dataset config.')
        raise e
      return self.builder_info
  
    def build(
        self,
        strategy: Optional[tf.distribute.Strategy] = None) -> tf.data.Dataset:
      """Construct a dataset end-to-end and return it using an optional strategy.
  
      Args:
        strategy: a strategy that, if passed, will distribute the dataset
          according to that strategy. If passed and `num_devices > 1`,
          `use_per_replica_batch_size` must be set to `True`.
  
      Returns:
        A TensorFlow dataset outputting batched images and labels.
      """
      if strategy:
        if strategy.num_replicas_in_sync != self.config.num_devices:
          logging.warn(
              'Passed a strategy with %d devices, but expected'
              '%d devices.', strategy.num_replicas_in_sync,
              self.config.num_devices)
        dataset = strategy.distribute_datasets_from_function(self._build)
      else:
        dataset = self._build()
  
      return dataset
  
    def _build(
        self,
        input_context: Optional[tf.distribute.InputContext] = None
    ) -> tf.data.Dataset:
      """Construct a dataset end-to-end and return it.
  
      Args:
        input_context: An optional context provided by `tf.distribute` for
          cross-replica training.
  
      Returns:
        A TensorFlow dataset outputting batched images and labels.
      """
      builders = {
          'tfds': self.load_tfds,
          'records': self.load_records,
          'synthetic': self.load_synthetic,
      }
  
      builder = builders.get(self.config.builder, None)
  
      if builder is None:
        raise ValueError('Unknown builder type {}'.format(self.config.builder))
  
      self.input_context = input_context
      dataset = builder()
      dataset = self.pipeline(dataset)
  
      return dataset
  
    def load_tfds(self) -> tf.data.Dataset:
      """Return a dataset loading files from TFDS."""
  
      logging.info('Using TFDS to load data.')
      builder = tfds.builder(self.config.name, data_dir=self.config.data_dir)
  
      if self.config.download:
        builder.download_and_prepare()
  
      decoders = {}
  
      if self.config.skip_decoding:
        decoders['image'] = tfds.decode.SkipDecoding()
  
      read_config = tfds.ReadConfig(
          interleave_cycle_length=10,
          interleave_block_length=1,
          input_context=self.input_context)
  
      dataset = builder.as_dataset(
          split=self.config.split,
          as_supervised=True,
          shuffle_files=True,
          decoders=decoders,
          read_config=read_config)
  
      return dataset
  
    def load_records(self) -> tf.data.Dataset:
      """Return a dataset loading files with TFRecords."""
      logging.info('Using TFRecords to load data.')
      if self.config.filenames is None:
        if self.config.data_dir is None:
          raise ValueError('Dataset must specify a path for the data files.')
  
        file_pattern = os.path.join(self.config.data_dir,
                                    '{}*'.format(self.config.split))
        dataset = tf.data.Dataset.list_files(file_pattern, shuffle=False)
      else:
        dataset = tf.data.Dataset.from_tensor_slices(self.config.filenames)
  
      return dataset
  
    def load_synthetic(self) -> tf.data.Dataset:
      """Return a dataset generating dummy synthetic data."""
      logging.info('Generating a synthetic dataset.')
  
      def generate_data(_):
        image = tf.zeros([self.image_size, self.image_size, self.num_channels],
                         dtype=self.dtype)
        label = tf.zeros([1], dtype=tf.int32)
        return image, label
  
      dataset = tf.data.Dataset.range(1)
      dataset = dataset.repeat()
      dataset = dataset.map(
          generate_data, num_parallel_calls=tf.data.experimental.AUTOTUNE)
      return dataset
  
    def pipeline(self, dataset: tf.data.Dataset) -> tf.data.Dataset:
      """Build a pipeline fetching, shuffling, and preprocessing the dataset.
  
      Args:
        dataset: A `tf.data.Dataset` that loads raw files.
  
      Returns:
        A TensorFlow dataset outputting batched images and labels.
      """
      if (self.config.builder != 'tfds' and self.input_context and
          self.input_context.num_input_pipelines > 1):
        dataset = dataset.shard(self.input_context.num_input_pipelines,
                                self.input_context.input_pipeline_id)
        logging.info(
            'Sharding the dataset: input_pipeline_id=%d '
            'num_input_pipelines=%d', self.input_context.num_input_pipelines,
            self.input_context.input_pipeline_id)
  
      if self.is_training and self.config.builder == 'records':
        # Shuffle the input files.
        dataset.shuffle(buffer_size=self.config.file_shuffle_buffer_size)
  
      if self.is_training and not self.config.cache:
        dataset = dataset.repeat()
  
      if self.config.builder == 'records':
        # Read the data from disk in parallel
        dataset = dataset.interleave(
            tf.data.TFRecordDataset,
            cycle_length=10,
            block_length=1,
            num_parallel_calls=tf.data.experimental.AUTOTUNE)
  
      if self.config.cache:
        dataset = dataset.cache()
  
      if self.is_training:
        dataset = dataset.shuffle(self.config.shuffle_buffer_size)
        dataset = dataset.repeat()
  
      # Parse, pre-process, and batch the data in parallel
      if self.config.builder == 'records':
        preprocess = self.parse_record
      else:
        preprocess = self.preprocess
      dataset = dataset.map(
          preprocess, num_parallel_calls=tf.data.experimental.AUTOTUNE)
  
      if self.input_context and self.config.num_devices > 1:
        if not self.config.use_per_replica_batch_size:
          raise ValueError(
              'The builder does not support a global batch size with more than '
              'one replica. Got {} replicas. Please set a '
              '`per_replica_batch_size` and enable '
              '`use_per_replica_batch_size=True`.'.format(
                  self.config.num_devices))
  
        # The batch size of the dataset will be multiplied by the number of
        # replicas automatically when strategy.distribute_datasets_from_function
        # is called, so we use local batch size here.
        dataset = dataset.batch(
            self.local_batch_size, drop_remainder=self.is_training)
      else:
        dataset = dataset.batch(
            self.global_batch_size, drop_remainder=self.is_training)
  
      # Prefetch overlaps in-feed with training
      dataset = dataset.prefetch(tf.data.experimental.AUTOTUNE)
  
      if self.config.tf_data_service:
        if not hasattr(tf.data.experimental, 'service'):
          raise ValueError('The tf_data_service flag requires Tensorflow version '
                           '>= 2.3.0, but the version is {}'.format(
                               tf.__version__))
        dataset = dataset.apply(
            tf.data.experimental.service.distribute(
                processing_mode='parallel_epochs',
                service=self.config.tf_data_service,
                job_name='resnet_train'))
        dataset = dataset.prefetch(buffer_size=tf.data.experimental.AUTOTUNE)
  
      return dataset
  
    def parse_record(self, record: tf.Tensor) -> Tuple[tf.Tensor, tf.Tensor]:
      """Parse an ImageNet record from a serialized string Tensor."""
      keys_to_features = {
          'image/encoded': tf.io.FixedLenFeature((), tf.string, ''),
          'image/format': tf.io.FixedLenFeature((), tf.string, 'jpeg'),
          'image/class/label': tf.io.FixedLenFeature([], tf.int64, -1),
          'image/class/text': tf.io.FixedLenFeature([], tf.string, ''),
          'image/object/bbox/xmin': tf.io.VarLenFeature(dtype=tf.float32),
          'image/object/bbox/ymin': tf.io.VarLenFeature(dtype=tf.float32),
          'image/object/bbox/xmax': tf.io.VarLenFeature(dtype=tf.float32),
          'image/object/bbox/ymax': tf.io.VarLenFeature(dtype=tf.float32),
          'image/object/class/label': tf.io.VarLenFeature(dtype=tf.int64),
      }
  
      parsed = tf.io.parse_single_example(record, keys_to_features)
  
      label = tf.reshape(parsed['image/class/label'], shape=[1])
  
      # Subtract one so that labels are in [0, 1000)
      label -= 1
  
      image_bytes = tf.reshape(parsed['image/encoded'], shape=[])
      image, label = self.preprocess(image_bytes, label)
  
      return image, label
  
    def preprocess(self, image: tf.Tensor,
                   label: tf.Tensor) -> Tuple[tf.Tensor, tf.Tensor]:
      """Apply image preprocessing and augmentation to the image and label."""
      if self.is_training:
        image = preprocessing.preprocess_for_train(
            image,
            image_size=self.image_size,
            mean_subtract=self.config.mean_subtract,
            standardize=self.config.standardize,
            dtype=self.dtype,
            augmenter=self.augmenter)
      else:
        image = preprocessing.preprocess_for_eval(
            image,
            image_size=self.image_size,
            num_channels=self.num_channels,
            mean_subtract=self.config.mean_subtract,
            standardize=self.config.standardize,
            dtype=self.dtype)
  
      label = tf.cast(label, tf.int32)
      if self.config.one_hot:
        label = tf.one_hot(label, self.num_classes)
        label = tf.reshape(label, [self.num_classes])
  
      return image, label
  
    @classmethod
    def from_params(cls, *args, **kwargs):
      """Construct a dataset builder from a default config and any overrides."""
      config = DatasetConfig.from_args(*args, **kwargs)
      return cls(config)
  