Basen on https://www.tensorflow.org/tutorials/generative/dcgan

#1 Libraries
from __future__ import absolute_import, division, print_function, unicode_literals
import sys
import glob
import imageio
import os
import time
import datetime
import pandas as pd
import numpy as np
from math import sqrt

import matplotlib.pyplot as plt

from pathlib import Path
from tensorflow.keras import layers
import tensorflow as tf

from src.datagens import WSIDatagen, DataSource
from src.training.experiments import SimpleExperiment, SummaryWriter
from src.training.utils import ExperimentLevel
from src.training.models import NASNetModel, ContextNASNetModel
from src.training.metrics import F1
from src.config import SLIDE_DIR

from summary import Summary
from gan_utils import plot_loss, sample_images, set_up_checkpoint

load
load
load
load
/mnt/data/home/xcechak/crc_ml_master/models/evaluations

# Configuration
noise_dim = 100
num_examples_to_generate = 16
BATCH_SIZE = 64
EPOCHS = 200000
STEPS_PER_EPOCH = 64
NUMBER_OF_BATCHES = 64
global SEED
SEED = tf.random.normal([num_examples_to_generate, noise_dim])

config = { 
    'sample_interval': 5,
    'save_interval': 50,
    'noise_dim': 100,
    'num_examples_to_generate': 16,
    'batch_size': 64,
    'epochs': 200000,
    'steps_per_epoch': 64,
    'nuumber_of_batches': 64,
}

# Nicely formatted time string
def hms_string(sec_elapsed):
    h = int(sec_elapsed / (60 * 60))
    m = int((sec_elapsed % (60 * 60)) / 60)
    s = sec_elapsed % 60
    return "{}:{:>02}:{:>05.2f}".format(h, m, s)

params = {
 'data': {
     'dataset': '/mnt/data/crc_ml/data/processed/Prostata/level1/context3/t96px/all/datasets/prospective.h5',
     'meta': {'context_size': 1, 'level': 1, 'tile_size': 96}},
}

# CREATE DATASOURCE
ds = DataSource()
ds.add_dataset(params['data']['dataset'], 'train')
train_ds, valid_ds = ds.split(0.03)

test_ds = DataSource()
test_ds.add_dataset(params['data']['dataset'], 'test')

WARNING:tensorflow:From /mnt/data/home/xcechak/crc_ml_master/src/datagens/_datasources.py:26: HDF5Matrix.__init__ (from tensorflow.python.keras.utils.io_utils) is deprecated and will be removed after 2020-05-30.
Instructions for updating:
Training with HDF5Matrix is not optimized for performance. Instead, we recommend using https://github.com/tensorflow/io to load your HDF5 data into a tf.data Dataset and passing that dataset to Keras.

# CREATE DATAGEN
datagen = WSIDatagen(**params['data']['meta'], slide_dir=SLIDE_DIR, batch_size=BATCH_SIZE)

# RETRIEVE GENERATORS (877 MB)
""" with train_augment_type=None argument the datagen produces only non-augmented images (no need to use .unaugmented() func) """
train_gen, valid_gen = datagen.get_training_generator((96,96), train_ds, train_steps = NUMBER_OF_BATCHES, valid_ds=valid_ds, valid_steps=1024, train_augment_type=None)

call on_epoch_end once more to get only selected labels
call on_epoch_end once more to get only selected labels

model_name = '96_vanilla_gan_{}'.format(datetime.datetime.now().strftime("%Y-%m-%d_%H:%M:%S"))
summary = Summary(model_name, config)
sw = summary.summary_writer
sw.update_log()

/mnt/data/home/xcechak/crc_ml_master/models/evaluations/96_vanilla_gan_2020-12-04_11:23:1404_11:23:14
Summary init done

class Vanilla_gan():
    def __init__(self, config, summary):
        self.summary = summary
#         ...
        
        # Build the generator and critic
        self.generator = self.make_generator_model()
        self.discriminator = self.make_discriminator_model()
        
        
        noise = tf.random.normal([1, 100])
        generated_image = self.generator(noise, training=False)
#         
        print('generated_image.shape ', generated_image.shape)
        print('generated_image_min ', np.min((generated_image[0, :, :,:] + 1) / 2))
        print('generated_image_max ', np.max((generated_image[0, :, :,:] + 1) / 2))
        plt.imshow((generated_image[0, :, :, :] + 1) / 2)
        plt.show()
#         
        decision = self.discriminator(generated_image)
        print("decision:", decision)
        
        self.cross_entropy = tf.keras.losses.BinaryCrossentropy(from_logits=True)
        
        self.generator_optimizer = tf.keras.optimizers.Adam(1e-4)
        self.discriminator_optimizer = tf.keras.optimizers.Adam(1e-4)

        self.summary.summary_writer.set_value('generator', value=self.generator.to_json())
        self.summary.summary_writer.update_log()
        self.summary.summary_writer.set_value('discriminator', value=self.discriminator.to_json())
        self.summary.summary_writer.update_log()
        
        set_up_checkpoint(self)

        global SEED
        SEED = tf.random.normal([num_examples_to_generate, noise_dim])

    
    def make_generator_model(self):
        model = tf.keras.Sequential()
        model.add(layers.Dense(units=8*8*256, use_bias=False, input_shape=(100,)))
        model.add(layers.BatchNormalization())
        model.add(layers.LeakyReLU())

        model.add(layers.Reshape((8, 8, 256)))
        assert model.output_shape == (None, 8, 8, 256) # Note: None is the batch size

        model.add(layers.Conv2DTranspose(128, (5, 5), strides=(1, 1), padding='same', use_bias=False))
        assert model.output_shape == (None, 8, 8, 128)
        model.add(layers.BatchNormalization())
        model.add(layers.LeakyReLU())

        model.add(layers.Conv2DTranspose(64, (5, 5), strides=(3, 3), padding='same', use_bias=False))
        assert model.output_shape == (None, 24, 24, 64)
        model.add(layers.BatchNormalization())
        model.add(layers.LeakyReLU())

        model.add(layers.Conv2DTranspose(64, (5, 5), strides=(2, 2), padding='same', use_bias=False))
        assert model.output_shape == (None, 48, 48, 64)
        model.add(layers.BatchNormalization())
        model.add(layers.LeakyReLU())

        model.add(layers.Conv2DTranspose(3, (5, 5), strides=(2, 2), padding='same', use_bias=False, activation='tanh'))
        assert model.output_shape == (None, 96, 96, 3)

        return model

    def make_discriminator_model(self):
        model = tf.keras.Sequential()
        model.add(layers.Conv2D(filters = 64, kernel_size = (5, 5), strides=(2, 2), padding='same', input_shape=[96, 96, 3]))
#         model.add(layers.LayerNormalization())
        model.add(layers.LeakyReLU())
        model.add(layers.Dropout(0.3))

        model.add(layers.Conv2D(128, (5, 5), strides=(2, 2), padding='same'))
#         model.add(layers.LayerNormalization())
        model.add(layers.LeakyReLU())
        model.add(layers.Dropout(0.3))

        model.add(layers.Flatten())
        model.add(layers.Dense(1))

        return model

    def discriminator_loss(self, real_output, fake_output):
        real_loss = self.cross_entropy(tf.ones_like(real_output), real_output)
        fake_loss = self.cross_entropy(tf.zeros_like(fake_output), fake_output)
        # should not be total_loss = (real_loss + fake_loss) * 0.5 ?
        total_loss = real_loss + fake_loss
        return total_loss


    def generator_loss(self, fake_output):
        return self.cross_entropy(tf.ones_like(fake_output), fake_output)


    @tf.function
    def train_step_generator(self, images):
        print('images.shape: ', images.shape)
        noise = tf.random.normal([BATCH_SIZE, noise_dim])

        with tf.GradientTape() as gen_tape, tf.GradientTape() as disc_tape:
            generated_images = self.generator(noise, training=True)
            fake_output = self.discriminator(generated_images, training=False)

            gen_loss = self.generator_loss(fake_output)  

        gradients_of_generator = gen_tape.gradient(gen_loss, self.generator.trainable_variables)
        self.generator_optimizer.apply_gradients(zip(gradients_of_generator, self.generator.trainable_variables))
        return gen_loss
    

    @tf.function
    def train_step_discriminator(self, images):
        print('images.shape: ', images.shape)
        noise = tf.random.normal([BATCH_SIZE, noise_dim])
        with tf.GradientTape() as gen_tape, tf.GradientTape() as disc_tape:
            generated_images = self.generator(noise, training=False)
            real_output = self.discriminator(images, training=True)
            fake_output = self.discriminator(generated_images, training=True)

            disc_loss = self.discriminator_loss(real_output, fake_output)

        gradients_of_discriminator = disc_tape.gradient(disc_loss, self.discriminator.trainable_variables)
        self.discriminator_optimizer.apply_gradients(zip(gradients_of_discriminator, self.discriminator.trainable_variables))
        return disc_loss


    def train(self, datagen, epochs):
        start = time.time()
        # statistics of loss functions
        gen_loss, disc_loss, gen_loss_stat, disc_loss_stat = [], [], [], []

        for epoch in range(epochs):
            epoch_start = time.time()
    #         datagen.unaugmented(epoch)[0] - images
    #         datagen.unaugmented(epoch)[1] - labels
    
            for batch_num in range(STEPS_PER_EPOCH):
                batch = datagen[batch_num][0]
                disc_loss_sample = self.train_step_discriminator(batch)
                gen_loss_sample = self.train_step_generator(batch)
                gen_loss.append(gen_loss_sample)
                disc_loss.append(disc_loss_sample)
    #             print('loss: disc: ', disc_loss_sample, 'gen: ', gen_loss_sample)

            g_loss = sum(gen_loss)/len(gen_loss)
            d_loss = sum(disc_loss)/len(disc_loss)
            epoch_elapsed = time.time()-epoch_start
            print (f'Epoch {epoch+1}, gen loss={g_loss},disc loss={d_loss}, {hms_string(epoch_elapsed)}')
            gen_loss_stat.append(g_loss)
            disc_loss_stat.append(d_loss)

            # Log
            # d_loss_list = list(map(lambda x: str(x.item()), d_loss))
            self.summary.summary_writer.add_value('loss', value=[epoch, str(d_loss.numpy()), str(g_loss.numpy())])
            self.summary.summary_writer.update_log()

            # save images and checkpoints
            if (epoch + 1) % config['sample_interval'] == 0:
                sample_images(self, epoch, SEED, num_examples_to_generate)
            
            if (epoch +1) % config['sample_interval'] == 0:
                plot_loss(self, [gen_loss_stat, disc_loss_stat], ["Generator", "Discriminator"], epoch)                
            
            self.checkpoint.step.assign_add(1)
            if (epoch + 1) % config['save_interval'] == 0:
                save_path = self.checkpoint_manager.save()
                print("Saved checkpoint for step {}: {}".format(int(self.checkpoint.step), save_path))
            datagen.on_epoch_end()
        
    # final images
        sample_images(self, epochs, SEED, num_examples_to_generate)
        return gen_loss_stat, disc_loss_stat

    
    def restore_and_continue_training(self):
        pass
    

    

#8 Training
gan = Vanilla_gan(config, summary)
gen_loss_stat, disc_loss_stat = gan.train(train_gen, EPOCHS)

# Restore latest checkpoint
gan.checkpoint.restore(tf.train.latest_checkpoint(gan.checkpoint_dir))

# Show images
gen_imgs = gan.generator(SEED, training = False)

samples_count = num_examples_to_generate if num_examples_to_generate < gen_imgs.shape[0] else gen_imgs.shape[0]
square_root = int(sqrt(samples_count))
fig, axs = plt.subplots(square_root, square_root)
x, y = None, None
for i in range(samples_count):
    x, y = i // square_root, i % square_root
    axs[x, y].imshow((gen_imgs[i] + 1) / 2)
    axs[x, y].axis('off')
plt.show()