gan

gan.ipynb

def plot(arr):
    # plot in 3D
    fig = plt.figure()
    ax = fig.add_subplot(1, 1, 1, projection='3d')
    ax.voxels((arr>0.7),antialiased=False, linewidth=0.0)
    plt.show()
import os
import numpy as np

from keras.models import Sequential, Model
from keras.layers import BatchNormalization, Input, Dense, Reshape, Flatten, Dropout
from keras.optimizers import Adam, RMSprop, SGD
from keras.layers.advanced_activations import LeakyReLU

import pymrt as mrt
import pymrt.geometry

import matplotlib.pyplot as plt
from skimage import measure

class GAN(object):


    def __init__(self, width=128, height=128, depth=128):
        self.width = width
        self.height = height
        self.depth = depth
        self.size = width * height * depth
        self.shape = (self.width, self.height, self.depth)

        self.OPTIMIZERD = Adam(lr=0.0005, decay=0.01)
        self.OPTIMIZER = Adam(lr=0.001, decay=8e-04)

        self.G = self.generator()
        self.D = self.discriminator()
        self.D.compile(loss='binary_crossentropy', optimizer=self.OPTIMIZERD, metrics=['accuracy'])

        z = Input(shape=(self.size,))
        self.stacked_G_D = Model(z, self.D(self.G(z)))
        self.D.trainable = False
        self.stacked_G_D.compile(loss='binary_crossentropy', optimizer=self.OPTIMIZER)

    def generator(self):
        model = Sequential()
        model.add(Dense(128, input_shape=(self.size,)))
        model.add(Dropout(0.5)) # noise noise noise
        model.add(Dense(self.size, activation='tanh'))
        model.add(Reshape(self.shape))
       
        

        noise = Input(shape=(self.size,))
        img = model(noise)

        return Model(noise, img)

    def discriminator(self):
        model = Sequential()

        model.add(Flatten(input_shape=self.shape))
        model.add(Dense(128, input_shape=self.shape))
        model.add(LeakyReLU(alpha=0.2))
        model.add(Dense(64))
        model.add(LeakyReLU(alpha=0.2))
        model.add(Dense(1, activation='sigmoid'))
        model.summary()

        img = Input(shape=self.shape)
        validity = model(img)
        
        return Model(img, validity)



    def train(self, X_train, epochs=10000, batch=2, save_interval=200):
        valid = np.ones((batch, 1))
        fake = np.zeros((batch, 1))
        for cnt in range(epochs):

            ## train discriminator
            random_index = np.random.randint(0, len(X_train), batch)
            legit_images = X_train[random_index].reshape(batch, self.width, self.height, self.depth)
            
            # generate some noise
            gen_noise = np.random.normal(0, 1, (batch, self.size))
            # generate a batch of new images
            syntetic_images = self.G.predict(gen_noise)
            
      
            d_loss = 0.5 * np.array(self.D.train_on_batch(legit_images, valid))
            d_loss += 0.5 * np.array(self.D.train_on_batch(syntetic_images, fake))
                                    

            # train generator

            noise = np.random.normal(0, 1, (batch, self.size))

            g_loss = self.stacked_G_D.train_on_batch(noise, valid)

            print ("%d [D loss: %f, acc.: %.2f%%] [G loss: %f]" % (cnt, d_loss[0], 100*d_loss[1], g_loss))

            if cnt % save_interval == 0 and cnt > 0:
                plot(syntetic_images[0])



    
if __name__ == '__main__':
    grid_size = 16

    arr = mrt.geometry.sphere(grid_size, 5, 0.5)
    #arr2 = mrt.geometry.cylinder(64, 16, 16, 0)

    X_train = np.asarray([arr,arr]).astype(np.float32)
    # Rescale -1 to 1
    #X_train = (X_train.astype(np.float32) - 127.5) / 127.5
    #X_train = np.expand_dims(X_train, axis=3)
    if True:
        plot(arr)
    gan = GAN(width=grid_size, height=grid_size, depth=grid_size)
    gan.train(X_train)