text: # Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved.
  #
  # This work is licensed under the Creative Commons Attribution-NonCommercial
  # 4.0 International License. To view a copy of this license, visit
  # http://creativecommons.org/licenses/by-nc/4.0/ or send a letter to
  # Creative Commons, PO Box 1866, Mountain View, CA 94042, USA.
  
  """Multi-resolution input data pipeline."""
  
  import os
  import glob
  import numpy as np
  import tensorflow as tf
  import dnnlib
  import dnnlib.tflib as tflib
  
  #----------------------------------------------------------------------------
  # Parse individual image from a tfrecords file.
  
  def parse_tfrecord_tf(record):
      features = tf.parse_single_example(record, features={
          'shape': tf.FixedLenFeature([3], tf.int64),
          'data': tf.FixedLenFeature([], tf.string)})
      data = tf.decode_raw(features['data'], tf.uint8)
      return tf.reshape(data, features['shape'])
  
  def parse_tfrecord_np(record):
      ex = tf.train.Example()
      ex.ParseFromString(record)
      shape = ex.features.feature['shape'].int64_list.value # temporary pylint workaround # pylint: disable=no-member
      data = ex.features.feature['data'].bytes_list.value[0] # temporary pylint workaround # pylint: disable=no-member
      return np.fromstring(data, np.uint8).reshape(shape)
  
  #----------------------------------------------------------------------------
  # Dataset class that loads data from tfrecords files.
  
  class TFRecordDataset:
      def __init__(self,
          tfrecord_dir,               # Directory containing a collection of tfrecords files.
          resolution      = None,     # Dataset resolution, None = autodetect.
          label_file      = None,     # Relative path of the labels file, None = autodetect.
          max_label_size  = 0,        # 0 = no labels, 'full' = full labels, <int> = N first label components.
          repeat          = True,     # Repeat dataset indefinitely.
          shuffle_mb      = 4096,     # Shuffle data within specified window (megabytes), 0 = disable shuffling.
          prefetch_mb     = 2048,     # Amount of data to prefetch (megabytes), 0 = disable prefetching.
          buffer_mb       = 256,      # Read buffer size (megabytes).
          num_threads     = 2):       # Number of concurrent threads.
  
          self.tfrecord_dir       = tfrecord_dir
          self.resolution         = None
          self.resolution_log2    = None
          self.shape              = []        # [channel, height, width]
          self.dtype              = 'uint8'
          self.dynamic_range      = [0, 255]
          self.label_file         = label_file
          self.label_size         = None      # [component]
          self.label_dtype        = None
          self._np_labels         = None
          self._tf_minibatch_in   = None
          self._tf_labels_var     = None
          self._tf_labels_dataset = None
          self._tf_datasets       = dict()
          self._tf_iterator       = None
          self._tf_init_ops       = dict()
          self._tf_minibatch_np   = None
          self._cur_minibatch     = -1
          self._cur_lod           = -1
  
          # List tfrecords files and inspect their shapes.
          assert os.path.isdir(self.tfrecord_dir)
          tfr_files = sorted(glob.glob(os.path.join(self.tfrecord_dir, '*.tfrecords')))
          assert len(tfr_files) >= 1
          tfr_shapes = []
          for tfr_file in tfr_files:
              tfr_opt = tf.python_io.TFRecordOptions(tf.python_io.TFRecordCompressionType.NONE)
              for record in tf.python_io.tf_record_iterator(tfr_file, tfr_opt):
                  tfr_shapes.append(parse_tfrecord_np(record).shape)
                  break
  
          # Autodetect label filename.
          if self.label_file is None:
              guess = sorted(glob.glob(os.path.join(self.tfrecord_dir, '*.labels')))
              if len(guess):
                  self.label_file = guess[0]
          elif not os.path.isfile(self.label_file):
              guess = os.path.join(self.tfrecord_dir, self.label_file)
              if os.path.isfile(guess):
                  self.label_file = guess
  
          # Determine shape and resolution.
          max_shape = max(tfr_shapes, key=np.prod)
          self.resolution = resolution if resolution is not None else max_shape[1]
          self.resolution_log2 = int(np.log2(self.resolution))
          self.shape = [max_shape[0], self.resolution, self.resolution]
          tfr_lods = [self.resolution_log2 - int(np.log2(shape[1])) for shape in tfr_shapes]
          assert all(shape[0] == max_shape[0] for shape in tfr_shapes)
          assert all(shape[1] == shape[2] for shape in tfr_shapes)
          assert all(shape[1] == self.resolution // (2**lod) for shape, lod in zip(tfr_shapes, tfr_lods))
          assert all(lod in tfr_lods for lod in range(self.resolution_log2 - 1))
  
          # Load labels.
          assert max_label_size == 'full' or max_label_size >= 0
          self._np_labels = np.zeros([1<<20, 0], dtype=np.float32)
          if self.label_file is not None and max_label_size != 0:
              self._np_labels = np.load(self.label_file)
              assert self._np_labels.ndim == 2
          if max_label_size != 'full' and self._np_labels.shape[1] > max_label_size:
              self._np_labels = self._np_labels[:, :max_label_size]
          self.label_size = self._np_labels.shape[1]
          self.label_dtype = self._np_labels.dtype.name
  
          # Build TF expressions.
          with tf.name_scope('Dataset'), tf.device('/cpu:0'):
              self._tf_minibatch_in = tf.placeholder(tf.int64, name='minibatch_in', shape=[])
              self._tf_labels_var = tflib.create_var_with_large_initial_value(self._np_labels, name='labels_var')
              self._tf_labels_dataset = tf.data.Dataset.from_tensor_slices(self._tf_labels_var)
              for tfr_file, tfr_shape, tfr_lod in zip(tfr_files, tfr_shapes, tfr_lods):
                  if tfr_lod < 0:
                      continue
                  dset = tf.data.TFRecordDataset(tfr_file, compression_type='', buffer_size=buffer_mb<<20)
                  dset = dset.map(parse_tfrecord_tf, num_parallel_calls=num_threads)
                  dset = tf.data.Dataset.zip((dset, self._tf_labels_dataset))
                  bytes_per_item = np.prod(tfr_shape) * np.dtype(self.dtype).itemsize
                  if shuffle_mb > 0:
                      dset = dset.shuffle(((shuffle_mb << 20) - 1) // bytes_per_item + 1)
                  if repeat:
                      dset = dset.repeat()
                  if prefetch_mb > 0:
                      dset = dset.prefetch(((prefetch_mb << 20) - 1) // bytes_per_item + 1)
                  dset = dset.batch(self._tf_minibatch_in)
                  self._tf_datasets[tfr_lod] = dset
              self._tf_iterator = tf.data.Iterator.from_structure(self._tf_datasets[0].output_types, self._tf_datasets[0].output_shapes)
              self._tf_init_ops = {lod: self._tf_iterator.make_initializer(dset) for lod, dset in self._tf_datasets.items()}
  
      # Use the given minibatch size and level-of-detail for the data returned by get_minibatch_tf().
      def configure(self, minibatch_size, lod=0):
          lod = int(np.floor(lod))
          assert minibatch_size >= 1 and lod in self._tf_datasets
          if self._cur_minibatch != minibatch_size or self._cur_lod != lod:
              self._tf_init_ops[lod].run({self._tf_minibatch_in: minibatch_size})
              self._cur_minibatch = minibatch_size
              self._cur_lod = lod
  
      # Get next minibatch as TensorFlow expressions.
      def get_minibatch_tf(self): # => images, labels
          return self._tf_iterator.get_next()
  
      # Get next minibatch as NumPy arrays.
      def get_minibatch_np(self, minibatch_size, lod=0): # => images, labels
          self.configure(minibatch_size, lod)
          if self._tf_minibatch_np is None:
              self._tf_minibatch_np = self.get_minibatch_tf()
          return tflib.run(self._tf_minibatch_np)
  
      # Get random labels as TensorFlow expression.
      def get_random_labels_tf(self, minibatch_size): # => labels
          if self.label_size > 0:
              with tf.device('/cpu:0'):
                  return tf.gather(self._tf_labels_var, tf.random_uniform([minibatch_size], 0, self._np_labels.shape[0], dtype=tf.int32))
          return tf.zeros([minibatch_size, 0], self.label_dtype)
  
      # Get random labels as NumPy array.
      def get_random_labels_np(self, minibatch_size): # => labels
          if self.label_size > 0:
              return self._np_labels[np.random.randint(self._np_labels.shape[0], size=[minibatch_size])]
          return np.zeros([minibatch_size, 0], self.label_dtype)
  
  #----------------------------------------------------------------------------
  # Base class for datasets that are generated on the fly.
  
  class SyntheticDataset:
      def __init__(self, resolution=1024, num_channels=3, dtype='uint8', dynamic_range=[0,255], label_size=0, label_dtype='float32'):
          self.resolution         = resolution
          self.resolution_log2    = int(np.log2(resolution))
          self.shape              = [num_channels, resolution, resolution]
          self.dtype              = dtype
          self.dynamic_range      = dynamic_range
          self.label_size         = label_size
          self.label_dtype        = label_dtype
          self._tf_minibatch_var  = None
          self._tf_lod_var        = None
          self._tf_minibatch_np   = None
          self._tf_labels_np      = None
  
          assert self.resolution == 2 ** self.resolution_log2
          with tf.name_scope('Dataset'):
              self._tf_minibatch_var = tf.Variable(np.int32(0), name='minibatch_var')
              self._tf_lod_var = tf.Variable(np.int32(0), name='lod_var')
  
      def configure(self, minibatch_size, lod=0):
          lod = int(np.floor(lod))
          assert minibatch_size >= 1 and 0 <= lod <= self.resolution_log2
          tflib.set_vars({self._tf_minibatch_var: minibatch_size, self._tf_lod_var: lod})
  
      def get_minibatch_tf(self): # => images, labels
          with tf.name_scope('SyntheticDataset'):
              shrink = tf.cast(2.0 ** tf.cast(self._tf_lod_var, tf.float32), tf.int32)
              shape = [self.shape[0], self.shape[1] // shrink, self.shape[2] // shrink]
              images = self._generate_images(self._tf_minibatch_var, self._tf_lod_var, shape)
              labels = self._generate_labels(self._tf_minibatch_var)
              return images, labels
  
      def get_minibatch_np(self, minibatch_size, lod=0): # => images, labels
          self.configure(minibatch_size, lod)
          if self._tf_minibatch_np is None:
              self._tf_minibatch_np = self.get_minibatch_tf()
          return tflib.run(self._tf_minibatch_np)
  
      def get_random_labels_tf(self, minibatch_size): # => labels
          with tf.name_scope('SyntheticDataset'):
              return self._generate_labels(minibatch_size)
  
      def get_random_labels_np(self, minibatch_size): # => labels
          self.configure(minibatch_size)
          if self._tf_labels_np is None:
              self._tf_labels_np = self.get_random_labels_tf(minibatch_size)
          return tflib.run(self._tf_labels_np)
  
      def _generate_images(self, minibatch, lod, shape): # to be overridden by subclasses # pylint: disable=unused-argument
          return tf.zeros([minibatch] + shape, self.dtype)
  
      def _generate_labels(self, minibatch): # to be overridden by subclasses
          return tf.zeros([minibatch, self.label_size], self.label_dtype)
  
  #----------------------------------------------------------------------------
  # Helper func for constructing a dataset object using the given options.
  
  def load_dataset(class_name='training.dataset.TFRecordDataset', data_dir=None, verbose=False, **kwargs):
      adjusted_kwargs = dict(kwargs)
      if 'tfrecord_dir' in adjusted_kwargs and data_dir is not None:
          adjusted_kwargs['tfrecord_dir'] = os.path.join(data_dir, adjusted_kwargs['tfrecord_dir'])
      if verbose:
          print('Streaming data using %s...' % class_name)
      dataset = dnnlib.util.get_obj_by_name(class_name)(**adjusted_kwargs)
      if verbose:
          print('Dataset shape =', np.int32(dataset.shape).tolist())
          print('Dynamic range =', dataset.dynamic_range)
          print('Label size    =', dataset.label_size)
      return dataset
  
  #----------------------------------------------------------------------------
  