# Path: awesome_gans/datasets.py
# class MNISTDataSet:
#     def __init__(self, use_split=False, split_rate=0.15, random_state=42, ds_path=None):
#         self.use_split = use_split
#         self.split_rate = split_rate
#         self.random_state = random_state
# 
#         self.ds_path = ds_path
# 
#         try:
#             assert self.ds_path
#         except AssertionError:
#             raise AssertionError("[-] MNIST DataSet Path is required!")
# 
#         from tensorflow.examples.tutorials.mnist import input_data
# 
#         self.data = input_data.read_data_sets(self.ds_path, one_hot=True)  # download MNIST
# 
#         # training data
#         self.train_data = self.data.train
# 
#         self.train_images = self.train_data.images
#         self.train_labels = self.train_data.labels
#         self.valid_images = None
#         self.valid_labels = None
# 
#         # test data
#         self.test_data = self.data.test
# 
#         self.test_images = self.test_data.images
#         self.test_labels = self.test_data.labels
# 
#         # split training data set into train, valid
#         if self.use_split:
#             self.train_images, self.valid_images, self.train_labels, self.valid_labels = train_test_split(
#                 self.train_images, self.train_labels, test_size=self.split_rate, random_state=self.random_state
#             )


import time
import numpy as np
import tensorflow as tf
import awesome_gans.gan.gan_model as gan
import awesome_gans.image_utils as iu
from awesome_gans.datasets import MNISTDataSet as DataSet
    # GPU configure
    config = tf.ConfigProto()
    config.gpu_options.allow_growth = True

    with tf.Session(config=config) as s:
        # GAN Model
        model = gan.GAN(s)

        # Initializing
        s.run(tf.global_variables_initializer())

        # Load model & Graph & Weights
        saved_global_step = 0
        ckpt = tf.train.get_checkpoint_state('./model/')
        if ckpt and ckpt.model_checkpoint_path:
            model.saver.restore(s, ckpt.model_checkpoint_path)

            saved_global_step = int(ckpt.model_checkpoint_path.split('/')[-1].split('-')[-1])
            print("[+] global step : %s" % saved_global_step, " successfully loaded")
        else:
            print('[-] No checkpoint file found')

        d_loss = 0.0
        d_overpowered = False
        for global_step in range(saved_global_step, train_step['global_step']):
            batch_x, _ = mnist.train.next_batch(model.batch_size)
            batch_x = batch_x.reshape(-1, model.n_input)
            batch_z = np.random.uniform(-1.0, 1.0, size=[model.batch_size, model.z_dim]).astype(np.float32)

            # Update D network

# Path: awesome_gans/datasets.py
# class MNISTDataSet:
#     def __init__(self, use_split=False, split_rate=0.15, random_state=42, ds_path=None):
#         self.use_split = use_split
#         self.split_rate = split_rate
#         self.random_state = random_state
# 
#         self.ds_path = ds_path
# 
#         try:
#             assert self.ds_path
#         except AssertionError:
#             raise AssertionError("[-] MNIST DataSet Path is required!")
# 
#         from tensorflow.examples.tutorials.mnist import input_data
# 
#         self.data = input_data.read_data_sets(self.ds_path, one_hot=True)  # download MNIST
# 
#         # training data
#         self.train_data = self.data.train
# 
#         self.train_images = self.train_data.images
#         self.train_labels = self.train_data.labels
#         self.valid_images = None
#         self.valid_labels = None
# 
#         # test data
#         self.test_data = self.data.test
# 
#         self.test_images = self.test_data.images
#         self.test_labels = self.test_data.labels
# 
#         # split training data set into train, valid
#         if self.use_split:
#             self.train_images, self.valid_images, self.train_labels, self.valid_labels = train_test_split(
#                 self.train_images, self.train_labels, test_size=self.split_rate, random_state=self.random_state
#             )


# Path: awesome_gans/gan/gan_train.py
import time
import numpy as np
import tensorflow as tf
import awesome_gans.gan.gan_model as gan
import awesome_gans.image_utils as iu
from awesome_gans.datasets import MNISTDataSet as DataSet

    # GPU configure
    config = tf.ConfigProto()
    config.gpu_options.allow_growth = True

    with tf.Session(config=config) as s:
        # GAN Model
        model = gan.GAN(s)

        # Initializing
        s.run(tf.global_variables_initializer())

        # Load model & Graph & Weights
        saved_global_step = 0
        ckpt = tf.train.get_checkpoint_state('./model/')
        if ckpt and ckpt.model_checkpoint_path:
            model.saver.restore(s, ckpt.model_checkpoint_path)

            saved_global_step = int(ckpt.model_checkpoint_path.split('/')[-1].split('-')[-1])
            print("[+] global step : %s" % saved_global_step, " successfully loaded")
        else:
            print('[-] No checkpoint file found')

        d_loss = 0.0
        d_overpowered = False
        for global_step in range(saved_global_step, train_step['global_step']):
            batch_x, _ = mnist.train.next_batch(model.batch_size)
            batch_x = batch_x.reshape(-1, model.n_input)
            batch_z = np.random.uniform(-1.0, 1.0, size=[model.batch_size, model.z_dim]).astype(np.float32)

            # Update D network

            if not d_overpowered: