<jupyter_start><jupyter_text>Utils<jupyter_code>def load_all_imgs(files, img_size):
    """Loads all facades images
    """
    a_list = []
    b_list = []
    for fn in files:
        img_arr = np.array(Image.open(fn))
        a_list.append(img_arr[:, :img_size])
        b_list.append(img_arr[:, img_size:])
    return np.array(a_list), np.array(b_list)
data_path = '/Users/stephanrasp/repositories/CGAN-towards-CBRAIN/datasets/facades/train/'<jupyter_output><empty_output><jupyter_text>Models<jupyter_code>from keras.models import Model
from keras.layers import *
import keras.backend as K
from keras.utils.vis_utils import model_to_dot
K.set_image_data_format('channels_last')
assert K.image_data_format() == 'channels_last', 'Require channels_last!'
def conv_block(x, filters, bn=True):
    """
    Convolution block: Conv2D -> Batch norm -> LeakyReLU

    Args:
        x: Input tensor [samples, channels/filters, height, width]
        filters: Number of output filters
        bn: Use batch norm

    Returns:
        x: Output tensor [samples, filters, height, width]
    """
    x = Conv2D(filters, kernel_size=3, strides=2, padding='same')(x)
    if bn:
        x = BatchNormalization(axis=-1)(x)   # Channel axis
    x = LeakyReLU(0.2)(x)
    return x
def upsample_block(x, x_down, filters, bn=True, dr=0.5):
    """
    Upsampling block for Unet

    Args:
        x:
        filters:
        bn:
        dr:

    Returns:
        x:
    """
    x = UpSampling2D(size=(2, 2))(x)
    x = Conv2D(filters, kernel_size=3, padding='same')(x)
    if bn:
        x = BatchNormalization(axis=-1)(x)
    if dr > 0:
        x = Dropout(dr)(x)
    return concatenate([x, x_down], axis=-1)   # Concatenate along channel axis
def create_Unet_generator(img_dim, n_down, n_filters_first=64, final_activation='tanh'):
    """
    Create the generator model: Unet
    For now only 3 channels to 3 channels.
    """

    # Define input
    inp = Input(shape=img_dim)
    x = inp

    # Go down
    down_list = []
    
    # Look at that
    filter_list = [n_filters_first * min(8, (2 ** i)) for i in range(n_down)]
    print(filter_list)
    for n, n_filters in enumerate(filter_list):
        #pdb.set_trace()
        x = conv_block(x, n_filters)
        down_list.append(x)
    
    # Go up
    for x_down, n_filters in zip(down_list[:-1][::-1], filter_list[:-1][::-1]):
        x = upsample_block(x, x_down, n_filters)

    # Final upsampling and activation
    x = UpSampling2D(size=(2, 2))(x)
    x = Conv2D(3, kernel_size=3, padding='same')(x)
    outp = Activation(final_activation)(x)

    # Create and return model
    return Model(inputs=inp, outputs=outp)
def create_PixelGAN_discriminator(img_dim, n_filters_first=64, n_layers=3):
    """
    Create the discriminator model: PixelGAN

    Args:
        img_dim: Tuple of image dimensions [channels, height, width]
        n_filters_first: Number of convolution filters in the first layer
        n_layers: number of convolution layers including first

    Returns:
        model: Keras model
    """

    # Get input tensor
    inp = Input(shape=img_dim)

    # First convolution layer
    x = conv_block(inp, n_filters_first, bn=False)   # False in pytorch implentation

    # Loop over next layers
    n_filters = n_filters_first
    for n in range(1, n_layers):
        x = conv_block(x, n_filters)
        n_filters = min(2**n, 8) * n_filters_first

    # Final layer to scalar
    outp = Dense(1, activation='sigmoid')(Flatten()(x))

    # Create and return model
    return Model(inputs=inp, outputs=outp)<jupyter_output><empty_output><jupyter_text>Test<jupyter_code>G = create_Unet_generator((256, 256, 3), 3)
G.summary()
SVG(model_to_dot(G).create(prog='dot', format='svg'))
D = create_PixelGAN_discriminator((256, 256, 3))
D.summary()<jupyter_output>_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
input_2 (InputLayer)         (None, 256, 256, 3)       0         
_________________________________________________________________
conv2d_7 (Conv2D)            (None, 128, 128, 64)      1792      
_________________________________________________________________
leaky_re_lu_4 (LeakyReLU)    (None, 128, 128, 64)      0         
_________________________________________________________________
conv2d_8 (Conv2D)            (None, 64, 64, 64)        36928     
_________________________________________________________________
batch_normalization_6 (Batch (None, 64, 64, 64)        256       
_________________________________________________________________
leaky_re_lu_5 (LeakyReLU)    (None, 64, 64, 64)        0         
_________________________________________________________________
conv2d_9 ([...]<jupyter_text>Data generator<jupyter_code>import threading
from glob import glob


# To make generators thread safe for multithreading
def threadsafe_generator(f):
    """A decorator that takes a generator function and makes it thread-safe.
    """
    def g(*a, **kw):
        return ThreadsafeIter(f(*a, **kw))
    return g


class ThreadsafeIter(object):
    """Takes an iterator/generator and makes it thread-safe by
    serializing call to the `next` method of given iterator/generator.
    https://github.com/fchollet/keras/issues/1638
    """
    def __init__(self, it):
        self.it = it
        self.lock = threading.Lock()

    def __iter__(self):
        return self

    def __next__(self):   # Py3
        with self.lock:
            return next(self.it)


def generator(a_imgs, b_imgs, bs, n_batches, shuffle=False):
    """First basic generator to work with facades dataset.
    """
    while True:
        for i in range(n_batches):
            x = a_imgs[i*bs:(i+1)*bs]
            y = b_imgs[i*bs:(i+1)*bs]
            yield x, y


class DataGenerator(object):
    def __init__(self, data_dir, img_size=256, bs=4, img_arrs=None):
        self.files = sorted(glob(data_dir + '*'))
        if img_arrs is None:
            self.a_imgs, self.b_imgs = load_all_imgs(self.files, img_size)
        else:
            self.a_imgs, self.b_imgs = img_arrs
        self.bs = bs
        self.n_samples = self.a_imgs.shape[0]
        self.n_batches = int(np.floor(self.n_samples / bs))
        # Normalize
        self.a_imgs = normalize(self.a_imgs); self.b_imgs = normalize(self.b_imgs)
        self.gen = generator(self.a_imgs, self.b_imgs, bs, self.n_batches)
def normalize(x):
    return (x.astype('float32') - 127.5) / 127.5
def denormalize(x):
    return x * 127.5 + 127.5<jupyter_output><empty_output><jupyter_text>GAN<jupyter_code>def l1_loss(y_true, y_pred):
    return K.sum(K.abs(y_pred - y_true), axis=-1)
from keras.optimizers import Adam
class GAN(object):
    """
    GAN class.
    """
    def __init__(self, verbose=1):
        """
        Initialize GAN object. Define attributes.

        Args:
            verbose: Verbosity level
        """
        # Initialize empty networks
        self.G = None    # Generator G(x)
        self.D = None    # Discriminator D(x)
        self.GD = None   # Combined D(G(x))

        # Initialize the data generators
        # Generators should be compatible to CBRAIN data generator
        self.train_generator=None
        self.valid_generator=None

        # Define basic properties of the network
        self.image_dim = (256, 256, 3)
        
        # Training information
        # Useful if training is not done in one go
        self.epoch_counter = 0
        self.train_history = OrderedDict({
            'train_discriminator_loss': [],
            'train_generator_pixel_loss': [],
            'train_generator_gan_loss': [],
            'valid_discriminator_loss': [],
            'valid_generator_pixel_loss': [],
            'valid_generator_gan_loss': [],
        })
        self.verbose = verbose

    def create_generator(self, **kwargs):
        """
        Create the generator network.
        """
        self.G = create_Unet_generator(**kwargs)

    def create_discriminator(self, **kwargs):
        """
        Create the discriminator network
        """
        self.D = create_PixelGAN_discriminator(**kwargs)

    def compile(self, loss_weights=[10, 1]):
        """
        Compile the networks and create the combined network.
        """
        opt = Adam(1e-3)
        # Compile the individual models
        self.G.compile(optimizer=opt, loss='mse')  # Loss does not matter for G
        self.D.compile(optimizer=opt, loss='binary_crossentropy')

        # Create and compile the combined model
        self.D.trainable = False
        G_inp = Input(shape=(256, 256, 3,))
        G_outp = self.G(G_inp)
        D_outp = self.D(self.G(G_inp))
        self.GD = Model(inputs=G_inp, outputs=[G_outp, D_outp])
        self.GD.compile(optimizer=opt, loss=[l1_loss, 'binary_crossentropy'])

    def load_data_generator(self, data_dir, bs, **kwargs):
        """
        Load the training and validation data generators for the requested
        dataset.

        Args:
            dataset: Name of dataset
            bs: Batch size
        """
        self.train_generator = DataGenerator(data_dir + '/train/')
        self.valid_generator = DataGenerator(data_dir + '/val/')

    def train(self, epochs):
        """
        Training operation. Note that batch size is defined in the data
        generator.

        Args:
            epochs: Number of epochs to train
        """

        if self.verbose > 0: pbar = tqdm(total=epochs * self.train_generator.n_batches)
        for e in range(self.epoch_counter, self.epoch_counter + epochs):
            dl, gl = [], []
            for b in range(self.train_generator.n_batches):
                if self.verbose > 0: pbar.update(1)
                dl, gl = self.train_step(self.train_generator.bs, dl, gl, b)
            self.epoch_counter += 1

            # End of epoch. Compute mean generator and discriminator loss
            self.train_history['train_discriminator_loss'].append(np.mean(dl))
            self.train_history['train_generator_pixel_loss'].append(np.mean(gl, 0)[0])
            self.train_history['train_generator_gan_loss'].append(np.mean(gl, 0)[1])
            fake = self.evaluate_test_losses()

            # Save images
            if e % 1 == 0:
                self.save_images(fake)

            # Update progressbar with latest losses
            if self.verbose > 0:
                pbar_dict = OrderedDict({k: v[-1] for k, v
                                         in self.train_history.items()})
                pbar.set_postfix(pbar_dict)
        if self.verbose > 0: pbar.close()

    def train_step(self, bs, dl, gl, ibatch):
        """One training step. May contain several discriminator steps."""

        # STEP 1: TRAIN DISCRIMINATOR
        self.D.trainable = True
        # Get images
        X, Y_true = next(self.train_generator.gen)
        if ibatch % 2 == 0:
            # Train alternatively on real and fake images
            # Create fake images
            Y = self.G.predict_on_batch(X)
            y = np.array(self.label('fake') * bs)
        else:
            Y = Y_true
            y = np.array(self.label('real') * bs)
        dl.append(self.D.train_on_batch(Y, y))

        # STEP 2: TRAIN GENERATOR
        self.D.trainable = False
        X, Y_true = next(self.train_generator.gen)
        gl.append(self.GD.train_on_batch(Y, [Y_true, np.array(self.label('real') * bs)]))
        return dl, gl

    def evaluate_test_losses(self):
        """Compute losses for test set and returns some fake images"""
        bs=4
        X, Y_true = (self.valid_generator.a_imgs, self.valid_generator.b_imgs)
        Y_fake = self.G.predict(X, batch_size=bs)
        Y_concat = np.concatenate([Y_true, Y_fake])
        y_concat = np.array(
            self.label('real') * self.valid_generator.n_samples +
            self.label('fake') * self.valid_generator.n_samples
        )
        self.train_history['valid_discriminator_loss'].append(
            self.D.evaluate(Y_concat, y_concat, batch_size=bs, verbose=0)
        )
        GD_losses = self.GD.evaluate(X, [Y_true, np.array(self.label('real') * 
                                         self.valid_generator.n_samples)])
        self.train_history['valid_generator_pixel_loss'].append(GD_losses[0])
        self.train_history['valid_generator_gan_loss'].append(GD_losses[1])
        return Y_fake

    def label(self, s):
        """Little helper function to return labels"""
        assert s in ['real', 'fake'], 'Wrong string for label function.'
        return [1] if s == 'real' else [0]

    def save_images(self, fake):
        """Saves some fake images"""
        s = ('/export/home/srasp/repositories/CGAN-towards-CBRAIN/images/' + 
               'plot_epoch_{0:04d}_generated'.format(self.epoch_counter))
        Image.fromarray(denormalize(fake[0]).astype(np.uint8)).save(s + '.png')

    def save_models(self):
        """
        Saves models and training history
        """
        s = self.model_dir + self.exp_id + '_'
        self.G.save(s + 'G.h5')
        self.D.save(s + 'D.h5')
        self.GD.save(s + 'GD.h5')
        # Save training history
        with open(s + 'history.pkl', 'wb') as f:
            pickle.dump(self.train_history, f)<jupyter_output><empty_output><jupyter_text>Try it out<jupyter_code>gan = GAN()
data_dir = '/export/home/srasp/repositories/CGAN-towards-CBRAIN/datasets/cityscapes/'
gan.train_generator = DataGenerator(data_dir + 'train/', img_arrs = train_arrs)
gan.valid_generator = DataGenerator(data_dir + 'val/', img_arrs = valid_arrs)
#tmp = DataGenerator(data_dir + '/val/')
gan.load_data_generator('/export/home/srasp/repositories/CGAN-towards-CBRAIN/datasets/cityscapes/', 
                        bs = 4)
gan.create_generator(img_dim=(256, 256, 3), n_down=3)
gan.create_discriminator(img_dim=(256, 256, 3))
gan.compile()
gan.train(epochs=1)
img_fn = '/export/home/srasp/repositories/CGAN-towards-CBRAIN/images/plot_epoch_0001_generated.png'
img = Image.open(img_fn); img
gan.train(epochs=5)
train_arrs = (gan.train_generator.a_imgs, gan.train_generator.b_imgs)
valid_arrs = (tmp.a_imgs, tmp.b_imgs)
gan.train_generator.a_imgs.shape
train_arrs_sample = (gan.train_generator.a_imgs[:100], gan.train_generator.b_imgs[:100])
valid_arrs_sample = (tmp.a_imgs[:20], tmp.b_imgs[:20])<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.11983929066223331, 'fraction_numerical': 0.01662510390689942, 'mean_word_length': 2.6716683621566633, 'pattern_counts': {'":': 0, '<': 24, '<?xml version=': 0, '>': 31, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '6', 'dir': 'jupyter-structured-clean-dedup', 'id': '363936', 'license': 'MIT', 'n_tokens_mistral': 4874, 'n_tokens_neox': 4602, 'n_words': 1022, 'path': 'notebooks/gan_development.ipynb', 'repo_name': 'raspstephan/CGAN-towards-CBRAIN', 'size': '14436'}