pix2pix_erosion_gan.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU",
    "gpuClass": "standard"
  },
  "cells": [
    {
      "cell_type": "code",
      "source": [
        "%cd '/content/drive/MyDrive/Colab Notebooks/1/'"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "7_oGw_o5LPwn",
        "outputId": "c16b6ccb-48e5-4bdf-ace9-c9053354240c"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "/content/drive/MyDrive/Colab Notebooks/1\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install tensorflow_addons"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "MqJr3yJtMnel",
        "outputId": "b870c4e5-9850-4eab-b60f-0d69fefe2313"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
            "Requirement already satisfied: tensorflow_addons in /usr/local/lib/python3.7/dist-packages (0.17.1)\n",
            "Requirement already satisfied: typeguard>=2.7 in /usr/local/lib/python3.7/dist-packages (from tensorflow_addons) (2.7.1)\n",
            "Requirement already satisfied: packaging in /usr/local/lib/python3.7/dist-packages (from tensorflow_addons) (21.3)\n",
            "Requirement already satisfied: pyparsing!=3.0.5,>=2.0.2 in /usr/local/lib/python3.7/dist-packages (from packaging->tensorflow_addons) (3.0.9)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import tensorflow as tf\n",
        "import tensorflow_addons as tfa\n",
        "import tensorflow.keras.layers as layers\n",
        "import matplotlib.pyplot as plt\n",
        "from datetime import datetime\n",
        "import argparse\n",
        "import cv2 as cv\n",
        "import os\n",
        "import numpy as np\n",
        "import shutil"
      ],
      "metadata": {
        "id": "RD8lvsuJLcWI"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# generator\n",
        "\n",
        "input_shape = (256, 256, 1)\n",
        "\n",
        "def downsample(filters, size, apply_batchnorm=True):\n",
        "\tinitializer = tf.random_normal_initializer(0., 0.02)\n",
        "\n",
        "\tresult = tf.keras.Sequential()\n",
        "\tresult.add(\n",
        "\t\ttf.keras.layers.Conv2D(filters, size, strides=2, padding='same',\n",
        "\t\t\t\t\t\t\t\tkernel_initializer=initializer, use_bias=False))\n",
        "\n",
        "\tif apply_batchnorm:\n",
        "\t\tresult.add(tf.keras.layers.BatchNormalization())\n",
        "\n",
        "\tresult.add(tf.keras.layers.LeakyReLU())\n",
        "\n",
        "\treturn result\n",
        "\n",
        "def upsample(filters, size, apply_dropout=False):\n",
        "\tinitializer = tf.random_normal_initializer(0., 0.02)\n",
        "\n",
        "\tresult = tf.keras.Sequential()\n",
        "\tresult.add(\n",
        "\ttf.keras.layers.Conv2DTranspose(filters, size, strides=2,\n",
        "\t\t\t\t\t\t\t\t\tpadding='same',\n",
        "\t\t\t\t\t\t\t\t\tkernel_initializer=initializer,\n",
        "\t\t\t\t\t\t\t\t\tuse_bias=False))\n",
        "\n",
        "\tresult.add(tf.keras.layers.BatchNormalization())\n",
        "\n",
        "\tif apply_dropout:\n",
        "\t\tresult.add(tf.keras.layers.Dropout(0.5))\n",
        "\n",
        "\tresult.add(tf.keras.layers.ReLU())\n",
        "\n",
        "\treturn result\n",
        "\n",
        "class Generator(tf.keras.models.Model):\n",
        "\tdef __init__(self):\n",
        "\t\tsuper(Generator, self).__init__()\n",
        "\n",
        "\t\toutput_channels = 1\n",
        "\t\tinputs = layers.Input(shape=input_shape)\n",
        "\n",
        "\t\tdown_stack = [\n",
        "\t\t\tdownsample(64, 4, apply_batchnorm=False),  # (bs, 128, 128, 64)\n",
        "\t\t\tdownsample(128, 4),  # (bs, 64, 64, 128)\n",
        "\t\t\tdownsample(256, 4),  # (bs, 32, 32, 256)\n",
        "\t\t\tdownsample(512, 4),  # (bs, 16, 16, 512)\n",
        "\t\t\tdownsample(512, 4),  # (bs, 8, 8, 512)\n",
        "\t\t\tdownsample(512, 4),  # (bs, 4, 4, 512)\n",
        "\t\t\tdownsample(512, 4),  # (bs, 2, 2, 512)\n",
        "\t\t\tdownsample(512, 4),  # (bs, 1, 1, 512)\n",
        "\t\t]\n",
        "\n",
        "\t\tup_stack = [\n",
        "\t\t\tupsample(512, 4, apply_dropout=True),  # (bs, 2, 2, 1024)\n",
        "\t\t\tupsample(512, 4, apply_dropout=True),  # (bs, 4, 4, 1024)\n",
        "\t\t\tupsample(512, 4, apply_dropout=True),  # (bs, 8, 8, 1024)\n",
        "\t\t\tupsample(512, 4),  # (bs, 16, 16, 1024)\n",
        "\t\t\tupsample(256, 4),  # (bs, 32, 32, 512)\n",
        "\t\t\tupsample(128, 4),  # (bs, 64, 64, 256)\n",
        "\t\t\tupsample(64, 4),  # (bs, 128, 128, 128)\n",
        "\t\t]\n",
        "\n",
        "\t\tinitializer = tf.random_normal_initializer(0., 0.02)\n",
        "\t\tlast = tf.keras.layers.Conv2DTranspose(output_channels, 4,\n",
        "\t\t\t\t\t\t\t\t\t\t\t\tstrides=2,\n",
        "\t\t\t\t\t\t\t\t\t\t\t\tpadding='same',\n",
        "\t\t\t\t\t\t\t\t\t\t\t\tkernel_initializer=initializer,\n",
        "\t\t\t\t\t\t\t\t\t\t\t\tactivation='tanh')  # (bs, 256, 256, 3)\n",
        "\n",
        "\t\tx = inputs\n",
        "\n",
        "\t\t# Downsampling through the model\n",
        "\t\tskips = []\n",
        "\t\tfor down in down_stack:\n",
        "\t\t\tx = down(x)\n",
        "\t\t\tskips.append(x)\n",
        "\n",
        "\t\tskips = reversed(skips[:-1])\n",
        "\n",
        "\t\t# Upsampling and establishing the skip connections\n",
        "\t\tfor up, skip in zip(up_stack, skips):\n",
        "\t\t\tx = up(x)\n",
        "\t\t\tx = tf.keras.layers.Concatenate()([x, skip])\n",
        "\n",
        "\t\tx = last(x)\n",
        "\n",
        "\n",
        "\t\tself.model = tf.keras.Model(inputs=inputs, outputs=x)\n",
        "\n",
        "\t@tf.function\n",
        "\tdef call(self, input, training=True):\n",
        "\t\treturn self.model(input, training=training)\n",
        "\n",
        "\n",
        "if False:\n",
        "\tgen = Generator()\n",
        "\n",
        "\tx = tf.random.uniform((32, input_shape[0], input_shape[1]), minval=-1, maxval=1)\n",
        "\tx = tf.ones((64, input_shape[0], input_shape[1]))\n",
        "\n",
        "\n",
        "\ty = gen(x)\n",
        "\tprint(y.shape)\n",
        "\n",
        "\tprint(gen.model.summary())"
      ],
      "metadata": {
        "id": "0UvwqqIbLc4d"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# discriminator\n",
        "\n",
        "input_shape = (256, 256, 1)\n",
        "\n",
        "def downsample(filters, size, apply_batchnorm=True):\n",
        "\tinitializer = tf.random_normal_initializer(0., 0.02)\n",
        "\n",
        "\tresult = tf.keras.Sequential()\n",
        "\tresult.add(\n",
        "\t\ttf.keras.layers.Conv2D(filters, size, strides=2, padding='same',\n",
        "\t\t\t\t\t\t\t\tkernel_initializer=initializer, use_bias=False))\n",
        "\n",
        "\tif apply_batchnorm:\n",
        "\t\tresult.add(tf.keras.layers.BatchNormalization())\n",
        "\n",
        "\tresult.add(tf.keras.layers.LeakyReLU())\n",
        "\n",
        "\treturn result\n",
        "\n",
        "def upsample(filters, size, apply_dropout=False):\n",
        "\tinitializer = tf.random_normal_initializer(0., 0.02)\n",
        "\n",
        "\tresult = tf.keras.Sequential()\n",
        "\tresult.add(\n",
        "\ttf.keras.layers.Conv2DTranspose(filters, size, strides=2,\n",
        "\t\t\t\t\t\t\t\t\tpadding='same',\n",
        "\t\t\t\t\t\t\t\t\tkernel_initializer=initializer,\n",
        "\t\t\t\t\t\t\t\t\tuse_bias=False))\n",
        "\n",
        "\tresult.add(tf.keras.layers.BatchNormalization())\n",
        "\n",
        "\tif apply_dropout:\n",
        "\t\tresult.add(tf.keras.layers.Dropout(0.5))\n",
        "\n",
        "\tresult.add(tf.keras.layers.ReLU())\n",
        "\n",
        "\treturn result\n",
        "\n",
        "class Discriminator(tf.keras.models.Model):\n",
        "\tdef __init__(self):\n",
        "\t\tsuper(Discriminator, self).__init__()\n",
        "\n",
        "\n",
        "\t\tinitializer = tf.random_normal_initializer(0., 0.02)\n",
        "\t\tinp = tf.keras.layers.Input(shape=input_shape, name='input_image')\n",
        "\t\ttar = tf.keras.layers.Input(shape=input_shape, name='target_image')\n",
        "\n",
        "\t\tx = tf.keras.layers.concatenate([inp, tar])  # (bs, 256, 256, channels*2)\n",
        "\n",
        "\t\tdown1 = downsample(64, 4, False)(x)  # (bs, 128, 128, 64)\n",
        "\t\tdown2 = downsample(128, 4)(down1)  # (bs, 64, 64, 128)\n",
        "\t\tdown3 = downsample(256, 4)(down2)  # (bs, 32, 32, 256)\n",
        "\n",
        "\t\tzero_pad1 = tf.keras.layers.ZeroPadding2D()(down3)  # (bs, 34, 34, 256)\n",
        "\t\tconv = tf.keras.layers.Conv2D(512, 4, strides=1,\n",
        "\t\t\t\t\t\t\t\t\t\tkernel_initializer=initializer,\n",
        "\t\t\t\t\t\t\t\t\t\tuse_bias=False)(zero_pad1)  # (bs, 31, 31, 512)\n",
        "\n",
        "\t\tbatchnorm1 = tf.keras.layers.BatchNormalization()(conv)\n",
        "\n",
        "\t\tleaky_relu = tf.keras.layers.LeakyReLU()(batchnorm1)\n",
        "\n",
        "\t\tzero_pad2 = tf.keras.layers.ZeroPadding2D()(leaky_relu)  # (bs, 33, 33, 512)\n",
        "\n",
        "\t\tlast = tf.keras.layers.Conv2D(1, 4, strides=1,\n",
        "\t\t\t\t\t\t\t\t\t\tkernel_initializer=initializer)(zero_pad2)  # (bs, 30, 30, 1)\n",
        "\n",
        "\t\tself.model = tf.keras.Model(inputs=[inp, tar], outputs=last)\n",
        "\n",
        "\t@tf.function\n",
        "\tdef call(self, input, training=True):\n",
        "\t\treturn self.model(input, training=training)\n",
        "\n",
        "\n",
        "if False:\n",
        "\tdisc = Discriminator()\n",
        "\n",
        "\tx = tf.random.uniform((32, input_shape[0], input_shape[1]), minval=-1, maxval=1)\n",
        "\tx = tf.ones((64, input_shape[0], input_shape[1]))\n",
        "\n",
        "\n",
        "\ty = disc([x,x])\n",
        "\tprint(y)\n",
        "\n",
        "\tprint(disc.model.summary())"
      ],
      "metadata": {
        "id": "raQgoukaLyLb"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# sketch\n",
        "\n",
        "def skeletonize(img):\n",
        "\tskel = np.zeros(img.shape, np.uint8)\n",
        "\n",
        "\t# Get a Cross Shaped Kernel\n",
        "\telement = cv.getStructuringElement(cv.MORPH_CROSS, (3,3))\n",
        "\n",
        "\t# Repeat steps 2-4\n",
        "\twhile True:\n",
        "\t\t#Step 2: Open the image\n",
        "\t\topen = cv.morphologyEx(img, cv.MORPH_OPEN, element)\n",
        "\t\t#Step 3: Substract open from the original image\n",
        "\t\ttemp = cv.subtract(img, open)\n",
        "\t\t#Step 4: Erode the original image and refine the skeleton\n",
        "\t\teroded = cv.erode(img, element)\n",
        "\t\tskel = cv.bitwise_or(skel,temp)\n",
        "\t\timg = eroded.copy()\n",
        "\t\t# Step 5: If there are no white pixels left ie.. the image has been completely eroded, quit the loop\n",
        "\t\tif cv.countNonZero(img)==0:\n",
        "\t\t\tbreak\n",
        "\treturn skel\n",
        "\n",
        "def gradients(image, use_sobel=True):\n",
        "\tif use_sobel:\n",
        "\t\tgx = cv.Sobel(image,cv.CV_64F,1,0,ksize=3,scale=0.5)\n",
        "\t\tgy = cv.Sobel(image,cv.CV_64F,0,1,ksize=3,scale=0.5)\n",
        "\telse:\n",
        "\t\tgx = cv.Scharr(image,cv.CV_64F,1,0)\n",
        "\t\tgy = cv.Scharr(image,cv.CV_64F,0,1)\n",
        "\n",
        "\tg = np.sqrt(gx**2 + gy**2)\n",
        "\ttheta = np.arctan2(gy, gx)\n",
        "\n",
        "\treturn g, theta\n",
        "\n",
        "def heightmap_sketch(src):\n",
        "\t#src = cv.imread(cv.samples.findFile(image))\n",
        "\tsrc = src.astype(np.uint8)\n",
        "\n",
        "\tres = src\n",
        "\n",
        "\n",
        "\tksize = 31\n",
        "\tres = cv.GaussianBlur(res, (ksize, ksize), 0)\n",
        "\tres,t = gradients(res)\n",
        "\t#imgs.append(('orig', res))\n",
        "\t#res = cv.Laplacian(res, cv.CV_64F, ksize=5)\n",
        "\tres = res.astype(np.uint8)\n",
        "\t#imgs.append(('orig', res))\n",
        "\t#res,y = gradients(res)\n",
        "\n",
        "\n",
        "\t#res = erode(res, 15, cv.MORPH_ELLIPSE)\n",
        "\t#res = morph(res, 3, 30, cv.MORPH_ELLIPSE)\n",
        "\n",
        "\t#res = cv.bitwise_not(res)\n",
        "\t#res = cv.subtract(src, res)\n",
        "\n",
        "\t#ret, thresh = cv.threshold(res, 25, 255, cv.THRESH_BINARY)\n",
        "\n",
        "\t#res = cv.subtract(src, res)\n",
        "\n",
        "\tksize = 5\n",
        "\tres = cv.GaussianBlur(res, (ksize, ksize), 0)\n",
        "\tres = cv.ximgproc.RidgeDetectionFilter_create(ksize=3).getRidgeFilteredImage(res)\n",
        "\t#imgs.append(('orig', res))\n",
        "\tksize = 7\n",
        "\tres = cv.GaussianBlur(res, (ksize, ksize), 0)\n",
        "\tres = erode(res, 1, cv.MORPH_ELLIPSE)\n",
        "\t#imgs.append(('orig', res))\n",
        "\tret, res = cv.threshold(res, 19, 255, cv.THRESH_BINARY)\n",
        "\t#imgs.append(('orig', res))\n",
        "\tres = cv.bitwise_and(src,res)\n",
        "\t#imgs.append(('orig', res))\n",
        "\tres = skeletonize(res)\n",
        "\tret, res = cv.threshold(res, 25, 0, cv.THRESH_TOZERO)\n",
        "\t#imgs.append(('orig', res))\n",
        "\t#res = cv.Canny(res,10,220, apertureSize=5)\n",
        "\n",
        "\t#res = np.array([[1,1,1],[1,100,200],[300,1,1]])\n",
        "\n",
        "\t#res = steepestAscent(src)\n",
        "\n",
        "\t#cv.imwrite('out.png', res)\n",
        "\n",
        "\n",
        "\treturn res\n",
        "\n",
        "\n",
        "\n",
        "def morph_shape(val):\n",
        "\tif val == 0:\n",
        "\t\treturn cv.MORPH_RECT\n",
        "\telif val == 1:\n",
        "\t\treturn cv.MORPH_CROSS\n",
        "\telif val == 2:\n",
        "\t\treturn cv.MORPH_ELLIPSE\n",
        "\n",
        "def morph(target, op, size, shape):\n",
        "\terosion_size = size\n",
        "\terosion_shape = shape\n",
        "\n",
        "\telement = cv.getStructuringElement(erosion_shape, (2 * erosion_size + 1, 2 * erosion_size + 1),\n",
        "\t\t\t\t\t\t\t\t\t   (erosion_size, erosion_size))\n",
        "\treturn cv.morphologyEx(target, op, element)\n",
        "\n",
        "def erode(target, size, shape):\n",
        "\terosion_size = size\n",
        "\terosion_shape = shape\n",
        "\n",
        "\telement = cv.getStructuringElement(erosion_shape, (2 * erosion_size + 1, 2 * erosion_size + 1),\n",
        "\t\t\t\t\t\t\t\t\t   (erosion_size, erosion_size))\n",
        "\treturn cv.erode(target, element)\n",
        "\n",
        "\n",
        "def dilate(target, size, shape):\n",
        "\tdilatation_size = size\n",
        "\tdilation_shape = shape\n",
        "\n",
        "\telement = cv.getStructuringElement(dilation_shape, (2 * dilatation_size + 1, 2 * dilatation_size + 1),\n",
        "\t\t\t\t\t\t\t\t\t   (dilatation_size, dilatation_size))\n",
        "\treturn cv.dilate(target, element)\n",
        "\n",
        "\n",
        "def convert_to_labels():\n",
        "\tdir_src = 'out256_clip'\n",
        "\tdir_label = 'out256_clip_label'\n",
        "\tfor file in os.listdir(dir_src):\n",
        "\t\tpath_src = dir_src + '/' + file\n",
        "\t\tpath_dst = dir_label + '/' + file\n",
        "\n",
        "\t\timage = Image.open(path_src)\n",
        "\t\timage = np.asarray(image)\n",
        "\n",
        "\t\tres = heightmap_sketch(image)\n",
        "\n",
        "\t\tres = Image.fromarray(res)\n",
        "\t\tres.save(path_dst)"
      ],
      "metadata": {
        "id": "HBe2WmSmj09v"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "'''\n",
        "\tcasts image data from uint8 to float and maps it into -1 to 1 range\n",
        "'''\n",
        "def img_to_network(img):\n",
        "\treturn (tf.cast(img, tf.float32) / 127.5) - 1.0\n",
        "\n",
        "'''\n",
        "\tcoverts image data from float in range -1 to 1 to uint8\n",
        "'''\n",
        "def restore_img(img):\n",
        "\treturn tf.cast((img + 1.0) * 127.5, tf.uint8)\n",
        "\n",
        "def split_image(img, tile_shape):\n",
        "\timg_shape = tf.shape(img)\n",
        "\ttile_rows = tf.reshape(img, [img_shape[0], -1, tile_shape[1], img_shape[2]])\n",
        "\tserial_tiles = tf.transpose(tile_rows, [1, 0, 2, 3])\n",
        "\treturn tf.reshape(serial_tiles, [-1, tile_shape[1], tile_shape[0], img_shape[2]])\n",
        "\n",
        "def heightmap_sketch_tensor(img):\n",
        " \treturn heightmap_sketch(img.numpy())\n",
        "\n",
        "def preprocess_dataset(img_path):\n",
        "\timg = tf.io.read_file(img_path)\n",
        "\timg = tf.io.decode_png(img, channels=1, dtype=tf.uint8)\n",
        "\n",
        "\thm = img_to_network(img)\n",
        "\thm_sketch = img_to_network(tf.py_function(func=heightmap_sketch_tensor, inp=[img], Tout=tf.uint8))\n",
        "\t#hm_sketch = tf.random.uniform(shape=(256,256))\n",
        "\treturn hm, tf.expand_dims(hm_sketch,axis=2)\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "def load_dataset(dataset_folder_paths):\n",
        "\timg_paths = []\n",
        "\tfor folder_path in dataset_folder_paths:\n",
        "\t\timg_paths += [folder_path + img_name for img_name in os.listdir(folder_path) if img_name.endswith('.png')]\n",
        "\n",
        "\tAUTOTUNE = tf.data.experimental.AUTOTUNE\n",
        "\n",
        "\tds_len = len(img_paths)\n",
        "\tds = tf.data.Dataset.from_tensor_slices(img_paths)\n",
        "\tds = ds.shuffle(ds_len, reshuffle_each_iteration=True)\n",
        "\tds = ds.map(preprocess_dataset, AUTOTUNE)\n",
        "\treturn ds, ds_len"
      ],
      "metadata": {
        "id": "yLrGV85uMAcD"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class DataAugmentation():\n",
        "\tdef __init__(self):\n",
        "\t\tshape = (256,256,1)\n",
        "\t\tself.prev_hm = tf.constant(-1, shape=shape, dtype=tf.float32)\n",
        "\t\tself.prev_hm_sketch = tf.constant(-1, shape=shape, dtype=tf.float32)\n",
        "\n",
        "\t\tself.rot_aug_prob = 0.35\n",
        "\t\tself.merge_aug_prob = 0.15\n",
        "\t\tself.identity_map_prob = 0.08\n",
        "\n",
        "\n",
        "\tdef augment(self, hm, hm_sketch):\n",
        "\t\tp = tf.random.uniform(shape=(), minval=0, maxval=1)\n",
        "\n",
        "\t\tif False: # TURN OFF AUG\n",
        "\t\t\tif p < self.identity_map_prob:\n",
        "\t\t\t\treturn self.aug_identity_map(hm, hm_sketch)\n",
        "\n",
        "\t\t\tif p < self.rot_aug_prob:\n",
        "\t\t\t\thm, hm_sketch = self.aug_rotate_rand(hm, hm_sketch)\n",
        "\n",
        "\t\t\tif p < self.merge_aug_prob:\n",
        "\t\t\t\thm, hm_sketch = self.aug_merge_maps(hm, hm_sketch)\n",
        "\n",
        "\t\t\tself.prev_hm = hm\n",
        "\t\t\tself.prev_hm_sketch = hm_sketch\n",
        "\t\treturn hm, hm_sketch\n",
        "\n",
        "\tdef aug_identity_map(self, hm, hm_sketch):\n",
        "\t\tvalue = tf.constant(-1, shape=hm.shape, dtype=tf.float32)\n",
        "\t\treturn value, value\n",
        "\n",
        "\tdef aug_rotate_rand(self, hm, hm_sketch):\n",
        "\t\trand_k = tf.random.uniform(shape=[], minval=1, maxval=4, dtype=tf.int32) # k = <1,2,3>\n",
        "\t\thm = tf.image.rot90(hm, k=rand_k)\n",
        "\t\thm_sketch = tf.image.rot90(hm_sketch, k=rand_k)\n",
        "\n",
        "\t\treturn hm, hm_sketch\n",
        "\n",
        "\tdef aug_merge_maps(self, hm, hm_sketch):\n",
        "\t\thm = tf.math.maximum(hm, self.prev_hm)\n",
        "\t\thm_sketch = tf.math.maximum(hm_sketch, self.prev_hm_sketch)\n",
        "\t\t# TODO actually math.max is theoretically less correct than doing sketch on combined but results seem better\n",
        "\t\t#hm_sketch = tf.map_fn(data.heightmap_sketch_tensor_network, data.restore_img(hm), dtype=tf.float32)\n",
        "\n",
        "\t\treturn hm, hm_sketch"
      ],
      "metadata": {
        "id": "pNHaMEGBMH-g"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def train(dataset, ds_len, epochs=100, preview_epochs=0, model_save_freq_epochs=30, logs_dir=None):\n",
        "\n",
        "\ttrain_summary_writer = tf.summary.create_file_writer(logs_dir)\n",
        "\tmetric_gen_loss = tf.keras.metrics.Mean('gen_loss', dtype=tf.float32)\n",
        "\tmetric_disc_loss = tf.keras.metrics.Mean('disc_loss', dtype=tf.float32)\n",
        "\n",
        "\tgen = Generator()\n",
        "\tdisc = Discriminator()\n",
        "\n",
        "\taug = DataAugmentation()\n",
        "\n",
        "\tlearning_rate = 0.0002\n",
        "\tgen_optimizer = tf.keras.optimizers.Adam(learning_rate, 0.5)\n",
        "\tdisc_optimizer = tf.keras.optimizers.Adam(learning_rate, 0.5)\n",
        "\n",
        "\tbce = tf.keras.losses.BinaryCrossentropy(from_logits=True)\n",
        "\tl1 = tf.keras.losses.MeanAbsoluteError()\n",
        "\tl2 = tf.keras.losses.MeanSquaredError()\n",
        "\n",
        "\tbatch_size = 4\n",
        "\n",
        "\tdataset = dataset.repeat()\n",
        "\tdataset = dataset.batch(batch_size)\n",
        "\tdataset = dataset.as_numpy_iterator()\n",
        "\n",
        "\ttrain_iter = ds_len // batch_size\n",
        "\n",
        "\t@tf.function\n",
        "\tdef train_generator(hm, hm_sketch):\n",
        "\t\twith tf.GradientTape() as tape:\n",
        "\t\t\tgenerated_hm = gen(hm_sketch)\n",
        "\n",
        "\t\t\tveracity_patch_gen = disc([hm_sketch, generated_hm])\n",
        "\n",
        "\t\t\tgen_loss = bce(tf.ones_like(veracity_patch_gen), veracity_patch_gen)\n",
        "\t\t\tl1_loss = l1(hm, generated_hm)\n",
        "\n",
        "\t\t\t#feature_extr_input = tf.concat([generated_hm, hm],axis=0)\n",
        "\t\t\t#content_loss, style_loss = sc_extr.losses(feature_extr_input)\n",
        "\n",
        "\t\t\tgenerator_loss = gen_loss + (100 * l1_loss)\n",
        "\n",
        "\t\tgradient = tape.gradient(generator_loss, gen.trainable_variables)\n",
        "\t\tgen_optimizer.apply_gradients(zip(gradient, gen.trainable_variables))\n",
        "\n",
        "\t\tmetric_gen_loss(generator_loss)\n",
        "\t\treturn generator_loss\n",
        "\n",
        "\t@tf.function\n",
        "\tdef train_discriminator(hm, hm_sketch):\n",
        "\t\twith tf.GradientTape() as tape:\n",
        "\t\t\tgenerated_hm = gen(hm_sketch)\n",
        "\n",
        "\t\t\tveracity_patch_real = disc([hm_sketch, hm])\n",
        "\t\t\tveracity_patch_gen = disc([hm_sketch, generated_hm])\n",
        "\n",
        "\t\t\treal_loss = bce(tf.ones_like(veracity_patch_real), veracity_patch_real)\n",
        "\t\t\tgen_loss = bce(tf.ones_like(veracity_patch_gen), veracity_patch_gen)\n",
        "\n",
        "\t\t\tdiscriminator_loss = real_loss + gen_loss\n",
        "\n",
        "\t\tgradient = tape.gradient(discriminator_loss, disc.trainable_variables)\n",
        "\t\tdisc_optimizer.apply_gradients(zip(gradient, disc.trainable_variables))\n",
        "\n",
        "\t\tmetric_disc_loss(discriminator_loss)\n",
        "\t\treturn discriminator_loss\n",
        "\n",
        "\tprint(f'training started {datetime.now().strftime(\"%H:%M:%S\")}')\n",
        "\tprint(f'Iterations per epoch: {train_iter}')\n",
        "\tgen_loss = 0\n",
        "\tdisc_loss = 0\n",
        "\n",
        "\tstart_epoch = 1\n",
        "\tfor epoch in range(start_epoch, epochs + 1):\n",
        "\t\tprint(f'Epoch {epoch}')\n",
        "\t\tfor it in range(train_iter):\n",
        "\t\t\thm, hm_sketch = dataset.next()\n",
        "\t\t\thm, hm_sketch = aug.augment(hm, hm_sketch)\n",
        "\t\t\tdisc_loss = train_discriminator(hm, hm_sketch)\n",
        "\n",
        "\t\t\tgen_loss = train_generator(hm, hm_sketch)\n",
        "\n",
        "\t\tprint(f'gloss:{metric_gen_loss.result()} dloss:{metric_disc_loss.result()}')\n",
        "\n",
        "\t\twith train_summary_writer.as_default():\n",
        "\t\t\ttf.summary.scalar('gen_loss', metric_gen_loss.result(), step=epoch)\n",
        "\t\t\ttf.summary.scalar('disc_loss', metric_disc_loss.result(), step=epoch)\n",
        "\t\t\tmetric_gen_loss.reset_states()\n",
        "\t\t\tmetric_disc_loss.reset_states()\n",
        "\n",
        "\n",
        "\t\tif epoch % model_save_freq_epochs == 0:\n",
        "\t\t\tprint(f'Saving models to disk... time: {datetime.now().strftime(\"%H:%M:%S\")}')\n",
        "\t\t\tgen.model.save((f'{logs_dir}/models/gen_{epoch}'))\n",
        "\t\t\tdisc.model.save((f'{logs_dir}/models/disc_{epoch}'))\n",
        "\n",
        "\n",
        "\t\tif preview_epochs and epoch % preview_epochs == 0:\n",
        "\t\t\ttests = []\n",
        "\t\t\ttests.append(hm_sketch[0])\n",
        "\n",
        "\t\t\tfor i in range(1, 4):\n",
        "\t\t\t\tsketch = tf.io.read_file(f'preview_input/input{i}.png')\n",
        "\t\t\t\tsketch = tf.io.decode_png(sketch, channels=1, dtype=tf.uint8)\n",
        "\t\t\t\tsketch = img_to_network(sketch)\n",
        "\t\t\t\ttests.append(sketch)\n",
        "\n",
        "\t\t\tfor i, sketch in enumerate(tests):\n",
        "\t\t\t\tsketch = tf.expand_dims(sketch, axis=0)\n",
        "\t\t\t\tgan_hm = gen(sketch)[0]\n",
        "\t\t\t\tcv.imwrite(f'{logs_dir}/output_e{epoch}_{i}.png', restore_img(gan_hm).numpy())\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "if __name__ == '__main__':\n",
        "\tinputs = [\n",
        "\t\t#'/media/krzysztof/g/Earthshaper/out256/',\n",
        "\t\t#'/media/krzysztof/g/Earthshaper/out256_clip/',\n",
        "\t\t'inputs/'\n",
        "\t]\n",
        "\n",
        "\texperiment_label = '1'\n",
        "\t# save run on colab\n",
        "\tlogs_dir=f'/content/train_logs/{experiment_label}'\n",
        "\tepochs=2000\n",
        "\tpreview_epochs=5\n",
        "\tmodel_save_freq_epochs=1000\n",
        "\n",
        "\n",
        "\tshutil.rmtree(logs_dir, ignore_errors=True)\n",
        "\n",
        "\ttry:\n",
        "\t\tos.mkdir(logs_dir, 0o777)\n",
        "\t\tos.mkdir(f\"{logs_dir}/models\", 0o777)\n",
        "\texcept FileExistsError:\n",
        "\t\tpass\n",
        "\n",
        "\tprint(f'Loading data from: {inputs}')\n",
        "\tds, ds_len = load_dataset(inputs)\n",
        "\n",
        "\ttrain(ds, ds_len,\n",
        "\t\tepochs=epochs,\n",
        "\t\tpreview_epochs=preview_epochs,\n",
        "\t\tmodel_save_freq_epochs=model_save_freq_epochs,\n",
        "\t\tlogs_dir=logs_dir)\n"
      ],
      "metadata": {
        "id": "Crhw-WXDLWLo",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        },
        "outputId": "7fe33518-03b7-45b3-9c7c-9faeff43175c"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Loading data from: ['inputs/']\n",
            "training started 13:13:41\n",
            "Iterations per epoch: 59\n",
            "Epoch 1\n",
            "gloss:30.084659576416016 dloss:0.06326471269130707\n",
            "Epoch 2\n",
            "gloss:26.013343811035156 dloss:7.237133104354143e-05\n",
            "Epoch 3\n",
            "gloss:26.00662612915039 dloss:3.3560361771378666e-05\n",
            "Epoch 4\n",
            "gloss:26.004165649414062 dloss:1.8850974811357446e-05\n",
            "Epoch 5\n",
            "gloss:26.003734588623047 dloss:1.1911878573300783e-05\n",
            "Epoch 6\n",
            "gloss:26.003061294555664 dloss:8.118426194414496e-06\n",
            "Epoch 7\n",
            "gloss:26.00372886657715 dloss:5.832075203215936e-06\n",
            "Epoch 8\n",
            "gloss:26.003093719482422 dloss:4.3587333493633196e-06\n",
            "Epoch 9\n",
            "gloss:26.003725051879883 dloss:3.356716661073733e-06\n",
            "Epoch 10\n",
            "gloss:26.00260353088379 dloss:2.6497143608139595e-06\n",
            "Epoch 11\n",
            "gloss:26.00334930419922 dloss:2.134182068402879e-06\n",
            "Epoch 12\n",
            "gloss:26.002756118774414 dloss:1.747896931192372e-06\n",
            "Epoch 13\n",
            "gloss:26.003704071044922 dloss:1.4524757716571912e-06\n",
            "Epoch 14\n",
            "gloss:26.0020751953125 dloss:1.2222537861816818e-06\n",
            "Epoch 15\n",
            "gloss:26.003860473632812 dloss:1.0400545988886734e-06\n",
            "Epoch 16\n",
            "gloss:26.002458572387695 dloss:8.934098332247231e-07\n",
            "Epoch 17\n",
            "gloss:26.002397537231445 dloss:7.73737781400996e-07\n",
            "Epoch 18\n",
            "gloss:26.001628875732422 dloss:6.750158831891895e-07\n",
            "Epoch 19\n",
            "gloss:25.998762130737305 dloss:5.929585995545494e-07\n",
            "Epoch 20\n",
            "gloss:26.012969970703125 dloss:5.240924565441674e-07\n",
            "Epoch 21\n",
            "gloss:26.020933151245117 dloss:4.651720644233137e-07\n",
            "Epoch 22\n",
            "gloss:26.02094268798828 dloss:4.1526260474711307e-07\n",
            "Epoch 23\n",
            "gloss:26.02094078063965 dloss:3.724784392034053e-07\n",
            "Epoch 24\n",
            "gloss:26.020936965942383 dloss:3.354711566316837e-07\n",
            "Epoch 25\n",
            "gloss:26.02094078063965 dloss:3.032756694665295e-07\n",
            "Epoch 26\n",
            "gloss:26.02094078063965 dloss:2.7515642386788386e-07\n",
            "Epoch 27\n",
            "gloss:26.020933151245117 dloss:2.504643248357752e-07\n",
            "Epoch 28\n",
            "gloss:26.020936965942383 dloss:2.2865118864956457e-07\n",
            "Epoch 29\n",
            "gloss:26.020936965942383 dloss:2.0934349720391765e-07\n",
            "Epoch 30\n",
            "gloss:26.020936965942383 dloss:1.921705887752978e-07\n",
            "Epoch 31\n",
            "gloss:26.02094268798828 dloss:1.768250683653605e-07\n",
            "Epoch 32\n",
            "gloss:26.020938873291016 dloss:1.6309047623508377e-07\n",
            "Epoch 33\n",
            "gloss:26.020944595336914 dloss:1.507324327576498e-07\n",
            "Epoch 34\n",
            "gloss:26.02094078063965 dloss:1.3960935518753104e-07\n",
            "Epoch 35\n",
            "gloss:26.020938873291016 dloss:1.295497469300244e-07\n",
            "Epoch 36\n",
            "gloss:26.020938873291016 dloss:1.2043827268826135e-07\n",
            "Epoch 37\n",
            "gloss:26.020944595336914 dloss:1.1214002171300308e-07\n",
            "Epoch 38\n",
            "gloss:26.020936965942383 dloss:1.0458496291221309e-07\n",
            "Epoch 39\n",
            "gloss:26.020938873291016 dloss:9.769002673465366e-08\n",
            "Epoch 40\n",
            "gloss:26.020938873291016 dloss:9.137666268088651e-08\n",
            "Epoch 41\n",
            "gloss:26.02094078063965 dloss:8.558806285918763e-08\n",
            "Epoch 42\n",
            "gloss:26.020936965942383 dloss:8.027178921565792e-08\n",
            "Epoch 43\n"
          ]
        },
        {
          "output_type": "error",
          "ename": "KeyboardInterrupt",
          "evalue": "ignored",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
            "\u001b[0;32m<ipython-input-34-d30cf47d13ff>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m()\u001b[0m\n\u001b[1;32m    144\u001b[0m                 \u001b[0mpreview_epochs\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mpreview_epochs\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    145\u001b[0m                 \u001b[0mmodel_save_freq_epochs\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mmodel_save_freq_epochs\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 146\u001b[0;31m \t\tlogs_dir=logs_dir)\n\u001b[0m",
            "\u001b[0;32m<ipython-input-34-d30cf47d13ff>\u001b[0m in \u001b[0;36mtrain\u001b[0;34m(dataset, ds_len, epochs, preview_epochs, model_save_freq_epochs, logs_dir)\u001b[0m\n\u001b[1;32m     79\u001b[0m                         \u001b[0mdisc_loss\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtrain_discriminator\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mhm\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mhm_sketch\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     80\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 81\u001b[0;31m                         \u001b[0mgen_loss\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtrain_generator\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mhm\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mhm_sketch\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     82\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     83\u001b[0m                 \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34mf'gloss:{metric_gen_loss.result()} dloss:{metric_disc_loss.result()}'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.7/dist-packages/tensorflow/python/util/traceback_utils.py\u001b[0m in \u001b[0;36merror_handler\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m    148\u001b[0m     \u001b[0mfiltered_tb\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    149\u001b[0m     \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 150\u001b[0;31m       \u001b[0;32mreturn\u001b[0m \u001b[0mfn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    151\u001b[0m     \u001b[0;32mexcept\u001b[0m \u001b[0mException\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0me\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    152\u001b[0m       \u001b[0mfiltered_tb\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0m_process_traceback_frames\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0me\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m__traceback__\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/def_function.py\u001b[0m in \u001b[0;36m__call__\u001b[0;34m(self, *args, **kwds)\u001b[0m\n\u001b[1;32m    913\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    914\u001b[0m       \u001b[0;32mwith\u001b[0m \u001b[0mOptionalXlaContext\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_jit_compile\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 915\u001b[0;31m         \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwds\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    916\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    917\u001b[0m       \u001b[0mnew_tracing_count\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mexperimental_get_tracing_count\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/def_function.py\u001b[0m in \u001b[0;36m_call\u001b[0;34m(self, *args, **kwds)\u001b[0m\n\u001b[1;32m    945\u001b[0m       \u001b[0;31m# In this case we have created variables on the first call, so we run the\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    946\u001b[0m       \u001b[0;31m# defunned version which is guaranteed to never create variables.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 947\u001b[0;31m       \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_stateless_fn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwds\u001b[0m\u001b[0;34m)\u001b[0m  \u001b[0;31m# pylint: disable=not-callable\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    948\u001b[0m     \u001b[0;32melif\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_stateful_fn\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    949\u001b[0m       \u001b[0;31m# Release the lock early so that multiple threads can perform the call\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/function.py\u001b[0m in \u001b[0;36m__call__\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   2953\u001b[0m     \u001b[0;32mwith\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_lock\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2954\u001b[0m       (graph_function,\n\u001b[0;32m-> 2955\u001b[0;31m        filtered_flat_args) = self._maybe_define_function(args, kwargs)\n\u001b[0m\u001b[1;32m   2956\u001b[0m     return graph_function._call_flat(\n\u001b[1;32m   2957\u001b[0m         filtered_flat_args, captured_inputs=graph_function.captured_inputs)  # pylint: disable=protected-access\n",
            "\u001b[0;32m/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/function.py\u001b[0m in \u001b[0;36m_maybe_define_function\u001b[0;34m(self, args, kwargs)\u001b[0m\n\u001b[1;32m   3242\u001b[0m     \u001b[0;32mif\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0minput_signature\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mNone\u001b[0m \u001b[0;32mor\u001b[0m \u001b[0margs\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0;32mNone\u001b[0m \u001b[0;32mor\u001b[0m \u001b[0mkwargs\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   3243\u001b[0m       \u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mkwargs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mflat_args\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfiltered_flat_args\u001b[0m \u001b[0;34m=\u001b[0m\u001b[0;31m \u001b[0m\u001b[0;31m\\\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 3244\u001b[0;31m           \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_function_spec\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcanonicalize_function_inputs\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   3245\u001b[0m     \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   3246\u001b[0m       \u001b[0mflat_args\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfiltered_flat_args\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;32mNone\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/function.py\u001b[0m in \u001b[0;36mcanonicalize_function_inputs\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   2764\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2765\u001b[0m     \u001b[0;32mif\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_input_signature\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 2766\u001b[0;31m       \u001b[0minputs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mflat_inputs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfiltered_flat_inputs\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0m_convert_numpy_inputs\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0minputs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   2767\u001b[0m       \u001b[0mkwargs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mflat_kwargs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfiltered_flat_kwargs\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0m_convert_numpy_inputs\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2768\u001b[0m       \u001b[0mflat_inputs\u001b[0m \u001b[0;34m+=\u001b[0m \u001b[0mflat_kwargs\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/function.py\u001b[0m in \u001b[0;36m_convert_numpy_inputs\u001b[0;34m(inputs)\u001b[0m\n\u001b[1;32m   2804\u001b[0m         raise TypeError(f\"The output of __array__ must be an np.ndarray, \"\n\u001b[1;32m   2805\u001b[0m                         f\"got {type(a)} from {value}.\")\n\u001b[0;32m-> 2806\u001b[0;31m       \u001b[0mflat_inputs\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mindex\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mconstant_op\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mconstant\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0ma\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   2807\u001b[0m       \u001b[0mfiltered_flat_inputs\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mappend\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mflat_inputs\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mindex\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2808\u001b[0m       \u001b[0mneed_packing\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mTrue\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.7/dist-packages/tensorflow/python/framework/constant_op.py\u001b[0m in \u001b[0;36mconstant\u001b[0;34m(value, dtype, shape, name)\u001b[0m\n\u001b[1;32m    266\u001b[0m   \"\"\"\n\u001b[1;32m    267\u001b[0m   return _constant_impl(value, dtype, shape, name, verify_shape=False,\n\u001b[0;32m--> 268\u001b[0;31m                         allow_broadcast=True)\n\u001b[0m\u001b[1;32m    269\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    270\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.7/dist-packages/tensorflow/python/framework/constant_op.py\u001b[0m in \u001b[0;36m_constant_impl\u001b[0;34m(value, dtype, shape, name, verify_shape, allow_broadcast)\u001b[0m\n\u001b[1;32m    277\u001b[0m       \u001b[0;32mwith\u001b[0m \u001b[0mtrace\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mTrace\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"tf.constant\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    278\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0m_constant_eager_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mctx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mvalue\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdtype\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mshape\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mverify_shape\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 279\u001b[0;31m     \u001b[0;32mreturn\u001b[0m \u001b[0m_constant_eager_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mctx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mvalue\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdtype\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mshape\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mverify_shape\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    280\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    281\u001b[0m   \u001b[0mg\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mops\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mget_default_graph\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.7/dist-packages/tensorflow/python/framework/constant_op.py\u001b[0m in \u001b[0;36m_constant_eager_impl\u001b[0;34m(ctx, value, dtype, shape, verify_shape)\u001b[0m\n\u001b[1;32m    302\u001b[0m \u001b[0;32mdef\u001b[0m \u001b[0m_constant_eager_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mctx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mvalue\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdtype\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mshape\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mverify_shape\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    303\u001b[0m   \u001b[0;34m\"\"\"Creates a constant on the current device.\"\"\"\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 304\u001b[0;31m   \u001b[0mt\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mconvert_to_eager_tensor\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mvalue\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mctx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdtype\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    305\u001b[0m   \u001b[0;32mif\u001b[0m \u001b[0mshape\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    306\u001b[0m     \u001b[0;32mreturn\u001b[0m \u001b[0mt\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.7/dist-packages/tensorflow/python/framework/constant_op.py\u001b[0m in \u001b[0;36mconvert_to_eager_tensor\u001b[0;34m(value, ctx, dtype)\u001b[0m\n\u001b[1;32m    100\u001b[0m       \u001b[0mdtype\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mdtypes\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mas_dtype\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdtype\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mas_datatype_enum\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    101\u001b[0m   \u001b[0mctx\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mensure_initialized\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 102\u001b[0;31m   \u001b[0;32mreturn\u001b[0m \u001b[0mops\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mEagerTensor\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mvalue\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mctx\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdevice_name\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdtype\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    103\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    104\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
          ]
        }
      ]
    }
  ]
}