{"cells":[{"metadata":{},"cell_type":"markdown","source":"# Contrastive Unpaired Translation\nThis Notebook implements the contrastive unpaired translation (CUT) of photos into Monet-style paintings.\nMore information about CUT can be found [here](https://taesung.me/ContrastiveUnpairedTranslation/) (videos, GitHub, link to paper, etc.).\nThe code is this notebook is taken from [this GitHub repository](https://github.com/cryu854/CUT), which contains a TensorFlow 2 implementation of CUT.\nThe last two cells to store the images is based on https://www.kaggle.com/tpothjuan/cyclegan-with-gpu-with-explanations (Maybe remove this last sentence if the 2nd link has exactly the same way of generating and storing the final images).\n\n**Notes:**\n* Run this notebook on a GPU.\n* Current speeds: (TBA)\n\n**Plans:**\n* Early stopping (but this is maybe not logical because the losses do not seem to stabalise between epochs)\n* Adapting the number of resnet layers\n* Choosing the best learning rates\n* Data-augmentation: more images, maybe zooming in.\n\n**Implemented:**\n* Data-augmentation: horizontally flipped images, removed two circular images.\n* Anti-aliasing: part of the code from the TensorFlow 2 implementation that this notebook is based on. Increased the score slightly (by 1 point), but also increased the running time by 30%. Did not remove the checkerboard pattern, although that is why we implemented it.\n* Validation data: assessing the performance of the model on unseen data shows how well it generalises."},{"metadata":{"trusted":true},"cell_type":"code","source":"# Maximum line-length (79) including the sharp and space for PEP-8:\n# 34567890123456789012345678901234567890123456789012345678901234567890123456789\n# Remember that the maximum line length for comments is 72.\n\n#### Customisation parameters #### Set these to what you prefer!\n\n# True for 'monet-jpg-improved', False for 'gan-getting-started':\naugmented_input = True\n# True to enable anti-alias, False to disable anti-alias:\nanti_alias = True\n# Max number of epochs:\nmax_epochs = 15\n# Threshold of run-time in seconds above which no new epochs start:\nmax_train_time = 14400\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"#### Imports ####\n# Kaggle:\nfrom kaggle_datasets import KaggleDatasets\n\n# Basic:\nimport time\nimport matplotlib.pyplot as plt\nimport numpy as np\n\n# Tensorflow:\nimport tensorflow as tf\nfrom tensorflow.keras.models import Model\nfrom tensorflow.keras.layers import (Input, Dense, Lambda, Layer, Conv2D,\n    Conv2DTranspose, BatchNormalization, Activation)\n# I think these can be removed:\n# from tensorflow import keras\n# from tensorflow.keras import layers\n\n# Outputting the generated images:\nimport os\nimport shutil\nfrom PIL import Image\n\n# Constants:\nAUTOTUNE = tf.data.experimental.AUTOTUNE\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Loading the data\n\nWe use the JPGs for training. This has to be used in combination with the GPU, unless the code is changed a lot further.\n"},{"metadata":{"trusted":true},"cell_type":"code","source":"@tf.function\ndef load_image(image_file, image_size=None, data_augmentation=False):\n    \"\"\" Load the image file.\"\"\"\n    image = tf.io.read_file(image_file)\n    image = tf.image.decode_png(image)\n    image = (tf.cast(image, tf.float32) / 127.5) - 1.0\n\n    if data_augmentation:\n        image = tf.image.random_flip_left_right(image)\n    if image_size is not None:\n        image = tf.image.resize(image, size=(image_size[0], image_size[1]))\n    if tf.shape(image)[-1] == 1:\n        image = tf.tile(image, [1,1,3])\n\n    return  image\n\ndef create_dataset(src_folder, tar_folder, batch_size):\n    \"\"\" Create tf.data.Dataset.\n    Input:\n    - src_folder, tar_folder: (String) folder path of training data and\n      their targets\n    - batch_size: the batch size for the generated dataset\n    \n    Output:\n    - train_dataset: (tf.data.Dataset) zipped dataset of 80 percent of\n      the contents of src_folder and tar_folder\n    - val_dataset: (tf.data.Dataset) zipped dataset of 20 percent of the\n      contents of src_ and tar_folder, as validation data\n    \"\"\"\n    # Create the shuffled datasets for the source and target:\n    src_dataset = tf.data.Dataset.list_files(src_folder+'/*.jpg', shuffle=False)\n    tar_dataset = tf.data.Dataset.list_files(tar_folder+'/*.jpg', shuffle=False)\n    \n    # Calculate size of training set for src and tar:\n    src_dataset_size = tf.data.Dataset.cardinality(src_dataset).numpy()\n    tar_dataset_size = tf.data.Dataset.cardinality(tar_dataset).numpy()\n    \n    train_src_size = int(0.8 * src_dataset_size)\n    train_tar_size = int(0.8 * tar_dataset_size)\n\n    # Shuffle with reshuffle=False, so order stays the same each time:\n    src_dataset = src_dataset.shuffle(src_dataset_size,\n                                      reshuffle_each_iteration=False)\n    tar_dataset = tar_dataset.shuffle(tar_dataset_size,\n                                      reshuffle_each_iteration=False)\n    \n    # Separate the training set and validation set:\n    train_src_dataset = src_dataset.take(train_src_size)\n    train_tar_dataset = tar_dataset.take(train_tar_size)\n    train_src_dataset = train_src_dataset.shuffle(train_src_size,\n        reshuffle_each_iteration=True)\n    train_tar_dataset = train_tar_dataset.shuffle(train_tar_size,\n        reshuffle_each_iteration=True)\n    \n    val_src_dataset = src_dataset.skip(train_src_size)\n    val_tar_dataset = tar_dataset.skip(train_tar_size)\n    \n    train_src_dataset = (\n        train_src_dataset.map(load_image, num_parallel_calls=AUTOTUNE)\n        .batch(batch_size, drop_remainder=True)\n        .prefetch(AUTOTUNE)\n    )\n    \n    train_tar_dataset = (\n        train_tar_dataset.map(load_image, num_parallel_calls=AUTOTUNE)\n        .batch(batch_size, drop_remainder=True)\n        .prefetch(AUTOTUNE)\n    )\n    \n    val_src_dataset = (\n        val_src_dataset.map(load_image, num_parallel_calls=AUTOTUNE)\n        .batch(batch_size, drop_remainder=True)\n        .prefetch(AUTOTUNE)\n    )\n    \n    val_tar_dataset = (\n        val_tar_dataset.map(load_image, num_parallel_calls=AUTOTUNE)\n        .batch(batch_size, drop_remainder=True)\n        .prefetch(AUTOTUNE)\n    )\n    \n    train_dataset = tf.data.Dataset.zip((train_src_dataset, train_tar_dataset))\n    val_dataset = tf.data.Dataset.zip((val_src_dataset, val_tar_dataset))\n\n    return train_dataset, val_dataset\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Set input folder paths:\nGCS_PATH_KAGGLE = KaggleDatasets().get_gcs_path('gan-getting-started')\ntrain_src_folder = GCS_PATH_KAGGLE + '/photo_jpg'\n\nif augmented_input:\n    GCS_PATH_AUGMENTED = KaggleDatasets().get_gcs_path('monet-jpg-improved')\n    train_tar_folder = GCS_PATH_AUGMENTED + '/monet_jpg_improved'\nelse:\n    train_tar_folder = GCS_PATH_KAGGLE + '/monet_jpg'\n\n# Create the datasets:\nbatch_size = 1\ntrain_dataset, val_dataset = create_dataset(train_src_folder,\n                                            train_tar_folder,\n                                            batch_size)\n\n# Get the image shapes (ignore batch size):\nsource_image, target_image = next(iter(train_dataset))\nsource_shape = source_image.shape[1:]\ntarget_shape = target_image.shape[1:]\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# CUT Model\n**Loss Functions**"},{"metadata":{"trusted":true},"cell_type":"code","source":"class GANLoss:\n    def __init__(self, gan_mode):\n        self.gan_mode = gan_mode\n        if gan_mode == 'lsgan':\n            self.loss = tf.keras.losses.MeanSquaredError()\n        else:\n            raise NotImplementedError(f'gan mode {gan_mode} not implemented.')\n\n    def __call__(self, prediction, target_is_real):\n        if self.gan_mode == 'lsgan':\n            if target_is_real:\n                loss = self.loss(tf.ones_like(prediction), prediction)\n            else:\n                loss = self.loss(tf.zeros_like(prediction), prediction)\n        return loss\n\n\nclass PatchNCELoss:\n    def __init__(self, nce_temp, nce_lambda):\n        # Potential: only supports for batch_size=1 now.\n        self.nce_temp = nce_temp\n        self.nce_lambda = nce_lambda\n        self.cross_entropy_loss = tf.keras.losses.CategoricalCrossentropy(\n            reduction=tf.keras.losses.Reduction.NONE, from_logits=True)\n\n    def __call__(self, source, target, netE, netF):\n        feat_source = netE(source, training=True)\n        feat_target = netE(target, training=True)\n\n        feat_source_pool, sample_ids = netF(feat_source,\n                                            patch_ids=None,\n                                            training=True)\n        feat_target_pool, _ = netF(feat_target,\n                                   patch_ids=sample_ids,\n                                   training=True)\n        \n        total_nce_loss = 0.0\n        for feat_s, feat_t in zip(feat_source_pool, feat_target_pool):\n            n_patches, dim = feat_s.shape\n\n            logit = tf.matmul(feat_s, tf.transpose(feat_t)) / self.nce_temp\n\n            # Diagonal entries are pos logits, the others are neg logits\n            diagonal = tf.eye(n_patches, dtype=tf.bool)\n            target = tf.where(diagonal, 1.0, 0.0)\n\n            loss = self.cross_entropy_loss(target, logit) * self.nce_lambda\n            total_nce_loss += tf.reduce_mean(loss)\n\n        return total_nce_loss / len(feat_source_pool)\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"**Layers**"},{"metadata":{},"cell_type":"markdown","source":"Up- and Downsampling "},{"metadata":{"trusted":true},"cell_type":"code","source":"def _setup_kernel(k):\n    k = np.asarray(k, dtype=np.float32)\n    if k.ndim == 1:\n        k = np.outer(k, k)\n    k /= np.sum(k)\n    assert k.ndim == 2\n    assert k.shape[0] == k.shape[1]\n    return k\n\ndef _shape(tf_expr, dim_idx):\n    if tf_expr.shape.rank is not None:\n        dim = tf_expr.shape[dim_idx]\n        if dim is not None:\n            return dim\n    return tf.shape(tf_expr)[dim_idx]\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def _upfirdn_2d_ref(x, k, upx, upy, downx, downy, padx0, padx1, pady0, pady1):\n    \"\"\"Slow reference implementation of `upfirdn_2d()` using standard\n    TensorFlow ops.\n    \"\"\"\n\n    x = tf.convert_to_tensor(x)\n    k = np.asarray(k, dtype=np.float32)\n    assert x.shape.rank == 4\n    inH = x.shape[1]\n    inW = x.shape[2]\n    minorDim = _shape(x, 3)\n    kernelH, kernelW = k.shape\n    assert inW >= 1 and inH >= 1\n    assert kernelW >= 1 and kernelH >= 1\n    assert isinstance(upx, int) and isinstance(upy, int)\n    assert isinstance(downx, int) and isinstance(downy, int)\n    assert isinstance(padx0, int) and isinstance(padx1, int)\n    assert isinstance(pady0, int) and isinstance(pady1, int)\n\n    # Upsample (insert zeros).\n    x = tf.reshape(x, [-1, inH, 1, inW, 1, minorDim])\n    x = tf.pad(x, [[0, 0], [0, 0], [0, upy - 1], [0, 0], [0, upx - 1], [0, 0]])\n    x = tf.reshape(x, [-1, inH * upy, inW * upx, minorDim])\n\n    # Pad (crop if negative).\n    x = tf.pad(x, [[0, 0],\n                   [max(pady0, 0), max(pady1, 0)],\n                   [max(padx0, 0), max(padx1, 0)],\n                   [0, 0]]\n              )\n    x = x[:,\n          max(-pady0, 0) : x.shape[1] - max(-pady1, 0),\n          max(-padx0, 0) : x.shape[2] - max(-padx1, 0),\n          :\n         ]\n\n    # Convolve with filter.\n    x = tf.transpose(x, [0, 3, 1, 2])\n    x = tf.reshape(x, [-1,\n                       1,\n                       inH * upy + pady0 + pady1,\n                       inW * upx + padx0 + padx1\n                      ]\n                  )\n    w = tf.constant(k[::-1, ::-1, np.newaxis, np.newaxis], dtype=x.dtype)\n    x = tf.nn.conv2d(x, w, strides=[1,1,1,1], padding='VALID',\n                     data_format='NCHW')\n    x = tf.reshape(x, [-1,\n                       minorDim,\n                       inH * upy + pady0 + pady1 - kernelH + 1,\n                       inW * upx + padx0 + padx1 - kernelW + 1\n                      ]\n                  )\n    x = tf.transpose(x, [0, 2, 3, 1])\n\n    # Downsample (throw away pixels).\n    return x[:, ::downy, ::downx, :]\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def _simple_upfirdn_2d(x, k, up=1, down=1, pad0=0, pad1=0, data_format='NCHW',\n                       impl='ref'):\n    assert data_format in ['NCHW', 'NHWC']\n    assert x.shape.rank == 4\n    y = x\n    if data_format == 'NCHW':\n        y = tf.reshape(y, [-1, _shape(y, 2), _shape(y, 3), 1])\n        #######MODIFIED#########\n    if impl== \"ref\":\n        y = _upfirdn_2d_ref(y, k, upx=up, upy=up, downx=down, downy=down,\n                            padx0=pad0, padx1=pad1, pady0=pad0, pady1=pad1)\n    else: \n         raise NotImplementedError(f'implentation {impl} not implemented.')\n    if data_format == 'NCHW':\n        y = tf.reshape(y, [-1, _shape(x, 1), _shape(y, 1), _shape(y, 2)])\n    return y\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def upsample_2d(x, k=None, factor=2, gain=1, data_format='NCHW', impl='cuda'):\n    r\"\"\"Upsample a batch of 2D images with the given filter.\n    Accepts a batch of 2D images of the shape `[N, C, H, W]` or\n    `[N, H, W, C]` and upsamples each image with the given filter.\n    The filter is normalized so that if the input pixels are constant,\n    they will be scaled by the specified `gain`.\n    Pixels outside the image are assumed to be zero, and the filter is\n    padded with zeros so that its shape is a multiple of the upsampling\n    factor.\n    Args:\n        x: Input tensor of the shape `[N, C, H, W]` or`[N, H, W, C]`.\n        k: FIR filter of the shape `[firH, firW]` or `[firN]`\n           (separable). The default is `[1] * factor`, which\n           corresponds to nearest-neighbor upsampling.\n        factor: Integer upsampling factor (default: 2).\n        gain: Scaling factor for signal magnitude (default: 1.0).\n        data_format: `'NCHW'` or `'NHWC'` (default: `'NCHW'`).\n        impl: Name of the implementation to use. Can ONLY be `\"ref\"`.\n    Returns:\n        Tensor of the shape `[N, C, H * factor, W * factor]` or\n        `[N, H * factor, W * factor, C]`, and same datatype as `x`.\n    \"\"\"\n\n    assert isinstance(factor, int) and factor >= 1\n    if k is None:\n        k = [1] * factor\n    k = _setup_kernel(k) * (gain * (factor ** 2))\n    p = k.shape[0] - factor\n    return _simple_upfirdn_2d(x, k, up=factor, pad0=(p+1)//2+factor-1,\n                              pad1=p//2, data_format=data_format, impl=impl)\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def downsample_2d(x, k=None, factor=2, gain=1, data_format='NCHW', impl='cuda'):\n    r\"\"\"Downsample a batch of 2D images with the given filter.\n    Accepts a batch of 2D images of the shape `[N, C, H, W]` or\n    `[N, H, W, C]` and downsamples each image with the given filter.\n    The filter is normalized so that if the input pixels are constant,\n    they will be scaled by the specified `gain`.\n    Pixels outside the image are assumed to be zero, and the filter\n    is padded with zeros so that its shape is a multiple of the\n    downsampling factor.\n    Args:\n        x: Input tensor of the shape `[N, C, H, W]` or `[N, H, W, C]`.\n        k: FIR filter of the shape `[firH, firW]` or `[firN]`\n           (separable). The default is `[1] * factor`, which corresponds\n           to average pooling.\n        factor: Integer downsampling factor (default: 2).\n        gain: Scaling factor for signal magnitude (default: 1.0).\n        data_format: `'NCHW'` or `'NHWC'` (default: `'NCHW'`).\n        impl: Name of the implementation to use. Can ONLY be `\"ref\"`.\n    Returns:\n        Tensor of the shape `[N, C, H // factor, W // factor]` or\n        `[N, H // factor, W // factor, C]`, and same datatype as `x`.\n    \"\"\"\n\n    assert isinstance(factor, int) and factor >= 1\n    if k is None:\n        k = [1] * factor\n    k = _setup_kernel(k) * gain\n    p = k.shape[0] - factor\n    return _simple_upfirdn_2d(x, k, down=factor, pad0=(p+1)//2, pad1=p//2,\n                              data_format=data_format, impl=impl)\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Other Layers"},{"metadata":{"trusted":true},"cell_type":"code","source":"class Padding2D(Layer):\n    \"\"\" 2D padding layer.\"\"\"\n    def __init__(self, padding=(1, 1), pad_type='constant', **kwargs):\n        assert pad_type in ['constant', 'reflect', 'symmetric']\n        super(Padding2D, self).__init__(**kwargs)\n        if type(padding) is int:\n            self.padding = (padding, padding)\n        else:\n            self.padding = tuple(padding)\n        self.pad_type = pad_type\n\n    def call(self, inputs, training=None):\n        padding_width, padding_height = self.padding\n        padding_tensor = [[0, 0],\n                          [padding_height, padding_height],\n                          [padding_width, padding_width],\n                          [0, 0],\n                         ]\n\n        return tf.pad(inputs, padding_tensor, mode=self.pad_type)\n\n\nclass InstanceNorm(Layer):\n    \"\"\" Instance Normalization layer\n    (https://arxiv.org/abs/1607.08022).\n    \"\"\"\n    def __init__(self, epsilon=1e-5, affine=False, **kwargs):\n        super(InstanceNorm, self).__init__(**kwargs)\n        self.epsilon = epsilon\n        self.affine = affine\n \n    def build(self, input_shape):\n        if self.affine:\n            self.gamma = self.add_weight(name='gamma',\n                shape=(input_shape[-1],),\n                initializer=tf.random_normal_initializer(0, 0.02),\n                trainable=True)\n            self.beta = self.add_weight(name='beta',\n                                        shape=(input_shape[-1],),\n                                        initializer=tf.zeros_initializer(),\n                                        trainable=True)\n\n    def call(self, inputs, training=None):\n        mean, var = tf.nn.moments(inputs, axes=[1, 2], keepdims=True)\n        x = tf.divide(tf.subtract(inputs, mean),\n                      tf.math.sqrt(tf.add(var, self.epsilon)))\n\n        if self.affine:\n            return self.gamma * x + self.beta\n        return x\n\n    \nclass AntialiasSampling(Layer):\n    \"\"\" Down/Up sampling layer with blur-kernel.\n    \"\"\"\n    def __init__(self,\n                 kernel_size,\n                 mode,\n                 impl, \n                 **kwargs):\n        super(AntialiasSampling, self).__init__(**kwargs)\n        if(kernel_size == 1):\n            self.k = np.array([1., ])\n        elif(kernel_size == 2):\n            self.k = np.array([1., 1.])\n        elif(kernel_size == 3):\n            self.k = np.array([1., 2., 1.])\n        elif(kernel_size == 4):\n            self.k = np.array([1., 3., 3., 1.])\n        elif(kernel_size == 5):\n            self.k = np.array([1., 4., 6., 4., 1.])\n        elif(kernel_size == 6):\n            self.k = np.array([1., 5., 10., 10., 5., 1.])\n        elif(kernel_size == 7):\n            self.k = np.array([1., 6., 15., 20., 15., 6., 1.])\n        self.mode = mode\n        self.impl = impl\n\n    def call(self, inputs, training=None):\n        if self.mode == 'up':\n            x = upsample_2d(inputs, k=self.k, data_format='NHWC',\n                            impl=self.impl)\n        elif self.mode == 'down':\n            x = downsample_2d(inputs, k=self.k, data_format='NHWC',\n                              impl=self.impl)\n        else:\n            raise ValueError(f'Unsupported sampling mode: {self.mode}')\n\n        return x\n    \nclass ConvBlock(Layer):\n    \"\"\"ConBlock layer consists of Conv2D + Normalization +\n    Activation.\n    \"\"\"\n    def __init__(self, filters, kernel_size, strides=(1,1), padding='valid',\n                 use_bias=True, norm_layer=None, activation='linear',\n                 **kwargs):\n        super(ConvBlock, self).__init__(**kwargs)\n        initializer = tf.random_normal_initializer(0., 0.02)\n        self.conv2d = Conv2D(filters,\n                             kernel_size,\n                             strides,\n                             padding,\n                             use_bias=use_bias,\n                             kernel_initializer=initializer)\n        self.activation = Activation(activation)\n        if norm_layer == 'batch':\n            self.normalization = BatchNormalization()\n        elif norm_layer == 'instance':\n            self.normalization = InstanceNorm(affine=False)\n        else:\n            self.normalization = tf.identity\n\n    def call(self, inputs, training=None):\n        x = self.conv2d(inputs)\n        x = self.normalization(x)\n        x = self.activation(x)\n        return x\n\n\nclass ConvTransposeBlock(Layer):\n    \"\"\" ConvTransposeBlock layer consists of Conv2DTranspose + \n    Normalization + Activation.\n    \"\"\"\n    def __init__(self, filters, kernel_size, strides=(1,1), padding='valid',\n                 use_bias=True, norm_layer=None, activation='linear',\n                 **kwargs):\n        super(ConvTransposeBlock, self).__init__(**kwargs)\n        initializer = tf.random_normal_initializer(0., 0.02)\n        self.convT2d = Conv2DTranspose(filters,\n                                       kernel_size,\n                                       strides,\n                                       padding,\n                                       use_bias=use_bias,\n                                       kernel_initializer=initializer)\n        self.activation = Activation(activation)\n        if norm_layer == 'batch':\n            self.normalization = BatchNormalization()\n        elif norm_layer == 'instance':\n            self.normalization = InstanceNorm(affine=False)\n        else:\n            self.normalization = tf.identity\n\n    def call(self, inputs, training=None):\n        x = self.convT2d(inputs)\n        x = self.normalization(x)\n        x = self.activation(x)\n        return x\n\n\nclass ResBlock(Layer):\n    \"\"\" ResBlock is a ConvBlock with skip connections.\n    Original Resnet paper (https://arxiv.org/pdf/1512.03385.pdf).\n    \"\"\"\n    def __init__(self, filters, kernel_size, use_bias, norm_layer, **kwargs):\n        super(ResBlock, self).__init__(**kwargs)\n        self.reflect_pad1 = Padding2D(1, pad_type='reflect')\n        self.conv_block1 = ConvBlock(filters,\n                                     kernel_size,\n                                     padding='valid',\n                                     use_bias=use_bias,\n                                     norm_layer=norm_layer,\n                                     activation='relu')\n        self.reflect_pad2 = Padding2D(1, pad_type='reflect')\n        self.conv_block2 = ConvBlock(filters,\n                                     kernel_size,\n                                     padding='valid',\n                                     use_bias=use_bias,\n                                     norm_layer=norm_layer)\n\n    def call(self, inputs, training=None):\n        x = self.reflect_pad1(inputs)\n        x = self.conv_block1(x)\n        x = self.reflect_pad2(x)\n        x = self.conv_block2(x)\n        return inputs + x\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"**The Model Itself**\n"},{"metadata":{"trusted":true},"cell_type":"code","source":"def Generator(input_shape, output_shape, norm_layer, use_antialias,\n              resnet_blocks, impl):\n    \"\"\" Create a Resnet-based generator.\n    Adapt from Justin Johnson's neural style transfer project\n    (https://github.com/jcjohnson/fast-neural-style).\n    For BatchNorm, we use learnable affine parameters and track running\n    statistics (mean/stddev). For InstanceNorm, we do not use learnable\n    affine parameters. We do not track running statistics. \n    \"\"\"\n    use_bias = (norm_layer == 'instance')\n    inputs = Input(shape=input_shape)\n    \n    x = Padding2D(3, pad_type='reflect')(inputs)\n    x = ConvBlock(64, 7, padding='valid', use_bias=use_bias,\n                  norm_layer=norm_layer, activation='relu')(x)\n   \n    if use_antialias:\n        x = ConvBlock(128, 3, padding='same', use_bias=use_bias,\n                      norm_layer=norm_layer, activation='relu')(x)\n        x = AntialiasSampling(4, mode='down', impl=impl)(x)\n        x = ConvBlock(256, 3, padding='same', use_bias=use_bias,\n                      norm_layer=norm_layer, activation='relu')(x)\n        x = AntialiasSampling(4, mode='down', impl=impl)(x)\n    else:\n        x = ConvBlock(128, 3, strides=2, padding='same', use_bias=use_bias,\n                      norm_layer=norm_layer, activation='relu')(x)\n        x = ConvBlock(256, 3, strides=2, padding='same', use_bias=use_bias,\n                      norm_layer=norm_layer, activation='relu')(x)\n\n    for _ in range(resnet_blocks):\n        x = ResBlock(256, 3, use_bias, norm_layer)(x)\n      \n    if use_antialias:\n        x = AntialiasSampling(4, mode='up', impl=impl)(x)\n        x = ConvBlock(128, 3, padding='same', use_bias=use_bias,\n                      norm_layer=norm_layer, activation='relu')(x)\n        x = AntialiasSampling(4, mode='up', impl=impl)(x)\n        x = ConvBlock(64, 3, padding='same', use_bias=use_bias,\n                      norm_layer=norm_layer, activation='relu')(x)\n    else:\n        x = ConvTransposeBlock(128, 3, strides=2, padding='same',\n                               use_bias=use_bias, norm_layer=norm_layer,\n                               activation='relu')(x)\n        x = ConvTransposeBlock(64, 3, strides=2, padding='same',\n                               use_bias=use_bias, norm_layer=norm_layer,\n                               activation='relu')(x)\n\n    x = Padding2D(3, pad_type='reflect')(x)\n    outputs = ConvBlock(output_shape[-1], 7, padding='valid',\n                        activation='tanh')(x)\n\n    return Model(inputs=inputs, outputs=outputs, name='generator')\n\ndef Discriminator(input_shape, norm_layer, use_antialias, impl):\n    \"\"\" Create a PatchGAN discriminator.\n    PatchGAN classifier described in the original pix2pix paper\n    (https://arxiv.org/abs/1611.07004).\n    Such a patch-level discriminator architecture has fewer parameters\n    than a full-image discriminator and can work on arbitrarily-sized\n    images in a fully convolutional fashion.\n    \"\"\"\n    use_bias = (norm_layer == 'instance')\n    inputs = Input(shape=input_shape)\n    \n    if use_antialias:\n        x = ConvBlock(64, 4, padding='same',\n                      activation=tf.nn.leaky_relu)(inputs)\n        x = AntialiasSampling(4, mode='down', impl=impl)(x)\n        x = ConvBlock(128, 4, padding='same', use_bias=use_bias,\n                      norm_layer=norm_layer, activation=tf.nn.leaky_relu)(x)\n        x = AntialiasSampling(4, mode='down', impl=impl)(x)\n        x = ConvBlock(256, 4, padding='same', use_bias=use_bias,\n                      norm_layer=norm_layer, activation=tf.nn.leaky_relu)(x)\n        x = AntialiasSampling(4, mode='down', impl=impl)(x)\n    else:\n        x = ConvBlock(64, 4, strides=2, padding='same',\n                      activation=tf.nn.leaky_relu)(inputs)\n        x = ConvBlock(128, 4, strides=2, padding='same', use_bias=use_bias,\n                      norm_layer=norm_layer, activation=tf.nn.leaky_relu)(x)\n        x = ConvBlock(256, 4, strides=2, padding='same', use_bias=use_bias,\n                      norm_layer=norm_layer, activation=tf.nn.leaky_relu)(x)\n\n    x = Padding2D(1, pad_type='constant')(x)\n    x = ConvBlock(512, 4, padding='valid', use_bias=use_bias,\n                  norm_layer=norm_layer, activation=tf.nn.leaky_relu)(x)\n    x = Padding2D(1, pad_type='constant')(x)\n    outputs = ConvBlock(1, 4, padding='valid')(x)\n\n    return Model(inputs=inputs, outputs=outputs, name='discriminator')\n\ndef Encoder(generator, nce_layers):\n    \"\"\" Create an Encoder that shares weights with the generator.\"\"\"\n    assert max(nce_layers) <= len(generator.layers) and min(nce_layers) >= 0\n    outputs = [generator.get_layer(index=idx).output for idx in nce_layers]\n    return Model(inputs=generator.input, outputs=outputs, name='encoder')\n\n\nclass PatchSampleMLP(Model):\n    \"\"\" Create a PatchSampleMLP.\n    Adapt from official CUT implementation\n    (https://github.com/taesungp/contrastive-unpaired-translation).\n    PatchSampler samples patches from pixel/feature-space. Two-layer\n    MLP projects both the input and output patches to a shared\n    embedding space.\n    \"\"\"\n    def __init__(self, units, num_patches, **kwargs):\n        super(PatchSampleMLP, self).__init__(**kwargs)\n        self.units = units\n        self.num_patches = num_patches\n        self.l2_norm = Lambda(lambda x: x * tf.math.rsqrt(tf.reduce_sum(\n            tf.square(x), axis=-1, keepdims=True) + 10-10))\n\n    def build(self, input_shape):\n        initializer = tf.random_normal_initializer(0., 0.02)\n        feats_shape = input_shape\n        for feat_id in range(len(feats_shape)):\n            mlp = tf.keras.models.Sequential([\n                    Dense(self.units, activation=\"relu\",\n                          kernel_initializer=initializer),\n                    Dense(self.units, kernel_initializer=initializer),\n                ])\n            setattr(self, f'mlp_{feat_id}', mlp)\n\n    def call(self, inputs, patch_ids=None, training=None):\n        feats = inputs\n        samples = []\n        ids = []\n        for feat_id, feat in enumerate(feats):\n            feat.set_shape([1, feat.shape[1], feat.shape[2], feat.shape[3]])\n            B, H, W, C = feat.shape\n            feat_reshape = tf.reshape(feat, [B, -1, C])\n\n            if patch_ids is not None:\n                patch_id = patch_ids[feat_id]\n            else:\n                patch_id = tf.random.shuffle(\n                    tf.range(H * W))[:min(self.num_patches, H * W)]\n\n            x_sample = tf.reshape(\n                tf.gather(feat_reshape, patch_id, axis=1), [-1, C])\n            mlp = getattr(self, f'mlp_{feat_id}')\n            x_sample = mlp(x_sample)\n            x_sample = self.l2_norm(x_sample)\n            samples.append(x_sample)\n            ids.append(patch_id)\n        return samples, ids\n\n\nclass CUT_model(Model):\n    \"\"\" Create a CUT/FastCUT model, described in the paper Contrastive\n    Learning for Unpaired Image-to-Image Translation. Taesung Park,\n    Alexei A. Efros, Richard Zhang, Jun-Yan Zhu. ECCV, 2020\n    (https://arxiv.org/abs/2007.15651).\n    \"\"\"\n    def __init__(self, source_shape, target_shape, cut_mode='cut',\n                 gan_mode='lsgan', use_antialias=True, norm_layer='instance',\n                 resnet_blocks=9, netF_units=256, netF_num_patches=256,\n                 nce_temp=0.07, nce_layers=[0,3,5,7,11], impl='ref', **kwargs):\n        assert cut_mode in ['cut']\n        assert gan_mode in ['lsgan']\n        assert norm_layer in [None, 'batch', 'instance']\n        assert netF_units > 0\n        assert netF_num_patches > 0\n        super(CUT_model, self).__init__(self, **kwargs)\n\n        self.gan_mode = gan_mode\n        self.nce_temp = nce_temp\n        self.nce_layers = nce_layers\n        self.netG = Generator(source_shape, target_shape, norm_layer,\n                              use_antialias, resnet_blocks, impl)\n        self.netD = Discriminator(target_shape, norm_layer, use_antialias,\n                                  impl)\n        self.netE = Encoder(self.netG, self.nce_layers)\n        self.netF = PatchSampleMLP(netF_units, netF_num_patches)\n\n        if cut_mode == 'cut':\n            self.nce_lambda = 1.0\n            self.use_nce_identity = True\n        else:\n            raise ValueError(cut_mode)\n\n    def compile(self, G_optimizer, F_optimizer, D_optimizer,):\n        super(CUT_model, self).compile()\n        self.G_optimizer = G_optimizer\n        self.F_optimizer = F_optimizer\n        self.D_optimizer = D_optimizer\n        self.gan_loss_func = GANLoss(self.gan_mode)\n        self.nce_loss_func = PatchNCELoss(self.nce_temp, self.nce_lambda)\n\n    @tf.function\n    def train_step(self, batch_data):\n        # A is source and B is target\n        real_A, real_B = batch_data\n        if self.use_nce_identity:\n            real = tf.concat([real_A, real_B], axis=0)\n        else:\n            real = real_A\n\n        with tf.GradientTape(persistent=True) as tape:\n            fake = self.netG(real, training=True)\n            fake_B = fake[:real_A.shape[0]]\n            if self.use_nce_identity:\n                idt_B = fake[real_A.shape[0]:]\n\n            # Calculate GAN loss for the discriminator\n            fake_score = self.netD(fake_B, training=True)\n            D_fake_loss = tf.reduce_mean(self.gan_loss_func(fake_score, False))\n            real_score = self.netD(real_B, training=True)\n            D_real_loss = tf.reduce_mean(self.gan_loss_func(real_score, True))\n            D_loss = (D_fake_loss + D_real_loss) * 0.5\n\n            # Calculate GAN loss and NCE loss for the generator\n            G_loss = tf.reduce_mean(self.gan_loss_func(fake_score, True))\n            NCE_loss = self.nce_loss_func(real_A, fake_B, self.netE, self.netF)\n            if self.use_nce_identity:\n                NCE_B_loss = self.nce_loss_func(real_B, idt_B, self.netE,\n                                                self.netF)\n                NCE_loss = (NCE_loss + NCE_B_loss) * 0.5\n            G_loss += NCE_loss\n\n        D_loss_grads = tape.gradient(D_loss, self.netD.trainable_variables)\n        self.D_optimizer.apply_gradients(\n            zip(D_loss_grads, self.netD.trainable_variables))\n        \n        G_loss_grads = tape.gradient(G_loss, self.netG.trainable_variables)\n        self.G_optimizer.apply_gradients(\n            zip(G_loss_grads, self.netG.trainable_variables))\n\n        F_loss_grads = tape.gradient(NCE_loss, self.netF.trainable_variables)\n        self.F_optimizer.apply_gradients(\n            zip(F_loss_grads, self.netF.trainable_variables))\n\n        del tape\n        return {'D_loss': D_loss, 'G_loss': G_loss, 'NCE_loss': NCE_loss}\n\n    @tf.function\n    def test_step(self, batch_data):\n        # A is source and B is target\n        real_A, real_B = batch_data\n        if self.use_nce_identity:\n            real = tf.concat([real_A, real_B], axis=0)\n        else:\n            real = real_A\n      \n        fake = self.netG(real, training=False)\n        fake_B = fake[:real_A.shape[0]]\n        if self.use_nce_identity:\n            idt_B = fake[real_A.shape[0]:]\n\n        # Calculate GAN loss for the discriminator\n        fake_score = self.netD(fake_B, training=False)\n        D_fake_loss = tf.reduce_mean(self.gan_loss_func(fake_score, False))\n        real_score = self.netD(real_B, training=False)\n        D_real_loss = tf.reduce_mean(self.gan_loss_func(real_score, True))\n        D_loss = (D_fake_loss + D_real_loss) * 0.5\n\n        # Calculate GAN loss and NCE loss for the generator\n        G_loss = tf.reduce_mean(self.gan_loss_func(fake_score, True))\n        NCE_loss = self.nce_loss_func(real_A, fake_B, self.netE, self.netF)\n        if self.use_nce_identity:\n            NCE_B_loss = self.nce_loss_func(real_B, idt_B, self.netE,\n                                            self.netF)\n            NCE_loss = (NCE_loss + NCE_B_loss) * 0.5\n        G_loss += NCE_loss\n\n        return {'D_loss': D_loss, 'G_loss': G_loss, 'NCE_loss': NCE_loss}","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"**Instantiate and train**"},{"metadata":{"trusted":true},"cell_type":"code","source":"mode = \"cut\"\nbatch_size = 1\nbeta_1 = 0.5\nbeta_2 = 0.999\nlr = 0.0002\nlr_decay_rate = 0.9\nlr_decay_step = 100000\n\n# Create model\ncut = CUT_model(source_shape, target_shape, cut_mode=mode,\n                use_antialias=anti_alias)\n\n# Define learning rate schedule\nlr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(\n    initial_learning_rate=lr, decay_steps=lr_decay_step,\n    decay_rate=lr_decay_rate, staircase=True)\n\n# Compile model\ncut.compile(G_optimizer=tf.keras.optimizers.Adam(learning_rate=lr_schedule,\n                                                 beta_1=beta_1, beta_2=beta_2),\n            F_optimizer=tf.keras.optimizers.Adam(learning_rate=lr_schedule,\n                                                 beta_1=beta_1, beta_2=beta_2),\n            D_optimizer=tf.keras.optimizers.Adam(learning_rate=lr_schedule,\n                                                 beta_1=beta_1, beta_2=beta_2),\n           )\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class GANMonitor(tf.keras.callbacks.Callback):\n    \"\"\"A callback to generate and save images after each epoch\"\"\"\n    def __init__(self, generator, val_dataset, out_dir, num_img=2):\n        self.num_img = num_img\n        self.generator = generator\n        self.val_dataset = val_dataset\n        self.out_dir = out_dir\n\n    def on_epoch_end(self, epoch, logs=None):\n        _, ax = plt.subplots(self.num_img, 4, figsize=(20, 10))\n        [ax[0, i].set_title(title) for i, title in enumerate(\n            ['Source', \"Translated\", \"Target\", \"Identity\"])]\n        for i, (source, target) in enumerate(\n                self.val_dataset.take(self.num_img)):\n            translated = self.generator(source, training=False)[0].numpy()\n            translated = (translated * 127.5 + 127.5).astype(np.uint8)\n            source = (source[0] * 127.5 + 127.5).numpy().astype(np.uint8)\n\n            idt = self.generator(target, training=False)[0].numpy()\n            idt = (idt * 127.5 + 127.5).astype(np.uint8)\n            target = (target[0] * 127.5 + 127.5).numpy().astype(np.uint8)\n\n            [ax[i, j].imshow(img) for j, img in enumerate(\n                [source, translated, target, idt])]\n            [ax[i, j].axis(\"off\") for j in range(4)]\n        plt.show()\n        plt.savefig(f'{self.out_dir}/epoch={epoch + 1}.png')\n        plt.close()\n\n\n# Create validating callback to generate output image every epoch\nout_dir = 'callbacks'\nif not(os.path.exists(out_dir)):\n    print(\"Creating folder...\")\n    os.makedirs(out_dir)\nplotter_callback = GANMonitor(cut.netG, val_dataset, out_dir)\n\n# Train the model for max. nr of epochs or until time-out:\nepoch = 0\ntime_start = time.time()\nhistory_val = {'D_loss': [], 'G_loss': [], 'NCE_loss': []}\nhistory_train = {'D_loss': [], 'G_loss': [], 'NCE_loss': []}\nwhile time.time() - time_start < max_train_time and epoch < max_epochs:\n    print(\"Epoch:\", epoch)\n    train_losses = cut.fit(train_dataset, epochs=1,\n                           callbacks=[plotter_callback], verbose=1)\n    history_train['D_loss'].append(train_losses.history['D_loss'][0])\n    history_train['G_loss'].append(train_losses.history['G_loss'][0])\n    history_train['NCE_loss'].append(train_losses.history['NCE_loss'][0])\n    val_losses = (cut.evaluate(val_dataset))\n    history_val['D_loss'].append(val_losses[0])\n    history_val['G_loss'].append(val_losses[1])\n    history_val['NCE_loss'].append(val_losses[2])\n    \n    print(\"Train:\", train_losses.history)\n    print(\"Val:\", val_losses)\n    epoch += 1\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"plt.plot(history_train['D_loss'], label=\"Training\")\nplt.plot(history_val['D_loss'], label=\"Validation\")\nplt.title('Discriminator Loss')\nplt.legend()\nplt.show()\n\nplt.plot(history_train['G_loss'], label=\"Training\")\nplt.plot(history_val['G_loss'], label=\"Validation\")\nplt.title('Generator Loss')\nplt.legend()\nplt.show()\n\nplt.plot(history_train['NCE_loss'], label=\"Training\")\nplt.plot(history_val['NCE_loss'], label=\"Validation\")\nplt.title('Patch NCE Loss')\nplt.legend()\nplt.show()\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"**Generate images and store them.**"},{"metadata":{"trusted":true},"cell_type":"code","source":"# Create exactly the same train data set as before, but separate from\n# the target, in order to use it as test data.\ntrain_src_dataset = tf.data.Dataset.list_files(train_src_folder+'/*.jpg',\n                                               shuffle=False)\ntrain_src_dataset = (\n    train_src_dataset.map(load_image, num_parallel_calls=AUTOTUNE)\n    .batch(batch_size, drop_remainder=True)\n    .prefetch(AUTOTUNE)\n)\n\nif not(os.path.exists('images')):\n    print(\"Creating folder...\")\n    os.makedirs('images') # Create folder to save generated images\n\ndef predict_and_save(input_ds, generator_model):\n    i = 1\n    for img in input_ds:\n        prediction = generator_model(img, training=False)\n        prediction = tf.cast((prediction * 127.5 + 127.5), tf.uint8)\n        im = tf.squeeze(prediction, 0)\n        im = im.numpy()\n        im = Image.fromarray(im)\n        im.save(\"images/\" + str(i) + '.jpg')\n        i += 1\n\npredict_and_save(train_src_dataset, cut.netG)\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"shutil.make_archive('/kaggle/working/images/', 'zip', 'images')\nshutil.rmtree(\"./images\")\n","execution_count":null,"outputs":[]}],"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat":4,"nbformat_minor":4}