# Path: disentanglement_lib/data/ground_truth/ground_truth_data.py
# class GroundTruthData(object):
#   def num_factors(self):
#   def factors_num_values(self):
#   def observation_shape(self):
#   def sample_factors(self, num, random_state):
#   def sample_observations_from_factors(self, factors, random_state):
#   def sample(self, num, random_state):
#   def sample_observations(self, num, random_state):
# 
# Path: disentanglement_lib/data/ground_truth/util.py
# def tf_data_set_from_ground_truth_data(ground_truth_data, random_seed):
#   def generator():
#   def __init__(self, factor_sizes, latent_factor_indices):
#   def num_latent_factors(self):
#   def sample_latent_factors(self, num, random_state):
#   def sample_all_factors(self, latent_factors, random_state):
#   def _sample_factor(self, i, num, random_state):
#   def __init__(self, factor_sizes, features):
#   def features_to_index(self, features):
#   def _features_to_state_space_index(self, features):
# class SplitDiscreteStateSpace(object):
# class StateSpaceAtomIndex(object):


import os
import numpy as np
import PIL
from disentanglement_lib.data.ground_truth import ground_truth_data
from disentanglement_lib.data.ground_truth import util
from six.moves import range
from tensorflow.compat.v1 import gfile

  The data set was originally introduced in "beta-VAE: Learning Basic Visual
  Concepts with a Constrained Variational Framework" and can be downloaded from
  https://github.com/deepmind/dsprites-dataset.

  The ground-truth factors of variation are (in the default setting):
  0 - shape (3 different values)
  1 - scale (6 different values)
  2 - orientation (40 different values)
  3 - position x (32 different values)
  4 - position y (32 different values)
  """

  def __init__(self, latent_factor_indices=None):
    # By default, all factors (including shape) are considered ground truth
    # factors.
    if latent_factor_indices is None:
      latent_factor_indices = list(range(6))
    self.latent_factor_indices = latent_factor_indices
    self.data_shape = [64, 64, 1]
    # Load the data so that we can sample from it.
    with gfile.Open(DSPRITES_PATH, "rb") as data_file:
      # Data was saved originally using python2, so we need to set the encoding.
      data = np.load(data_file, encoding="latin1", allow_pickle=True)
      self.images = np.array(data["imgs"])
      self.factor_sizes = np.array(
          data["metadata"][()]["latents_sizes"], dtype=np.int64)
    self.full_factor_sizes = [1, 3, 6, 40, 32, 32]
    self.factor_bases = np.prod(self.factor_sizes) / np.cumprod(
        self.factor_sizes)

# Path: disentanglement_lib/data/ground_truth/ground_truth_data.py
# class GroundTruthData(object):
#   def num_factors(self):
#   def factors_num_values(self):
#   def observation_shape(self):
#   def sample_factors(self, num, random_state):
#   def sample_observations_from_factors(self, factors, random_state):
#   def sample(self, num, random_state):
#   def sample_observations(self, num, random_state):
# 
# Path: disentanglement_lib/data/ground_truth/util.py
# def tf_data_set_from_ground_truth_data(ground_truth_data, random_seed):
#   def generator():
#   def __init__(self, factor_sizes, latent_factor_indices):
#   def num_latent_factors(self):
#   def sample_latent_factors(self, num, random_state):
#   def sample_all_factors(self, latent_factors, random_state):
#   def _sample_factor(self, i, num, random_state):
#   def __init__(self, factor_sizes, features):
#   def features_to_index(self, features):
#   def _features_to_state_space_index(self, features):
# class SplitDiscreteStateSpace(object):
# class StateSpaceAtomIndex(object):


# Path: disentanglement_lib/data/ground_truth/dsprites.py
import os
import numpy as np
import PIL
from disentanglement_lib.data.ground_truth import ground_truth_data
from disentanglement_lib.data.ground_truth import util
from six.moves import range
from tensorflow.compat.v1 import gfile


  The data set was originally introduced in "beta-VAE: Learning Basic Visual
  Concepts with a Constrained Variational Framework" and can be downloaded from
  https://github.com/deepmind/dsprites-dataset.

  The ground-truth factors of variation are (in the default setting):
  0 - shape (3 different values)
  1 - scale (6 different values)
  2 - orientation (40 different values)
  3 - position x (32 different values)
  4 - position y (32 different values)
  """

  def __init__(self, latent_factor_indices=None):
    # By default, all factors (including shape) are considered ground truth
    # factors.
    if latent_factor_indices is None:
      latent_factor_indices = list(range(6))
    self.latent_factor_indices = latent_factor_indices
    self.data_shape = [64, 64, 1]
    # Load the data so that we can sample from it.
    with gfile.Open(DSPRITES_PATH, "rb") as data_file:
      # Data was saved originally using python2, so we need to set the encoding.
      data = np.load(data_file, encoding="latin1", allow_pickle=True)
      self.images = np.array(data["imgs"])
      self.factor_sizes = np.array(
          data["metadata"][()]["latents_sizes"], dtype=np.int64)
    self.full_factor_sizes = [1, 3, 6, 40, 32, 32]
    self.factor_bases = np.prod(self.factor_sizes) / np.cumprod(
        self.factor_sizes)

    self.state_space = util.SplitDiscreteStateSpace(self.factor_sizes,