code: 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.
      max_images=None,  # Maximum number of images to use, None = use all images.
      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 = []  # [channels, height, width]
      self.dtype = "uint8"
      self.dynamic_range = [0, 255]
      self.label_file = label_file
      self.label_size = None  # components
      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(self.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 << 30, 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]
      if (
          max_images is not None
          and self._np_labels.shape[0] > max_images
      ):
          self._np_labels = self._np_labels[:max_images]
      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,
              )
              if max_images is not None:
                  dset = dset.take(max_images)
              dset = dset.map(
                  self.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()
          }
  
file_path: insightface/reconstruction/ostec/external/stylegan2/training/dataset.py
imports: import os
  import glob
  import numpy
  import tensorflow
  import dnnlib
  import dnnlib.tflib
label_desc: Use image node_id 1 to create a new TFRecordDataset object  with example: obj = TFRecordDataset(tfrecord_dir, resolution, label_file, max_label_size, max_images, repeat, shuffle_mb, prefetch_mb, buffer_mb, num_threads)