test1

Author: MinaRe

README.md

@@ -0,0 +1,2 @@
+echo # voxelGAN
+# voxelGAN

config.py

@@ -0,0 +1,23 @@
+import os
+from keras import backend as K
+
+K.set_image_data_format('channels_last')
+
+# general
+data_dir = 'data/train/**/*t1.nii.gz'
+predict_dir = 'data/predict/**/*t1.nii.gz'
+checkpoint_dir = 'checkpoints/'
+
+directory = os.path.dirname(checkpoint_dir)
+if not os.path.exists(directory):
+    os.makedirs(directory)
+
+# bdsscgan
+input = 4
+output = 4
+size = 32
+epochs = 50
+kernel_depth = 32
+
+checkpoint_gen_name = checkpoint_dir + 'gen.hdf5'
+checkpoint_disc_name = checkpoint_dir + 'disc.hdf5'

model.py

@@ -0,0 +1,73 @@
+import tensorflow as tf
+#keras = tf.keras
+from keras.models import Model
+from keras.layers import Bidirectional, Input, Concatenate, Cropping3D, Dense, Flatten, TimeDistributed, ConvLSTM2D, LeakyReLU
+from keras.layers.core import Dropout, Activation, Reshape
+from keras.layers.convolutional import Conv3D, MaxPooling3D, UpSampling3D, Cropping3D
+from keras.layers.normalization import BatchNormalization
+from keras.layers.merge import concatenate
+import keras.backend as K
+
+
+def conv_layer(layer, depth, size):
+    conv = Conv3D(depth, size, padding='same')(layer)
+    conv = LeakyReLU(0.2)(conv)
+    return BatchNormalization()(conv)
+    
+def Generator(input_shape, output, kernel_depth, kernel_size=3):
+    # 32x32x32x4
+    input = Input(shape=input_shape)
+    
+    conv_32 = conv_layer(input, 1 * kernel_depth, kernel_size)
+    pool_16 = MaxPooling3D()(conv_32)
+
+    conv_16 = conv_layer(pool_16, 2 * kernel_depth, kernel_size)
+    pool_8 = MaxPooling3D()(conv_16)
+
+    conv_8 = conv_layer(pool_8, 4 * kernel_depth, kernel_size)
+
+    up_16 = concatenate([UpSampling3D()(conv_8), conv_16])
+    up_conv_16 = conv_layer(up_16, 2 * kernel_depth, kernel_size)
+
+    up_32 = concatenate([UpSampling3D()(up_conv_16), conv_32])
+    up_conv_32 = conv_layer(up_32, 1 * kernel_depth, kernel_size)
+    
+    final1 = concatenate([up_conv_32, input])
+    final2 = Conv3D(output, 1, activation='softmax')(final1)
+    
+    model = Model(input, final2, name="Generator")
+    return model
+   
+def Discriminator(input_shape, generator_shape, kernel_depth, kernel_size=5):
+    real_input = Input(shape=input_shape)
+    generator_input = Input(shape=generator_shape)    
+    input = Concatenate()([real_input, generator_input])
+   
+    conv_32 = conv_layer(input, 1 * kernel_depth, kernel_size)
+    pool_16 = MaxPooling3D()(conv_32)
+
+    conv_16 = conv_layer(pool_16, 2 * kernel_depth, kernel_size)
+    pool_8 = MaxPooling3D()(conv_16)
+
+    conv_8 = conv_layer(pool_8, 4 * kernel_depth, kernel_size)
+    pool_4 = MaxPooling3D()(conv_8)
+    
+    x = Flatten()(pool_4)
+    x = Dense(2, activation="softmax")(x)
+    
+    model = Model([real_input, generator_input], x, name="Discriminator")
+    return model
+  
+def Combine(gen, disc, input_shape, new_sequence):
+    input = Input(shape=input_shape)
+    generated_image = gen(input)
+
+    reshaped = Reshape(new_sequence)(generated_image)
+    
+    DCGAN_output = disc([input, reshaped])
+
+    DCGAN = Model(inputs=[input],
+                  outputs=[generated_image, DCGAN_output],
+                  name="Combined")
+
+    return DCGAN

predict.py

@@ -0,0 +1,37 @@
+from keras.optimizers import Adadelta
+
+from model import *
+from utils import *
+from config import *
+
+opt = Adadelta()
+  
+gen = Generator((size, size, size, input), output, kernel_depth)
+gen.compile(loss='mae', optimizer=opt)
+gen.load_weights(checkpoint_gen_name)
+
+# List sequences  
+sequences = prepare_data(predict_dir)
+print(sequences)
+
+progbar = keras.utils.Progbar(len(sequences))
+
+for s in range(len(sequences)):
+    
+    
+    progbar.add(1)
+    sequence = sequences[s]
+    x, idx = load2(sequence, size)
+    y = []
+    
+    for i in range(len(x)):
+    
+        # gen
+        fake = gen.predict(x[i])
+        print(fake.shape)
+        
+        y.append(fake)
+        
+    store(sequence, y, idx)
+    
+    #save_image(x[i] / 2 + 0.5, y[i], re_shape(generated_y), prediction_dir + "test{}.png".format(s))

train.py

@@ -0,0 +1,88 @@
+from keras.callbacks import TensorBoard
+from keras.optimizers import Adadelta
+
+from model import *
+from utils import *
+from config import *
+
+# Create optimizers
+opt_dcgan = Adadelta()
+opt_discriminator = Adadelta()
+  
+gen = Generator((size, size, size, input), output, kernel_depth)
+gen.compile(loss='mae', optimizer=opt_discriminator)
+
+disc = Discriminator((size, size, size, input), (size, size, size, output), kernel_depth)
+disc.trainable = False
+
+combined = Combine(gen, disc, (size, size, size, input), (size, size, size, output))
+loss = [selective_crossentropy, 'binary_crossentropy']
+loss_weights = [10, 1]
+combined.compile(loss=loss, loss_weights=loss_weights, optimizer=opt_dcgan)
+
+disc.trainable = True
+disc.compile(loss='binary_crossentropy', optimizer=opt_discriminator)
+
+if os.path.isfile(checkpoint_gen_name):
+    gen.load_weights(checkpoint_gen_name)
+if os.path.isfile(checkpoint_disc_name):
+    disc.load_weights(checkpoint_disc_name)
+
+# List sequences  
+sequences = prepare_data(data_dir)
+print(sequences)
+
+real_y = np.reshape(np.array([0, 1]), (1, 2))
+fake_y = np.reshape(np.array([1, 0]), (1, 2))
+
+#log = open("train.log",'w')
+
+tensorlog = TensorBoard(log_dir='./logs', histogram_freq=0, batch_size=1, write_graph=True, write_grads=True, write_images=True)
+tensorlog.set_model(gen)
+
+for e in range(epochs):
+    print("Epoch {}".format(e))
+    random.shuffle(sequences)
+    
+    # select a fraction
+    train_offset = int(len(sequences) * 0.9)
+    train_sequence = sequences[:train_offset]
+    
+    progbar = keras.utils.Progbar(len(train_sequence))
+    
+    for s in range(len(train_sequence)):
+        
+        progbar.add(1)
+        sequence = train_sequence[s]
+        x, y, idx = load(sequence, size)
+        
+        for i in range(len(x)):
+        
+            # train disc on real
+            disc.train_on_batch([x[i], y[i]], real_y)
+        
+            # gen fake
+            fake = gen.predict(x[i])
+        
+            # train disc on fake
+            disc.train_on_batch([x[i], fake], fake_y)
+        
+            # train combined    
+            disc.trainable = False
+            combined.train_on_batch(x[i], [y[i], real_y])
+            disc.trainable = True
+            
+            #log.write(str(e) + ", " + str(s) + ", " + str(dr_loss) + ", " + str(df_loss) + ", " + str(g_loss[0]) + ", " + str(g_loss[1]) + ", " + str(opt_dcgan.get_config()["lr"]) + "\n")
+            
+    # output random result
+    #val_sequence = sequences[train_offset:]
+    #generated_y = gen.predict(x[random_index])
+    #save_image(strip(x[random_index]) / 2 + 0.5, y[random_index], re_shape(generated_y), "validation/e{}_{}.png".format(e, s))
+        
+    # save weights
+    gen.save_weights(checkpoint_gen_name, overwrite=True)
+    disc.save_weights(checkpoint_disc_name, overwrite=True)
+    
+    tensorlog.on_epoch_end(e)
+    
+tensorlog.on_train_end()