arnheim_3.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "spHtkST4105t"
      },
      "source": [
        "Copyright 2021 DeepMind Technologies Limited\n",
        "\n",
        "Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "you may not use this file except in compliance with the License.\n",
        "You may obtain a copy of the License at [https://www.apache.org/licenses/LICENSE-2.0](https://www.apache.org/licenses/LICENSE-2.0).\n",
        "Unless required by applicable law or agreed to in writing, software\n",
        "distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "See the License for the specific language governing permissions and\n",
        "limitations under the License.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zfSc66QYgnfR"
      },
      "source": [
        "# Arnheim 3 - Collage\n",
        "\n",
        "**Piotr Mirowski, Dylan Banarse, Mateusz Malinowski, Yotam Doron, Oriol Vinyals, Simon Osindero, Chrisantha Fernando**\n",
        "\n",
        "DeepMind, 2021\n",
        "\n",
        "![picture](https://github.com/deepmind/arnheim/raw/main/images/arnheim3_examples.png)\n",
        "Clockwise from top left: \"Sri Lankan objects\" (200 transparent patches); \"Waves\" (70 masked transparent patches with background); \"Fruit bowl\" (100 opaque patches); \"Fruit bowl\" (100 transparent patches); \"Face\" (7 opaque patches); \"Swans\" (100 masked transparent patches); \"Chicken\" (70 masked transparent patches); \"Dancer\" (40 transparent patches). See description in the [videos](https://www.youtube.com/watch?v=HKDQsrO5xF4&list=PLKhLdFXp1JN5SEV56w9OWWsT5pAz9z7G_) for settings.\n",
        "\n",
        "##An Exploration of Architectures and Losses for Painting and Drawing\n",
        "\n",
        "Arnheim 3 is an algorithm which generates collages by training (by gradient descent) a network that applies affine transformations, i.e translation, scaling, rotation, and shear, to a set of image patches. The set of image patches is subject to evolution in the outer optimisation loop. \n",
        "\n",
        "The signal for how good an image is comes from CLIP, a text-image dual encoder. This work simplifies and extends Arnheim 2 which also used CLIP but generated SVG strokes using a more complex hierarchical stroke grammar. \n",
        "\n",
        "Here you can experiment with a variety of rendering methods for combining patches in a learnable way.\n",
        "\n",
        "##Quickstart\n",
        "1. Click \"Connect\" in the top right corner.\n",
        "1. Select __\"Runtime -> Run all\"__.\n",
        "\n",
        "Play around:\n",
        "* Experiment with basic settings in the the __Configure Collage__ cell.\n",
        "* Under the __Advanced Parameters__ heading are several sections for more detailed control over collage creatio. Read the paper for insight into the different settings.\n",
        "* After changing an setting, select menu option __\"Runtime -> Run after\"__ to run all subsequent cells to generate a collage.\n",
        "\n",
        "**Note that the Colab can easily run out of memory with large populations, many patches and large patch sizes! If you start to encounter CUDA memory issues, try lowering the number of patches and restarting the Colab.**"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "#More details\n"
      ],
      "metadata": {
        "id": "e2NL9__xkEeQ"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "##New Features\n",
        "1. Tiling\n",
        "\n",
        "  Multiple images can now be tiled to create arbitrary large images. The individual images (referred to as *tiles*) are drawn sequentially starting at the top left. All the tiles overlap each other so the drawing process can blend content of neighbouring tiles. \n",
        "\n",
        "1. Compositional Images\n",
        "\n",
        "  Uses 3x3 prompts covering over-lapping regions of the image to specify different content across the whole image. The main prompt guides the direction of overall image.\n",
        "\n",
        "1. Coloured Background\n",
        "\n",
        "  User-selectable background colour or use of uploaded images.\n",
        "\n",
        "1. Interactive Patch Placement\n",
        "\n",
        "  Stop the \"Create collage loop\" cell at any time and run the \"Tinker with patches\" cell below it to adjust individual patches with sliders. Then re-run the \"Create collage loop\" cell to continue generation.\n",
        "\n",
        "##Tips\n",
        "\n",
        "**Compositional** uses 10 parallel CLIP evaluators; nine in a 3x3 overlapping configuration covering the image, and the tenth evaluating the whole image. Each region of the image can be given a different prompt, together with the global prompt. For example, a global prompt of \"A realistic landscape\" can be combined with the following local prompts and settings:\n",
        "* NUM_PATCHES = 70\n",
        "* COMPOSITIONAL_IMAGE = ON\n",
        "* PROMPT_X0_Y0 = \"a photorealistic sky with sun\"\n",
        "* PROMPT_X1_Y0 = \"a photorealistic sky\"\n",
        "* PROMPT_X2_Y0 = \"a photorealistic sky with moon\"\n",
        "* PROMPT_X0_Y1 = \"a photorealistic tree\"\n",
        "* PROMPT_X1_Y1 = \"a photorealistic tree\"\n",
        "* PROMPT_X2_Y1 = \"a photorealistic tree\"\n",
        "* PROMPT_X0_Y2 = \"a photorealistic field\"\n",
        "* PROMPT_X1_Y2 = \"a photorealistic field\"\n",
        "* PROMPT_X2_Y2 = \"a photorealistic chicken\"\n",
        "\n",
        "This process is more memory intensive so reducing the number of patches per tile helps avoid out of memory errors. \n",
        "\n",
        "**Tiling** produces hard edges if patches go outside the tile canvas. To alleviate this restrict the patch translation and keep them relatively small, try using these settings:\n",
        "\n",
        "* MIN_TRANS = -0.66\n",
        "* MAX_TRANS = 0.8\n",
        "* PATCH_MAX_PROPORTION = 5\n",
        "* FIXED_PATCH_SCALE = OFF\n",
        "\n",
        "**opacity rendering** uses alpha and depth to render semi-opaque overlapping patches which allow gradients to be used during learning. The translucency is reduced over the course of learning to end with opaque patches. When using a small number of patches evolution can perform better than learning alone. For example, with only 7 patches, a population of 10 with the Evolutionary Strategies method applied at every step can yield good results. The settings to get the face image above were:\n",
        "\n",
        "* GLOBAL_PROMPT = “Face”\n",
        "* 7 patches\n",
        "* opacity\n",
        "* 400 steps\n",
        "* ES evolution every step\n",
        "* POP_SIZE = 10\n",
        "* ROT_POS_MUTATION = 0.05\n",
        "* SCALE_MUTATION = 0.02\n",
        "* PATCH_MUTATION = 0.2\n",
        "\n",
        "**Transparency rendering** works well as gradients are more effective. Note that colours are additive so setting INITIAL_MIN_RGB=0.1 and INITIAL_MAX_RGB=0.5 helps reduce bleaching. Something to try could be:\n",
        "\n",
        "* 80 patches\n",
        "* Transparency\n",
        "* INITIAL_MIN_RGB = 0.1\n",
        "* INITIAL_MAX_RGB = 0.5\n",
        "* 15000 steps\n",
        "* Microbial GA every 100 steps\n",
        "* POP_SIZE = 2\n",
        "* LEARNING_RATE = 0.07\n",
        "* PATCH_MUTATION_PROBABILITY = 1\n",
        "* GLOBAL_PROMPT = \"Swans on a pond\"\n",
        "\n",
        "**Masked Transparency rendering** also works well as gradients are more effective.  Something to try could be:\n",
        "\n",
        "* 200 patches\n",
        "* Masked Transparency clipped\n",
        "* INITIAL_MIN_RGB = 0.7\n",
        "* INITIAL_MAX_RGB = 1.0\n",
        "* 15000 steps\n",
        "* Microbial GA every 100 steps\n",
        "* POP_SIZE = 2\n",
        "* LEARNING_RATE = 0.07\n",
        "* PATCH_MUTATION_PROBABILITY = 1\n",
        "* GLOBAL_PROMPT = \"Swans on a pond\"\n"
      ],
      "metadata": {
        "id": "RXXqg-1cj4lM"
      }
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BfEKdIbK4VZa"
      },
      "source": [
        "# Preliminaries"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5dyyH781qzIC",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "#@title Installation of libraries {vertical-output: true}\n",
        "\n",
        "!nvidia-smi -L\n",
        "\n",
        "import subprocess\n",
        "\n",
        "CUDA_version = [s for s in subprocess.check_output([\"nvcc\", \"--version\"]).decode(\"UTF-8\").split(\", \") if s.startswith(\"release\")][0].split(\" \")[-1]\n",
        "print(\"CUDA version:\", CUDA_version)\n",
        "\n",
        "if CUDA_version == \"10.0\":\n",
        "  torch_version_suffix = \"+cu100\"\n",
        "elif CUDA_version == \"10.1\":\n",
        "  torch_version_suffix = \"+cu101\"\n",
        "elif CUDA_version == \"10.2\":\n",
        "  torch_version_suffix = \"\"\n",
        "else:\n",
        "  torch_version_suffix = \"+cu110\"\n",
        "\n",
        "%cd /content/\n",
        "!pip install cssutils\n",
        "!pip install torch-tools\n",
        "!pip install   visdom\n",
        "!pip install kornia\n",
        "!pip install ftfy regex tqdm\n",
        "!pip install git+https://github.com/openai/CLIP.git --no-deps\n",
        "!pip3 install imageio==2.4.1\n",
        "!git clone https://github.com/deepmind/arnheim.git"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "c8bdUyJs4hq3"
      },
      "source": [
        "## Imports and libraries"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Imports {vertical-output: true}\n",
        "import clip\n",
        "import copy\n",
        "import cv2\n",
        "import datetime\n",
        "import glob\n",
        "from google.colab import drive\n",
        "from google.colab import files\n",
        "from google.colab.patches import cv2_imshow\n",
        "import imageio\n",
        "import io\n",
        "from kornia.color import hsv\n",
        "from matplotlib import pyplot as plt\n",
        "import moviepy.editor as mvp\n",
        "from moviepy.video.io.ffmpeg_writer import FFMPEG_VideoWriter\n",
        "import numpy as np\n",
        "import os\n",
        "import pathlib\n",
        "import random\n",
        "import requests\n",
        "from skimage.transform import resize\n",
        "import time\n",
        "import torch\n",
        "import torch.nn.functional as F\n",
        "import torchvision.transforms as transforms\n",
        "import yaml\n",
        "os.environ[\"FFMPEG_BINARY\"] = \"ffmpeg\"\n",
        "print(\"Torch version:\", torch.__version__)\n",
        "\n",
        "import arnheim.arnheim_3.src\n",
        "import arnheim.arnheim_3.src.collage as collage\n",
        "import arnheim.arnheim_3.src.video_utils as video_utils"
      ],
      "metadata": {
        "id": "aODZ9Ir_Dhu8",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "yIVTTQO-lLCx"
      },
      "outputs": [],
      "source": [
        "#@title Initialise and load CLIP model {vertical-output: true}\n",
        "torch_device=\"cuda\"\n",
        "device = torch.device(torch_device)\n",
        "CLIP_MODEL = \"ViT-B/32\"\n",
        "print(f\"Downloading CLIP model {CLIP_MODEL}...\")\n",
        "clip_model, _ = clip.load(CLIP_MODEL, device, jit=False)\n",
        "\n",
        "LOCAL_PATCH_SET = None\n",
        "LOCAL_PATCH_FILE = None\n",
        "\n",
        "!nvidia-smi"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Configure Collage"
      ],
      "metadata": {
        "id": "YIKa6msef164"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#@markdown ##Prompt\n",
        "#@markdown Enter a **global** description of the image, e.g. 'a photorealistic chicken':\n",
        "# PROMPT = \"A photorealistic chicken\"  #@param {type:\"string\"}\n",
        "\n",
        "GLOBAL_PROMPT = \"a photorealistic chicken\"  #@param {type:\"string\"}\n",
        "\n",
        "#@markdown ##Rendering\n",
        "#@markdown * **opacity** - patches are mostly opaque.\n",
        "#@markdown * **masked_transparency_clipped** - blended patches appear opaque on background.\n",
        "#@markdown * **transparency** - colours are added so black is transparent.\n",
        "#@markdown * **masked transparency normed** - very translucent blending.\n",
        "RENDER_METHOD = \"masked_transparency_clipped\"  #@param [\"opacity\", \"masked_transparency_clipped\", \"transparency\", \"masked_transparency_normed\"]\n",
        "\n",
        "#@markdown ##Patch settings\n",
        "#@markdown Select a patch set  - select example sets or your own:\n",
        "EXAMPLE_PATCH_SET = \"Animals\" #@param ['Fruit and veg', 'Sea glass', 'Handwritten MNIST', 'Animals', 'Waste', 'Human artefacts', 'Leaves', 'Broken plate', 'NONE OF ABOVE (use advanced options)']\n",
        "\n",
        "#@markdown Number of patches to use in image:\n",
        "NUM_PATCHES =   100#@param {type:\"integer\"}\n",
        "\n",
        "#**opacity** patches overlay each other using a combination of alpha and depth,\n",
        "#**transparency** _adds_ patch colours (black therefore appearing transparent),\n",
        "#**masked transparency normed** blends patches using a normalised alpha channel where areas of maximum patch overlap are opaque and all other areas are translucent.\n",
        "#and **masked transparency clipped** blends patches using a clipped alpha channel where all regions with alpha > 1 are opaque.\n",
        "\n",
        "#@markdown ##Optimisation steps\n",
        "#@markdown More optimisation steps generally produce better results but take longer:\n",
        "OPTIM_STEPS = 4000  #@param{type:\"slider\", min:200, max:20000, step:100}\n",
        "\n",
        "#@markdown ##Monitor and visualisation\n",
        "#@markdown How often to show progress image:\n",
        "TRACE_EVERY =   50#@param {type:\"integer\"}\n",
        "#@markdown How often to create a frame for video animation:\n",
        "VIDEO_STEPS =   100#@param {type:\"integer\"}\n",
        "\n"
      ],
      "metadata": {
        "cellView": "form",
        "id": "-xJvvYNP7iQZ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Advanced Settings"
      ],
      "metadata": {
        "id": "aGK__cS0HJD6"
      }
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XlVNTb_1NelG"
      },
      "source": [
        "## Spatial and Colour Transforms"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "j70Ff0qRNBjh",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "#@title Collage configuration\n",
        "COLOUR_TRANSFORMATIONS = \"RGB space\"  #@param [\"none\", \"RGB space\", \"HSV space\"]\n",
        "#@markdown Invert image colours to have a white background?\n",
        "INVERT_COLOURS = False #@param {type:\"boolean\"}\n",
        "\n",
        "CANVAS_WIDTH = 224\n",
        "CANVAS_HEIGHT = 224\n",
        "MULTIPLIER_BIG_IMAGE = 4"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "H5vRneu99YwV",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "#@title Affine transform settings\n",
        "\n",
        "#@markdown Initial translation bounds for X and Y.\n",
        "INITIAL_MIN_TRANS = -1.0  #@param{type:\"slider\", min:-1.0, max:1.0, step:0.01}\n",
        "INITIAL_MAX_TRANS = 1.0  #@param{type:\"slider\", min:-1.0, max:1.0, step:0.01}\n",
        "#@markdown Translation bounds for X and Y.\n",
        "MIN_TRANS = -1  #@param{type:\"slider\", min:-1.0, max:1.0, step:0.01}\n",
        "MAX_TRANS = 1  #@param{type:\"slider\", min:-1.0, max:1.0, step:0.01}\n",
        "#@markdown Scale bounds (> 1 means zoom out and < 1 means zoom in).\n",
        "MIN_SCALE =   1#@param {type:\"number\"}\n",
        "MAX_SCALE =   2#@param {type:\"number\"\n",
        "#@markdown Bounds on ratio between X and Y scale (default 1).\n",
        "MIN_SQUEEZE =   0.5#@param {type:\"number\"}\n",
        "MAX_SQUEEZE =   2.0#@param {type:\"number\"}\n",
        "#@markdown Shear deformation bounds (default 0).\n",
        "MIN_SHEAR = -0.2  #@param{type:\"slider\", min:-1.0, max:1.0, step:0.01}\n",
        "MAX_SHEAR = 0.2  #@param{type:\"slider\", min:-1.0, max:1.0, step:0.01}\n",
        "#@markdown Rotation bounds.\n",
        "MIN_ROT_DEG = -180 #@param{type:\"slider\", min:-180, max:180, step:1}\n",
        "MAX_ROT_DEG = 180 #@param{type:\"slider\", min:-180, max:180, step:1}\n",
        "MIN_ROT = MIN_ROT_DEG * np.pi / 180.0\n",
        "MAX_ROT = MAX_ROT_DEG * np.pi / 180.0\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cMn71brOJb1Z",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "#@title Colour transform settings\n",
        "\n",
        "#@markdown RGB\n",
        "MIN_RGB = -0.21  #@param {type:\"slider\", min: -1, max: 1, step: 0.01}\n",
        "MAX_RGB = 1.0  #@param {type:\"slider\", min: 0, max: 1, step: 0.01}\n",
        "INITIAL_MIN_RGB = 0.8  #@param {type:\"slider\", min: 0, max: 1, step: 0.01}\n",
        "INITIAL_MAX_RGB = 0.9  #@param {type:\"slider\", min: 0, max: 1, step: 0.01}\n",
        "#@markdown HSV\n",
        "MIN_HUE = 0.  #@param {type:\"slider\", min: 0, max: 1, step: 0.01}\n",
        "MAX_HUE_DEG = 360 #@param {type:\"slider\", min: 0, max: 360, step: 1}\n",
        "MAX_HUE = MAX_HUE_DEG * np.pi / 180.0\n",
        "MIN_SAT = 0.  #@param {type:\"slider\", min: 0, max: 1, step: 0.01}\n",
        "MAX_SAT = 1.  #@param {type:\"slider\", min: 0, max: 1, step: 0.01}\n",
        "MIN_VAL = 0.  #@param {type:\"slider\", min: -1, max: 1, step: 0.01}\n",
        "MAX_VAL = 1.  #@param {type:\"slider\", min: 0, max: 1, step: 0.01}"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Optimisation Settings"
      ],
      "metadata": {
        "id": "wG5OSH6zG-Or"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "FFZ31zE0AAKJ",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "#@title Training settings\n",
        "\n",
        "# Reasonable defaults:\n",
        "# OPTIM_STEP = 10000 to 20000\n",
        "# LEARNING_RATE = 0.05\n",
        "# NUM_AUGS = 16\n",
        "# GRADIENT_CLIPPING = 10.0\n",
        "# USE_NORMALIZED_CLIP = True\n",
        "\n",
        "LEARNING_RATE = 0.05    #@param{type:\"slider\", min:0.0, max:0.6, step:0.01}\n",
        "#@markdown Number of augmentations to use in evaluation:\n",
        "USE_IMAGE_AUGMENTATIONS = True #@param{type:\"boolean\"}\n",
        "NUM_AUGS = 16  #@param {type:\"integer\"}\n",
        "\n",
        "#@markdown Normalize colours for CLIP, generally leave this as True:\n",
        "USE_NORMALIZED_CLIP = False  #@param {type:\"boolean\"}\n",
        "\n",
        "#@markdown Gradient clipping during optimisation:\n",
        "GRADIENT_CLIPPING = 10.0  #@param {type:\"number\"}\n",
        "\n",
        "#@markdown Initial random search size (1 means no search):\n",
        "INITIAL_SEARCH_SIZE = 1 #@param {type:\"slider\", min:1, max:50, step:1}\n",
        "\n",
        "#@markdown Number of gradient steps in initial search (1 means no optimisation):\n",
        "INITIAL_SEARCH_NUM_STEPS = 1 #@param {type:\"slider\", min:1, max:100, step:1}\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "z-yFYf0vAS42",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "#@title Evolution settings\n",
        "\n",
        "# Reasonable defaults:\n",
        "# POP_SIZE = 2\n",
        "# EVOLUTION_FREQUENCY = 100\n",
        "# MUTION SCALES = ~0.1\n",
        "# MAX_MULTIPLE_VISUALISATIONS = 7\n",
        "\n",
        "#@markdown For evolution set POP_SIZE greater than 1:\n",
        "POP_SIZE =    2  #@param{type:\"slider\", min:1, max:100}\n",
        "EVOLUTION_FREQUENCY =  100#@param {type:\"integer\"}\n",
        "\n",
        "#@markdown ### Genetic algorithm methods\n",
        "\n",
        "#@markdown **Microbial** - loser of randomly selected pair is replaced by mutated winner. A low selection pressure.\n",
        "\n",
        "#@markdown **Evolutionary Strategies** - mutantions of the best individual replace the rest of the population. Much higher selection pressure than Microbial GA.\n",
        "GA_METHOD = \"Microbial\"  #@param [\"Evolutionary Strategies\", \"Microbial\"]\n",
        "#@markdown ### Mutation levels\n",
        "#@markdown Scale mutation applied to position and rotation, scale, distortion, colour and patch swaps.\n",
        "POS_AND_ROT_MUTATION_SCALE = 0.02  #@param{type:\"slider\", min:0.0, max:0.3, step:0.01}\n",
        "SCALE_MUTATION_SCALE = 0.02  #@param{type:\"slider\", min:0.0, max:0.3, step:0.01}\n",
        "DISTORT_MUTATION_SCALE = 0.02  #@param{type:\"slider\", min:0.0, max:0.3, step:0.01}\n",
        "COLOUR_MUTATION_SCALE = 0.02  #@param{type:\"slider\", min:0.0, max:0.3, step:0.01}\n",
        "PATCH_MUTATION_PROBABILITY = 1  #@param{type:\"slider\", min:0.0, max:1.0, step:0.1}\n",
        "#@markdown Limit the number of individuals shown during training.\n",
        "MAX_MULTIPLE_VISUALISATIONS =   5#@param {type:\"integer\"}\n",
        "#@markdown Save video of population sample over time.\n",
        "POPULATION_VIDEO = True  #@param (type:\"boolean\")\n",
        "\n",
        "USE_EVOLUTION = POP_SIZE > 1"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "37P41H1vGu-0",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "# @title Saving images on Drive\n",
        "#@markdown Displayed results can also be stored on Google Drive.\n",
        "STORE_ON_GOOGLE_DRIVE = False  #@param {type:\"boolean\"}\n",
        "GOOGLE_DRIVE_RESULTS_DIR = \"\"  #@param {type:\"string\"}\n",
        "\n",
        "MOUNT_DIR = \"/content/drive\"\n",
        "\n",
        "if STORE_ON_GOOGLE_DRIVE:\n",
        "  from google.colab import drive\n",
        "  drive.mount(MOUNT_DIR)\n",
        "  DIR_RESULTS = pathlib.PurePath(MOUNT_DIR, \"MyDrive\", GOOGLE_DRIVE_RESULTS_DIR)\n",
        "  print(f\"Storing results on Google Drive in {DIR_RESULTS}\")\n",
        "else:\n",
        "  DIR_RESULTS = \"output_\"\n",
        "  DIR_RESULTS += datetime.datetime.strftime(\n",
        "      datetime.datetime.now(), '%Y%m%d_%H%M%S')\n",
        "  print(f\"Storing results in Colab in {DIR_RESULTS}\")\n",
        "\n",
        "pathlib.Path(DIR_RESULTS).mkdir(parents=True, exist_ok=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YbGkH7v-_--H"
      },
      "source": [
        "## Patch Settings"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "LXuH-1hcAl2M",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "#@title Select segmented patches\n",
        "\n",
        "#@markdown Load patch sets from elsewhere:\n",
        "ADVANCED_PATCH_SET = \"Upload to Colab\" #@param [\"Upload to Colab\", \"Load from URL\", \"Load from Google Drive\", \"Multiple (below)\"]\n",
        "#@markdown URL if downloading .npy file from website:\n",
        "URL_TO_PATCH_FILE = \"\" #@param {type:\"string\"}\n",
        "#@markdown Path if loading .npy file from Google Drive:\n",
        "DRIVE_PATH_TO_PATCH_FILE = \"\" #@param {type:\"string\"}\n",
        "#@markdown Reuse uploaded patch file if \"Upload to Colab\" selected:\n",
        "REUSE_UPLOAD = True  #@param {type: \"boolean\"}\n",
        "\n",
        "examples = {\"Fruit and veg\" : \"fruit.npy\", \n",
        "            \"Sea glass\" : \"shore_glass.npy\",\n",
        "            \"Handwritten MNIST\" : \"handwritten_mnist.npy\",\n",
        "            \"Animals\" : \"animals.npy\",\n",
        "            \"Animal Forms\": \"animal_forms.npy\",\n",
        "            \"Plant Forms\": \"plant_forms.npy\",\n",
        "            \"Waste\": \"waste.npy\",\n",
        "            \"Human artefacts\": \"human_artefacts.npy\",\n",
        "            \"Leaves\": \"open_leaves.npy\",\n",
        "            \"Broken plate\": \"broken_plate.npy\",\n",
        "            \"Natural forms\": \"natural_forms.npy\"\n",
        "            }\n",
        "\n",
        "# Example patch set selection overrides settings here.\n",
        "if EXAMPLE_PATCH_SET in examples:\n",
        "  repo_root = \"https://storage.googleapis.com/dm_arnheim_3_assets\"\n",
        "  URL_TO_PATCH_FILE=f\"{repo_root}/collage_patches/{examples[EXAMPLE_PATCH_SET]}\"\n",
        "  PATCH_SET = \"Load from URL\"\n",
        "else:\n",
        "  PATCH_SET = ADVANCED_PATCH_SET\n",
        "\n",
        "if PATCH_SET == \"Load from Google Drive\":\n",
        "  drive.mount(MOUNT_DIR)\n",
        "  URL_TO_PATCH_FILE = str(pathlib.PurePath(MOUNT_DIR, \"MyDrive\", \n",
        "                          DRIVE_PATH_TO_PATCH_FILE))\n",
        "  patchset_filename = os.path.basename(URL_TO_PATCH_FILE)\n",
        "  # Copy patch set file locally\n",
        "  !cp {URL_TO_PATCH_FILE} {patchset_filename}\n",
        "  print(\"Using patch set\", URL_TO_PATCH_FILE)\n",
        "\n",
        "if PATCH_SET == \"Upload to Colab\":\n",
        "  if (REUSE_UPLOAD\n",
        "      and LOCAL_PATCH_SET is not None\n",
        "      and LOCAL_PATCH_FILE is not None):\n",
        "    PATCH_SET = LOCAL_PATCH_SET\n",
        "    URL_TO_PATCH_FILE = LOCAL_PATCH_FILE\n",
        "  else: \n",
        "    from google.colab import files\n",
        "    uploaded = files.upload()\n",
        "    for k, v in uploaded.items():\n",
        "      !rm -f {k}\n",
        "      open(k, 'wb').write(v)\n",
        "      print('User patch file saved to \"{name}\" with length {length} bytes'.format(\n",
        "          name=k, length=len(v)))\n",
        "      PATCH_SET = k\n",
        "      URL_TO_PATCH_FILE = f\"./{k}\"\n",
        "      LOCAL_PATCH_SET = PATCH_SET\n",
        "      LOCAL_PATCH_FILE = URL_TO_PATCH_FILE"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Edit this cell for multiple patch set support\n",
        "\n",
        "# Define lists below to use different patch settings for each tile.\n",
        "# Settings are used in order for each tile, with the list repeated as necessary.\n",
        "# SETTING THESE WILL OVERRIDE THE PATCH SETTINGS IN THE FOLLOWING CELL.\n",
        "# Set these to empty strings (\"\") to disable their use.\n",
        "#\n",
        "# For example \n",
        "# MULTIPLE_PATCH_SET=[\"shore_glass.npy\", \"animals.npy\"]\n",
        "# will use patches shore_glass for the first tile and animals for the second.\n",
        "# Because the list is repeated if necessary, if there are more tiles then \n",
        "# shore_glass will be used for all the odd tiles animals for all the even.\n",
        "\n",
        "# Use the npy file names here. They are loaded from PATCH_REPO_ROOT set here.\n",
        "PATCH_REPO_ROOT=\"https://storage.googleapis.com/dm_arnheim_3_assets/collage_patches\"\n",
        "MULTIPLE_PATCH_SET=[\"human_artefacts.npy\", \"waste.npy\", \"animal_form.npy\", \"vegetal_form.npy\"]  # e.g.  [\"shore_glass.npy\", \"animals.npy\"]\n",
        "MULTIPLE_FIXED_SCALE_PATCHES=\"\"  # e.g.  [true, true, false]\n",
        "MULTIPLE_FIXED_SCALE_COEFF=\"\"  # e.g.  [0.8, 0.3]\n",
        "MULTIPLE_PATCH_MAX_PROPORTION=\"\"  # e.g. [3, 5, 5]"
      ],
      "metadata": {
        "id": "x3V1KFxUyiEg",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "yg4ed9tyi6vZ",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "#@title Image patch sizing for low- and high-res.\n",
        "\n",
        "#@markdown Scale patches to within (image size / PATCH_MAX_PROPORTION). \n",
        "#@markdown E.g. 5 produces small patches good for tiled images.\n",
        "PATCH_MAX_PROPORTION =  4  #@param{type:\"slider\", min:2, max:8, step:1}\n",
        "\n",
        "#@markdown Alternatively, scale all patches by same amount.\n",
        "FIXED_SCALE_PATCHES = False #@param {type:\"boolean\"}\n",
        "FIXED_SCALE_COEFF =   0.3#@param {type:\"number\"}\n",
        "\n",
        "#@markdown Brighten patches.\n",
        "NORMALIZE_PATCH_BRIGHTNESS = False  #@param {type: \"boolean\"}\n",
        "\n",
        "PATCH_WIDTH_MIN = 16  \n",
        "PATCH_HEIGHT_MIN = 16 \n"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Background, Composition and Tiling\n",
        "and Background Settings"
      ],
      "metadata": {
        "id": "g35H20mf61r5"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "p6drbSbshNN0",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "# @title Configure background\n",
        "\n",
        "#@markdown Configure a background, e.g. uploaded picture or solid colour.\n",
        "# NOTE!! Check code in the rest of this cell if modifying these text strings.\n",
        "BACKGROUND = \"None (black)\" #@param [\"None (black)\", \"Solid colour below\", \"Upload image to Colab\", \"Load image from URL\"]\n",
        "# BACKGROUND = \"Load image from Google Drive\" #@param [\"None (black)\", \"Solid colour below\", \"Upload image to Colab\", \"Load image from URL\", \"Load image from Google Drive\"]\n",
        "#@markdown Background usage: Global = use image across whole image; Local = reuse same image for every tile.\n",
        "BACKGROUND_USE = \"Global\" #@param [\"Global\", \"Local\"]\n",
        "\n",
        "#@markdown Colour configuration for solid colour background.\n",
        "BACKGROUND_RED = 195 #@param {type:\"slider\", min:0, max:255, step:1}\n",
        "BACKGROUND_GREEN = 181 #@param {type:\"slider\", min:0, max:255, step:1}\n",
        "BACKGROUND_BLUE = 172 #@param {type:\"slider\", min:0, max:255, step:1}\n",
        "\n",
        "#@markdown URL if downloading image file from website:\n",
        "BACKGROUND_IMAGE_URL = \"\" #@param {type:\"string\"}\n",
        "#markdown Path if loading image file from Google Drive:\n",
        "# BACKGROUND_IMAGE_DRIVE_PATH = \"Art/Collage/Backgrounds/biggest_chicken_ever.jpg\" #@param {type:\"string\"}\n",
        "\n",
        "PROMPTS = [GLOBAL_PROMPT]\n",
        "\n",
        "background_image = None\n",
        "\n",
        "def cached_url_download(url, format):\n",
        "  cache_filename = os.path.basename(url)\n",
        "  cache = pathlib.Path(cache_filename)\n",
        "  if not cache.is_file():\n",
        "    print(\"Downloading \" + cache_filename)\n",
        "    r = requests.get(url)\n",
        "    bytesio_object = io.BytesIO(r.content)\n",
        "    with open(cache_filename, \"wb\") as f:\n",
        "        f.write(bytesio_object.getbuffer())\n",
        "  else:\n",
        "    print(\"Using cached version of \" + cache_filename)\n",
        "  if format == \"numpy\":\n",
        "    return np.load(cache, allow_pickle=True)\n",
        "  elif format == \"image as RGB\":\n",
        "    return load_image(cache_filename, show=True)\n",
        "\n",
        "def upload_files():\n",
        "  # Upload and save to Colab's disk.\n",
        "  uploaded = files.upload()\n",
        "  # Save to disk\n",
        "  for k, v in uploaded.items():\n",
        "    open(k, 'wb').write(v)\n",
        "  return list(uploaded.keys())\n",
        "\n",
        "def load_image(filename, as_cv2_image=False, show=False):\n",
        "  # Load an image as [0,1] RGB numpy array or cv2 image format.\n",
        "  img = cv2.imread(filename)\n",
        "  if show:\n",
        "    cv2_imshow(img)\n",
        "  if as_cv2_image:\n",
        "    return img  # With colour format BGR\n",
        "  img = np.asarray(img)\n",
        "  return img[..., ::-1] / 255.  # Reverse colour dim to convert BGR to RGB\n",
        "\n",
        "if BACKGROUND == \"None (black)\":\n",
        "  # 'No background' is actually a black background.\n",
        "  BACKGROUND = \"Solid colour below\"\n",
        "  BACKGROUND_RED = 0 \n",
        "  BACKGROUND_GREEN = 0\n",
        "  BACKGROUND_BLUE = 0\n",
        "\n",
        "if BACKGROUND == \"Load image from URL\":\n",
        "  background_image = cached_url_download(BACKGROUND_IMAGE_URL,\n",
        "                                         format=\"image as RGB\")\n",
        "elif BACKGROUND == \"Solid colour below\":\n",
        "  background_image = np.ones((10, 10, 3), dtype=np.float32)\n",
        "  background_image[:, :, 0] = BACKGROUND_RED\n",
        "  background_image[:, :, 1] = BACKGROUND_GREEN\n",
        "  background_image[:, :, 2] = BACKGROUND_BLUE\n",
        "  background_image /= 255.\n",
        "  print('Defined background colour ({}, {}, {})'.format(\n",
        "      BACKGROUND_RED, BACKGROUND_GREEN, BACKGROUND_BLUE))\n",
        "elif BACKGROUND == \"Load image from Google Drive\":\n",
        "  drive.mount(MOUNT_DIR)\n",
        "  data_file = pathlib.PurePath(MOUNT_DIR, \"MyDrive\", \n",
        "                               BACKGROUND_IMAGE_DRIVE_PATH)\n",
        "  print(\"Reading\", data_file)\n",
        "  background_image = load_image(data_file)\n",
        "else:  # \"Upload image to Colab\"\n",
        "  backgrounds = upload_files()\n",
        "  background_image = load_image(backgrounds[0], show=True)\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# @title Composition prompts (i.e. for regions within a tile)\n",
        "\n",
        "#@markdown Use additional prompts for each region:\n",
        "COMPOSITIONAL_IMAGE = False #@param {type:\"boolean\"}\n",
        "\n",
        "#@markdown **Single image composition prompts** (i.e. no tiling) for the 3x3 regions (left to right, starting at the top):\n",
        "PROMPT_x0_y0 = \"a photorealistic sky with sun\"   #@param {type:\"string\"}\n",
        "PROMPT_x1_y0 = \"a photorealistic sky\"   #@param {type:\"string\"}\n",
        "PROMPT_x2_y0 = \"a photorealistic sky with moon\"   #@param {type:\"string\"}\n",
        "PROMPT_x0_y1 = \"a photorealistic tree\"   #@param {type:\"string\"}\n",
        "PROMPT_x1_y1 = \"a photorealistic tree\"   #@param {type:\"string\"}\n",
        "PROMPT_x2_y1 = \"a photorealistic tree\"   #@param {type:\"string\"}\n",
        "PROMPT_x0_y2 = \"a photorealistic field\"   #@param {type:\"string\"}\n",
        "PROMPT_x1_y2 = \"a photorealistic field\"   #@param {type:\"string\"}\n",
        "PROMPT_x2_y2 = \"a photorealistic chicken\"   #@param {type:\"string\"}\n",
        "\n",
        "#@markdown **Tiled images composition prompts** (use when tiling)\n",
        "\n",
        "#@markdown This string is formated to autogenerate compositional prompts for each tile, using each tile's prompt. e.g. \"close-up of {}\"\n",
        "TILE_PROMPT_FORMATING = \"close-up of {}\"  #@param {type:\"string\"}\n",
        "\n",
        "# Example prompt lists for different settings, where\n",
        "# PROMPT = \"Roman\"\n",
        "# TILE_PROMPT_FORMATING = \"close-up of {}\"\n",
        "# TILE_PROMPT_STRING = \"sun | clouds | sky / fields | fields | trees\"\n",
        "\n",
        "# 1. Single image with **global** prompt\n",
        "#   * Tile 0 prompts: ['Roman']\n",
        "# 1. Single image with **composition** prompts (tested)\n",
        "#   * Tile 0 prompts: ['close-up of Roman', 'close-up of Roman', 'close-up of Roman', 'close-up of Roman', 'close-up of Roman', 'close-up of Roman', 'close-up of Roman', 'close-up of Roman', 'close-up of Roman', 'Roman']\n",
        "# 1. Tiled images with **global** prompt for each tile.\n",
        "#   * Tile 0 prompts: ['Roman']\n",
        "#   * Tile 1 prompts: ['Roman']\n",
        "#   * Tile 2 prompts: ['Roman']\n",
        "#   * Tile 3 prompts: ['Roman']\n",
        "#   * Tile 4 prompts: ['Roman']\n",
        "#   * Tile 5 prompts: ['Roman']\n",
        "# 1. Tiled images with **global** prompt for each tile.\n",
        "#   * Tile 0 prompts: ['sun']\n",
        "#   * Tile 1 prompts: ['clouds']\n",
        "#   * Tile 2 prompts: ['sky']\n",
        "#   * Tile 3 prompts: ['fields']\n",
        "#   * Tile 4 prompts: ['fields']\n",
        "#   * Tile 5 prompts: ['trees']\n",
        "# 1. Tiled images with separate **composition** prompts for each tile.\n",
        "#   * Tile 0 prompts: ['close-up of sun', 'close-up of sun', 'close-up of sun', 'close-up of sun', 'close-up of sun', 'close-up of sun', 'close-up of sun', 'close-up of sun', 'close-up of sun', 'sun']\n",
        "#   * Tile 1 prompts: ['close-up of clouds', 'close-up of clouds', 'close-up of clouds', 'close-up of clouds', 'close-up of clouds', 'close-up of clouds', 'close-up of clouds', 'close-up of clouds', 'close-up of clouds', 'clouds']\n",
        "#   * Tile 2 prompts: ['close-up of sky', 'close-up of sky', 'close-up of sky', 'close-up of sky', 'close-up of sky', 'close-up of sky', 'close-up of sky', 'close-up of sky', 'close-up of sky', 'sky']\n",
        "#   * Tile 3 prompts: ['close-up of fields', 'close-up of fields', 'close-up of fields', 'close-up of fields', 'close-up of fields', 'close-up of fields', 'close-up of fields', 'close-up of fields', 'close-up of fields', 'fields']\n",
        "#   * Tile 4 prompts: ['close-up of fields', 'close-up of fields', 'close-up of fields', 'close-up of fields', 'close-up of fields', 'close-up of fields', 'close-up of fields', 'close-up of fields', 'close-up of fields', 'fields']\n",
        "#   * Tile 5 prompts: ['close-up of trees', 'close-up of trees', 'close-up of trees', 'close-up of trees', 'close-up of trees', 'close-up of trees', 'close-up of trees', 'close-up of trees', 'close-up of trees', 'trees']\n",
        "# [188]\n",
        "# \n",
        "# 1. Tiled images with **global** **composition** prompts for each tile.\n",
        "#   * Tile 0 prompts: ['close-up of Roman', 'close-up of Roman', 'close-up of Roman', 'close-up of Roman', 'close-up of Roman', 'close-up of Roman', 'close-up of Roman', 'close-up of Roman', 'close-up of Roman', 'Roman']\n",
        "#   * Tile 1 prompts: ['close-up of Roman', 'close-up of Roman', 'close-up of Roman', 'close-up of Roman', 'close-up of Roman', 'close-up of Roman', 'close-up of Roman', 'close-up of Roman', 'close-up of Roman', 'Roman']\n",
        "#   * Tile 2 prompts: ['close-up of Roman', 'close-up of Roman', 'close-up of Roman', 'close-up of Roman', 'close-up of Roman', 'close-up of Roman', 'close-up of Roman', 'close-up of Roman', 'close-up of Roman', 'Roman']\n",
        "#   * Tile 3 prompts: ['close-up of Roman', 'close-up of Roman', 'close-up of Roman', 'close-up of Roman', 'close-up of Roman', 'close-up of Roman', 'close-up of Roman', 'close-up of Roman', 'close-up of Roman', 'Roman']\n",
        "#   * Tile 4 prompts: ['close-up of Roman', 'close-up of Roman', 'close-up of Roman', 'close-up of Roman', 'close-up of Roman', 'close-up of Roman', 'close-up of Roman', 'close-up of Roman', 'close-up of Roman', 'Roman']\n",
        "#   * Tile 5 prompts: ['close-up of Roman', 'close-up of Roman', 'close-up of Roman', 'close-up of Roman', 'close-up of Roman', 'close-up of Roman', 'close-up of Roman', 'close-up of Roman', 'close-up of Roman', 'Roman']"
      ],
      "metadata": {
        "id": "Iz10BbmO4w5n",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# @title Tile prompts and tiling settings\n",
        "\n",
        "TILE_IMAGES = False #@param {type:\"boolean\"}\n",
        "\n",
        "TILES_WIDE = 2  #@param {type:\"slider\", min:1, max:10, step:1}\n",
        "TILES_HIGH = 2  #@param {type:\"slider\", min:1, max:10, step:1}\n",
        "\n",
        "# Turn off tiling if either boolean is set or width/height set to 1.\n",
        "if not TILE_IMAGES or (TILES_WIDE == 1 and TILES_HIGH == 1):\n",
        "  TILES_WIDE = 1\n",
        "  TILES_HIGH = 1\n",
        "  TILE_IMAGES = False\n",
        "  \n",
        "#@markdown **Prompt(s) for tiles**\n",
        "\n",
        "#@markdown **Global tile prompt** uses GLOBAL_PROMPT (previous cell) for *all* tiles (e.g. \"Roman mosaic of an unswept floor\").\n",
        "GLOBAL_TILE_PROMPT = False #@param {type:\"boolean\"}\n",
        "\n",
        "#@markdown Otherwise, specify **separate prompt for each tile** (overriding GLOBAL_PROMPT) with columns separated by | and / to delineate new row.\n",
        "\n",
        "#@markdown E.g. multiple prompts for a 3x2 \"landscape\" image : \"sun | clouds | sky / fields | fields | trees\".\n",
        "\n",
        "TILE_PROMPT_STRING = \"photorealistic sun | photorealistic clouds / photograph of colorful buildings | crowds of people\"   #@param {type:\"string\"}\n",
        "\n",
        "if not TILE_IMAGES or GLOBAL_TILE_PROMPT:\n",
        "  TILE_PROMPTS = [GLOBAL_PROMPT] * TILES_HIGH * TILES_WIDE\n",
        "else:\n",
        "  TILE_PROMPTS = []\n",
        "  count_y = 0\n",
        "  count_x = 0\n",
        "  for row in TILE_PROMPT_STRING.split(\"/\"):\n",
        "    for prompt in row.split(\"|\"):\n",
        "      prompt = prompt.strip()\n",
        "      TILE_PROMPTS.append(prompt)\n",
        "      count_x += 1\n",
        "    if count_x != TILES_WIDE:\n",
        "      raise ValueError(f\"Insufficient prompts for row {count_y}; expected {TILES_WIDE} but got {count_x}\")\n",
        "    count_x = 0\n",
        "    count_y += 1\n",
        "  if count_y != TILES_HIGH:\n",
        "    raise ValueError(f\"Insufficient prompt rows; expected {TILES_HIGH} but got {count_y}\")\n",
        "\n",
        "print(\"Tile prompts: \", TILE_PROMPTS)\n"
      ],
      "metadata": {
        "cellView": "form",
        "id": "QFvC6js6LVuo"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Initial Checkpoint"
      ],
      "metadata": {
        "id": "ZHAkxmLIPC7v"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Select initial checkpoint\n",
        "\n",
        "# Default filename for the initial checkpoint.\n",
        "DEFAULT_INIT_CHECKPOINT = \"generator.pt\"\n",
        "\n",
        "#@markdown Use an existing `generator.pt` file as checkpoint (if it is present in the local colab), upload it from the computer or upload it from Google Drive:\n",
        "CHOICE_INIT_CHECKPOINT = \"Use existing checkpoint file (if present)\"  #@param [\"Use existing checkpoint file (if present)\", \"Upload checkpoint file to colab\", \"Upload checkpoint file from Google Drive\"]\n",
        "\n",
        "#@markdown Enable this to use that checkpoint to initialise the collage.\n",
        "USE_INIT_CHECKPOINT = False #@param {type:\"boolean\"}\n",
        "\n",
        "#@markdown Path if loading `generator.pt` file from Google Drive:\n",
        "DRIVE_PATH_TO_INIT_CHECKPOINT = \"\" #@param {type:\"string\"}\n",
        "\n",
        "if CHOICE_INIT_CHECKPOINT == \"Upload checkpoint file from Google Drive\":\n",
        "  print(f\"Mounting directory {DRIVE_PATH_TO_INIT_CHECKPOINT}\")\n",
        "  drive.mount(MOUNT_DIR)\n",
        "  URL_TO_INIT_CHECKPOINT = str(pathlib.PurePath(MOUNT_DIR, \"MyDrive\",\n",
        "                                                DRIVE_PATH_TO_INIT_CHECKPOINT))\n",
        "  init_checkpoint_filename = os.path.basename(URL_TO_INIT_CHECKPOINT)\n",
        "  # Copy patch set file locally\n",
        "  try:\n",
        "    INIT_CHECKPOINT = DEFAULT_INIT_CHECKPOINT\n",
        "    print(f\"Copying checkpoint {INIT_CHECKPOINT} from {URL_TO_INIT_CHECKPOINT}\")\n",
        "    !cp \"{URL_TO_INIT_CHECKPOINT}/{INIT_CHECKPOINT}\" {INIT_CHECKPOINT}\n",
        "  except:\n",
        "    USE_INIT_CHECKPOINT = False\n",
        "    INIT_CHECKPOINT = ''\n",
        "    print(f\"Could not copy file from {URL_TO_INIT_CHECKPOINT}\")\n",
        "\n",
        "if CHOICE_INIT_CHECKPOINT == \"Upload checkpoint file to colab\":\n",
        "  uploaded = files.upload()\n",
        "\n",
        "if USE_INIT_CHECKPOINT:\n",
        "  INIT_CHECKPOINT = DEFAULT_INIT_CHECKPOINT\n",
        "else:\n",
        "  INIT_CHECKPOINT = ''\n",
        "\n",
        "# Check that the file is present.\n",
        "path_ckpt = \"/content/\" + INIT_CHECKPOINT\n",
        "is_checkpoint_present = os.path.isfile(\"/content/\" + INIT_CHECKPOINT)\n",
        "if USE_INIT_CHECKPOINT is True and is_checkpoint_present is True:\n",
        "  print(f\"Initialising collage from checkpoint stored in file {path_ckpt}\")\n",
        "if USE_INIT_CHECKPOINT is True and is_checkpoint_present is False:\n",
        "  USE_INIT_CHECKPOINT = False\n",
        "  INIT_CHECKPOINT = \"\"\n",
        "  print(f\"Checkpoint {path_ckpt} is missing.\")\n",
        "if USE_INIT_CHECKPOINT is False:\n",
        "  INIT_CHECKPOINT = \"\"\n",
        "  print(f\"Not using a checkpoint to initialise collage.\")\n"
      ],
      "metadata": {
        "cellView": "form",
        "id": "7qiqvoMyPOgv"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "#Make Collage"
      ],
      "metadata": {
        "id": "P9DR3C-PGe8I"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Create config\n",
        "\n",
        "# Do not edit this directly as it may not have an effect as some assets will\n",
        "# have already been created at this point, e.g. the background.\n",
        "\n",
        "# Safety check on checkpoint.\n",
        "is_checkpoint_present = os.path.isfile(\"/content/\" + INIT_CHECKPOINT)\n",
        "if is_checkpoint_present is False:\n",
        "  INIT_CHECKPOINT = \"\"\n",
        "\n",
        "config = dict(\n",
        "  background_blue=BACKGROUND_BLUE,\n",
        "  background_green=BACKGROUND_GREEN,\n",
        "  background_red=BACKGROUND_RED,\n",
        "  background_url=BACKGROUND_IMAGE_URL,\n",
        "  background_use=BACKGROUND_USE,\n",
        "  canvas_height=CANVAS_HEIGHT,\n",
        "  canvas_width=CANVAS_WIDTH,\n",
        "  clean_up=False,\n",
        "  clip_model=CLIP_MODEL,\n",
        "  colour_mutation_scale=COLOUR_MUTATION_SCALE,\n",
        "  colour_transformations=COLOUR_TRANSFORMATIONS,\n",
        "  compositional_image=COMPOSITIONAL_IMAGE,\n",
        "  cuda=True,\n",
        "  distort_mutation_scale=DISTORT_MUTATION_SCALE,\n",
        "  evolution_frequency=EVOLUTION_FREQUENCY,\n",
        "  fixed_scale_coeff=FIXED_SCALE_COEFF,\n",
        "  fixed_scale_patches=FIXED_SCALE_PATCHES,\n",
        "  ga_method=GA_METHOD,\n",
        "  global_prompt=GLOBAL_PROMPT,\n",
        "  global_tile_prompt=GLOBAL_TILE_PROMPT,\n",
        "  gradient_clipping=GRADIENT_CLIPPING,\n",
        "  gui=True,\n",
        "  high_res_multiplier=MULTIPLIER_BIG_IMAGE,\n",
        "  init_checkpoint=INIT_CHECKPOINT,\n",
        "  initial_max_rgb=INITIAL_MAX_RGB,\n",
        "  initial_min_rgb=INITIAL_MIN_RGB,\n",
        "  initial_search_size=INITIAL_SEARCH_SIZE,\n",
        "  initial_search_num_steps=INITIAL_SEARCH_NUM_STEPS,\n",
        "  invert_colours=INVERT_COLOURS,\n",
        "  learning_rate=LEARNING_RATE,\n",
        "  max_block_size_high_res=2000,\n",
        "  max_hue_deg=MAX_HUE_DEG,\n",
        "  max_multiple_visualizations=MAX_MULTIPLE_VISUALISATIONS,\n",
        "  max_rgb=MAX_RGB,\n",
        "  max_rot_deg=MAX_ROT_DEG,\n",
        "  max_sat=MAX_SAT,\n",
        "  max_scale=MAX_SCALE,\n",
        "  max_shear=MAX_SHEAR,\n",
        "  max_squeeze=MAX_SQUEEZE,\n",
        "  max_trans=MAX_TRANS,\n",
        "  max_trans_init=INITIAL_MAX_TRANS,\n",
        "  max_val=MAX_VAL,\n",
        "  min_hue_deg=MAX_HUE_DEG,\n",
        "  min_rgb=MIN_RGB,\n",
        "  min_rot_deg=MIN_ROT_DEG,\n",
        "  min_sat=MIN_SAT,\n",
        "  min_scale=MIN_SCALE,\n",
        "  min_shear=MIN_SHEAR,\n",
        "  min_squeeze=MIN_SQUEEZE,\n",
        "  min_trans=MIN_TRANS,\n",
        "  min_trans_init=INITIAL_MIN_TRANS,\n",
        "  min_val=MIN_VAL,\n",
        "  multiple_patch_set=MULTIPLE_PATCH_SET,  # e.g.  [\"shore_glass.npy\", \"animals.npy\"]\n",
        "  multiple_fixed_scale_patches=MULTIPLE_FIXED_SCALE_PATCHES,  # e.g.  [true, true, false]\n",
        "  multiple_fixed_scale_coeff=MULTIPLE_FIXED_SCALE_COEFF,  # e.g.  [0.8, 0.3]\n",
        "  multiple_patch_max_proportion=MULTIPLE_PATCH_MAX_PROPORTION,  # e.g. [3, 5, 5]\n",
        "  normalize_patch_brightness=NORMALIZE_PATCH_BRIGHTNESS,\n",
        "  num_augs=NUM_AUGS,\n",
        "  num_patches=NUM_PATCHES,\n",
        "  optim_steps=OPTIM_STEPS,\n",
        "  output_dir=DIR_RESULTS,\n",
        "  patch_height_min=PATCH_HEIGHT_MIN,\n",
        "  patch_max_proportion=PATCH_MAX_PROPORTION,\n",
        "  patch_mutation_probability=PATCH_MUTATION_PROBABILITY,\n",
        "  patch_repo_root=PATCH_REPO_ROOT,\n",
        "  patch_set=PATCH_SET,\n",
        "  patch_width_min=PATCH_WIDTH_MIN,\n",
        "  pop_size=POP_SIZE,\n",
        "  population_video=POPULATION_VIDEO,\n",
        "  pos_and_rot_mutation_scale=POS_AND_ROT_MUTATION_SCALE,\n",
        "  prompt_x0_y0=PROMPT_x0_y0,\n",
        "  prompt_x0_y1=PROMPT_x0_y1,\n",
        "  prompt_x0_y2=PROMPT_x0_y2,\n",
        "  prompt_x1_y0=PROMPT_x1_y0,\n",
        "  prompt_x1_y1=PROMPT_x1_y1,\n",
        "  prompt_x1_y2=PROMPT_x1_y2,\n",
        "  prompt_x2_y0=PROMPT_x2_y0,\n",
        "  prompt_x2_y1=PROMPT_x2_y1,\n",
        "  prompt_x2_y2=PROMPT_x2_y2,\n",
        "  render_method=RENDER_METHOD,\n",
        "  save_all_arrays=False,\n",
        "  scale_mutation_scale=SCALE_MUTATION_SCALE,\n",
        "  tile_images=TILE_IMAGES,\n",
        "  tile_prompt_formating=TILE_PROMPT_FORMATING,\n",
        "  tile_prompt_string=TILE_PROMPT_STRING,\n",
        "  tiles_high=TILES_HIGH,\n",
        "  tiles_wide=TILES_WIDE,\n",
        "  torch_device=torch_device,\n",
        "  trace_every=TRACE_EVERY,\n",
        "  url_to_patch_file=URL_TO_PATCH_FILE,\n",
        "  use_image_augmentations=USE_IMAGE_AUGMENTATIONS,\n",
        "  use_normalized_clip=USE_NORMALIZED_CLIP,\n",
        "  video_steps=VIDEO_STEPS,\n",
        ")\n",
        "\n",
        "\n"
      ],
      "metadata": {
        "id": "mh_eYcz1phjX",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Initialisation\n",
        "\n",
        "#TODO(dylski) Move this code into a module for Colab and main.py to share.\n",
        "# Adjust config for compositional image.\n",
        "if config[\"compositional_image\"] == True:\n",
        "  print(\"Generating compositional image\")\n",
        "  config['canvas_width'] *= 2\n",
        "  config['canvas_height'] *= 2\n",
        "  config['high_res_multiplier'] = int(config['high_res_multiplier'] / 2)\n",
        "  print(\"Using one image augmentations for compositional image creation.\")\n",
        "  config[\"use_image_augmentations\"] = True\n",
        "  config[\"num_augs\"] = 1\n",
        "\n",
        "# Turn off tiling if either boolean is set or width/height set to 1.\n",
        "if (not config[\"tile_images\"] or\n",
        "    (config[\"tiles_wide\"] == 1 and config[\"tiles_high\"] == 1)):\n",
        "  print(\"No tiling.\")\n",
        "  config[\"tiles_wide\"] = 1\n",
        "  config[\"tiles_high\"] = 1\n",
        "  config[\"tile_images\"] = False\n",
        "\n",
        "# Default output dir. Make sure it is a new directory for each generation.\n",
        "if len(config[\"output_dir\"]) == 0 or STORE_ON_GOOGLE_DRIVE is False:\n",
        "  config[\"output_dir\"] = \"output_\"\n",
        "  config[\"output_dir\"] += datetime.datetime.strftime(\n",
        "      datetime.datetime.now(), '%Y%m%d_%H%M%S')\n",
        "  config[\"output_dir\"] += '/'\n",
        "\n",
        "# Make output dir.\n",
        "output_dir = config[\"output_dir\"]\n",
        "print(f\"Storing results in {output_dir}\\n\")\n",
        "pathlib.Path(output_dir).mkdir(parents=True, exist_ok=True)\n",
        "\n",
        "# Save the config.\n",
        "config_filename = config[\"output_dir\"] + \"/\" + \"config.yaml\"\n",
        "with open(config_filename, \"w\") as f:\n",
        "  yaml.dump(config, f, default_flow_style=False, allow_unicode=True)\n",
        "\n",
        "# TODO(dylski) Put this into the package code.\n",
        "# Tiling.\n",
        "if not config[\"tile_images\"] or config[\"global_tile_prompt\"]:\n",
        "  tile_prompts = (\n",
        "    [config[\"global_prompt\"]] * config[\"tiles_high\"] * config[\"tiles_wide\"])\n",
        "else:\n",
        "  tile_prompts = []\n",
        "  count_y = 0\n",
        "  count_x = 0\n",
        "  for row in config[\"tile_prompt_string\"].split(\"/\"):\n",
        "    for prompt in row.split(\"|\"):\n",
        "      prompt = prompt.strip()\n",
        "      tile_prompts.append(prompt)\n",
        "      count_x += 1\n",
        "    if count_x != config[\"tiles_wide\"]:\n",
        "      w = config[\"tiles_wide\"]\n",
        "      raise ValueError(\n",
        "        f\"Insufficient prompts for row {count_y}; expected {w}, got {count_x}\")\n",
        "    count_x = 0\n",
        "    count_y += 1\n",
        "  if count_y != config[\"tiles_high\"]:\n",
        "    h = config[\"tiles_high\"]\n",
        "    raise ValueError(f\"Insufficient prompt rows; expected {h}, got {count_y}\")\n",
        "\n",
        "\n",
        "print(\"Tile prompts: \", tile_prompts)\n",
        "# Prepare duplicates of config data if required for tiles.\n",
        "tile_count = 0\n",
        "all_prompts = []\n",
        "for y in range(config[\"tiles_high\"]):\n",
        "  for x in range(config[\"tiles_wide\"]):\n",
        "    list_tile_prompts = []\n",
        "    if config[\"compositional_image\"]:\n",
        "      if config[\"tile_images\"]:\n",
        "        list_tile_prompts = [\n",
        "            config[\"tile_prompt_formating\"].format(tile_prompts[tile_count])\n",
        "            ] * 9\n",
        "      else:\n",
        "        list_tile_prompts = [\n",
        "            config[\"prompt_x0_y0\"], config[\"prompt_x1_y0\"],\n",
        "            config[\"prompt_x2_y0\"],\n",
        "            config[\"prompt_x0_y1\"], config[\"prompt_x1_y1\"],\n",
        "            config[\"prompt_x2_y1\"],\n",
        "            config[\"prompt_x0_y2\"], config[\"prompt_x1_y2\"],\n",
        "            config[\"prompt_x2_y2\"]]\n",
        "    list_tile_prompts.append(tile_prompts[tile_count])\n",
        "    tile_count += 1\n",
        "    all_prompts.append(list_tile_prompts)\n",
        "print(f\"All prompts: {all_prompts}\")\n",
        "\n",
        "ct = collage.CollageTiler(\n",
        "    all_prompts, background_image, clip_model, device, config)\n",
        "ct.initialise()\n",
        "\n",
        "generator = ct._collage_maker.generator\n",
        "img = generator({'gamma': 0})\n",
        "img = img.permute(0, 3, 1, 2)\n",
        "_ = video_utils.show_and_save(img, config=config)"
      ],
      "metadata": {
        "id": "QKfQHdXtvspC",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Create collage loop\n",
        "#@markdown To edit patches interrupt this cell and run the one below this. Re-run this cell afterwards to continue generating the image.\n",
        "output = ct.loop()"
      ],
      "metadata": {
        "id": "Y6QyfhZ5VBsJ",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Tinker with patches\n",
        "#@markdown Enable this cell to allow patch editing:\n",
        "PATCH_TINKERING = False #@param {type:\"boolean\"}\n",
        "\n",
        "#@markdown Interupt the cell above mid-optimisation and run this cell to manually adjust the patches. Run it several times to adjust different patches. Then re-run the cell above to continue optimising.\n",
        "\n",
        "if PATCH_TINKERING:\n",
        "  from ipywidgets import interactive\n",
        "  import IPython.display\n",
        "  from google.colab.output import eval_js\n",
        "  import base64\n",
        "  \n",
        "  # Render the current collage(s).\n",
        "  generator = ct._collage_maker.generator\n",
        "  step = ct._collage_maker.step\n",
        "  params = {'gamma': step / OPTIM_STEPS}\n",
        "  img = generator(params)\n",
        "  img = img.permute(0, 3, 1, 2)  # NHWC -> NCHW\n",
        "  print('Current collage(s)')\n",
        "  res_img = show_stitched_batch(img,\n",
        "                                max_display=MAX_MULTIPLE_VISUALISATIONS,\n",
        "                                show=False)\n",
        "  filename_temp = f\"./temp.png\"\n",
        "  res_img = cv2.cvtColor(res_img, cv2.COLOR_BGR2RGB) * 255\n",
        "  cv2.imwrite(filename_temp, res_img)\n",
        "  \n",
        "  # HTML code to plot the image and detect the mouse cursor.\n",
        "  canvas_html = \"\"\"\n",
        "  <canvas width=%d height=%d></canvas>\n",
        "  <script>\n",
        "    var filename_image = \"%s\"\n",
        "    var canvas = document.querySelector('canvas')\n",
        "    var ctx = canvas.getContext('2d')\n",
        "    ctx.lineWidth = 1\n",
        "    var mouse = {x: 0, y: 0}\n",
        "    canvas.addEventListener('mousemove', function(e) {\n",
        "      mouse.x = e.pageX - this.offsetLeft\n",
        "      mouse.y = e.pageY - this.offsetTop\n",
        "    })\n",
        "    canvas.onmousedown = ()=>{\n",
        "      ctx.beginPath()\n",
        "      ctx.moveTo(mouse.x, mouse.y)\n",
        "      canvas.addEventListener('mousemove', onPaint)\n",
        "    }\n",
        "    var onPaint = ()=>{\n",
        "      ctx.lineTo(mouse.x, mouse.y)\n",
        "      ctx.stroke()\n",
        "    }\n",
        "    var data = new Promise(resolve=>{\n",
        "      canvas.onmouseup = ()=>{\n",
        "        canvas.removeEventListener('mousemove', onPaint)\n",
        "        resolve(mouse)\n",
        "      }\n",
        "    })\n",
        "    function draw_collage_image() {\n",
        "      collage_image = new Image();\n",
        "      collage_image.src = filename_image;\n",
        "      collage_image.onload = function(){\n",
        "        ctx.drawImage(collage_image, 0, 0);\n",
        "      }\n",
        "    }\n",
        "    draw_collage_image();\n",
        "  </script>\n",
        "  \"\"\"\n",
        "  \n",
        "  im = IPython.display.Image(filename_temp, embed=True)\n",
        "  # IPython.display.display(im)\n",
        "  filename_embed = 'data:image/png;base64,'\n",
        "  filename_embed += base64.b64encode(im.data).decode('ascii')\n",
        "  \n",
        "  # Display an HTML canvas with the image.\n",
        "  canvas = IPython.display.HTML(\n",
        "      canvas_html % (CANVAS_WIDTH * POP_SIZE, CANVAS_HEIGHT, filename_embed))\n",
        "  print('Click with the mouse on the desired image and patch:')\n",
        "  IPython.display.display(canvas)\n",
        "  \n",
        "  # Select the image and pixel coordinates.\n",
        "  def draw():\n",
        "    print('draw()')\n",
        "    mouse = eval_js('data')\n",
        "    return mouse\n",
        "  mouse = draw()\n",
        "  pop_id_mouse = int(np.floor(mouse['x'] / CANVAS_WIDTH))\n",
        "  x_mouse = int(mouse['x'] % CANVAS_WIDTH)\n",
        "  y_mouse = int(mouse['y'])\n",
        "  print(f'Selected image {pop_id_mouse} at ({x_mouse}, {y_mouse})')\n",
        "  \n",
        "  def find_patch(generator, id, u, v):\n",
        "    # Render only the spatial transforms of the patches.\n",
        "    rendered_patches = generator.spatial_transformer(generator.patches)\n",
        "    rendered_patches = rendered_patches.detach().cpu().numpy()\n",
        "    patch_id = np.argmax(rendered_patches[id, :, 3, u, v] * rendered_patches[id, :, 4, u, v])\n",
        "    return patch_id\n",
        "  \n",
        "  # Select the patch.\n",
        "  patch_id = find_patch(generator, pop_id_mouse, y_mouse, x_mouse)\n",
        "  print(f'Found matching patch {patch_id}')\n",
        "  \n",
        "  # Extract the patch's current affine transform paramaters.\n",
        "  with torch.no_grad():\n",
        "    x0 = generator.spatial_transformer.translation[pop_id_mouse, patch_id, 0, 0]\n",
        "    x0 = float(x0.detach().cpu().numpy())\n",
        "    y0 = generator.spatial_transformer.translation[pop_id_mouse, patch_id, 1, 0]\n",
        "    y0 = float(y0.detach().cpu().numpy())\n",
        "    rot0 = generator.spatial_transformer.rotation[pop_id_mouse, patch_id, 0, 0]\n",
        "    rot0 = float(rot0.detach().cpu().numpy())\n",
        "    scale0 = generator.spatial_transformer.scale[pop_id_mouse, patch_id, 0, 0]\n",
        "    scale0 = float(scale0.detach().cpu().numpy())\n",
        "    squeeze0 = generator.spatial_transformer.squeeze[pop_id_mouse, patch_id, 0, 0]\n",
        "    squeeze0 = float(squeeze0.detach().cpu().numpy())\n",
        "    shear0 = generator.spatial_transformer.shear[pop_id_mouse, patch_id, 0, 0]\n",
        "    shear0 = float(shear0.detach().cpu().numpy())\n",
        "    patch_info = {'pop_id': pop_id_mouse, 'patch_id': patch_id,\n",
        "                  'x0': x0, 'y0': y0, 'rot0': rot0,\n",
        "                  'scale0': scale0, 'squeeze0': squeeze0, 'shear0': shear0,\n",
        "                  'x': x0, 'y': y0, 'rot': rot0,\n",
        "                  'scale': scale0, 'squeeze': squeeze0, 'shear': shear0}\n",
        "  \n",
        "  def show_modified(dx, dy, drot, dscale, dsqueeze, dshear):\n",
        "    \"\"\"Visualization callback function with affine transform deltas.\"\"\"\n",
        "    with torch.no_grad():\n",
        "      x = patch_info['x0'] - dx\n",
        "      y = patch_info['y0'] + dy\n",
        "      rot = patch_info['rot0'] - drot\n",
        "      scale = patch_info['scale0'] - dscale\n",
        "      squeeze = patch_info['squeeze0'] + dsqueeze\n",
        "      shear = patch_info['shear0'] + dshear\n",
        "      generator.spatial_transformer.translation[pop_id_mouse, patch_id, 0, 0] = x\n",
        "      generator.spatial_transformer.translation[pop_id_mouse, patch_id, 1, 0] = y\n",
        "      generator.spatial_transformer.rotation[pop_id_mouse, patch_id, 0, 0] = rot\n",
        "      generator.spatial_transformer.scale[pop_id_mouse, patch_id, 0, 0] = scale\n",
        "      generator.spatial_transformer.squeeze[pop_id_mouse, patch_id, 0, 0] = squeeze\n",
        "      generator.spatial_transformer.shear[pop_id_mouse, patch_id, 0, 0] = shear\n",
        "    patch_info['x'] = x\n",
        "    patch_info['y'] = y\n",
        "    patch_info['rot'] = rot\n",
        "    patch_info['shear'] = shear\n",
        "    patch_info['squeeze'] = squeeze\n",
        "    patch_info['shear'] = shear\n",
        "    params = {'gamma': step / OPTIM_STEPS}\n",
        "    img = generator(params)\n",
        "    img = img.permute(0, 3, 1, 2)  # NHWC -> NCHW\n",
        "    _ = show_stitched_batch(img,\n",
        "                            max_display=MAX_MULTIPLE_VISUALISATIONS)\n",
        "  \n",
        "  # Interactive editing of the patch's affine transform parameters.\n",
        "  interactive_plot = interactive(show_modified,\n",
        "                                dx=(-MAX_TRANS * 2, MAX_TRANS * 2, 0.01),\n",
        "                                dy=(-MAX_TRANS * 2, MAX_TRANS * 2, 0.01),\n",
        "                                drot=(-MAX_ROT * 2, MAX_ROT * 2, 0.01),\n",
        "                                dscale=(-MAX_SCALE * 2, MAX_SCALE * 2, 0.01),\n",
        "                                dsqueeze=(-MAX_SQUEEZE * 2, MAX_SQUEEZE * 2, 0.01),\n",
        "                                dshear=(-MAX_SHEAR * 2, MAX_SHEAR * 2, 0.01))\n",
        "  output = interactive_plot.children[-1]\n",
        "  output.layout.height = '350px'\n",
        "else:\n",
        "  interactive_plot = \"Patch tinkering not enabled.\"\n",
        "interactive_plot\n",
        "  "
      ],
      "metadata": {
        "cellView": "form",
        "id": "OZpKu3XcdogG"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Render high res image and finish up.\n",
        "\n",
        "ct.assemble_tiles()"
      ],
      "metadata": {
        "cellView": "form",
        "id": "F4ley6lsEtPO"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Save and download assets\n",
        "\n",
        "#@markdown Enable this to allow everything to be zipped up and downloaded\n",
        "DOWNLOAD_FILES = True #@param {type:\"boolean\"}\n",
        "\n",
        "if DOWNLOAD_FILES:\n",
        "  zipname = f\"{config['output_dir'].rstrip('/')}.zip\"\n",
        "  print(f\"Output {config['output_dir']} will be downladed as {zipname}\")\n",
        "  !zip -r {zipname} {ct._output_dir}\n",
        "  from google.colab import files\n",
        "  files.download(zipname)\n"
      ],
      "metadata": {
        "id": "YVst6OhPVx_s",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Copy model checkpoint to use as initialisation for next collage\n",
        "\n",
        "#@markdown Enable this to copy file `generator.pt` from the current output directory. This model file contains the learned spatial and colour transforms from the current collage. If `USE_INIT_CHECKPOINT` is set to `True`, then next collage generation will try to use that model as initialisation.\n",
        "COPY_CHECKPOINT_FOR_INITIALISATION = True #@param {type:\"boolean\"}\n",
        "\n",
        "if COPY_CHECKPOINT_FOR_INITIALISATION:\n",
        "  !cp \"{ct._output_dir}/generator.pt\" .\n",
        "  print(f\"Copied {ct._output_dir}/generator.py to root directory.\")"
      ],
      "metadata": {
        "cellView": "form",
        "id": "3zBG6s1WiwEB"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "G264zsJ0XV-3"
      },
      "outputs": [],
      "source": [
        "raise ValueError(\"Stop here.\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9biyGTC-6DOy"
      },
      "source": [
        "#Patch file Creation\n",
        "\n",
        "Patch files can be created from PNGs (best) and JPGS. Both methods produce a `.npy` a file that contains all the patches. They can be uploaded and used in this Colab by\n",
        "\n",
        "* Set `Configure -> Patch Settings -> Select a patch set -> EXAMPLE_PATCH_SET` to `None of the above`\n",
        "* Set `Advanced Settings -> Patch Settings -> Select Segmented Patches -> Load patch sets from elsewhere -> ADVANCED_PATCH_SET` to `Upload to Colab`\n",
        "\n",
        "* PNGs reqiure the subject to have been manually segmented and on a transparent background. The cells below can be used with such PNGs.\n",
        "\n",
        "* JPGs use the [Arnheim 3 Patch Maker](https://colab.research.google.com/github/deepmind/arnheim/blob/main/arnheim_3_patch_maker.ipynb) Colab which uses a segmentation algorithm to attempt to cut out the subject. This can be hit or miss. Instructions are in the Colab.\n",
        "\n",
        "## Note:\n",
        "Make sure you run the `Preliminaries` calls at the top of the Colab before running these."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SjYI9mBDHSxm",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "#@title Functions to create patch files.\n",
        "\n",
        "def upload_files(image_path):\n",
        "  \"\"\"Upload files to target directory.\"\"\"\n",
        "  print(f\"Uploading files to {image_path}\")\n",
        "  !rm -rf {image_path}\n",
        "  !mkdir {image_path}\n",
        "  uploaded = files.upload()\n",
        "  for k, v in uploaded.items():\n",
        "    open(image_path + \"/\" + k, 'wb').write(v)\n",
        "  return list(uploaded.keys())\n",
        "\n",
        "\n",
        "def convert_pngs(image_path, destination_file):\n",
        "  \"\"\"Convert pngs to single numpy array file.\"\"\"\n",
        "  png_imgs = []\n",
        "  for png_im_path in glob.glob(image_path + \"/*.png\"):\n",
        "    png_im = imageio.imread(png_im_path)\n",
        "    png_imgs.append(png_im)\n",
        "  \n",
        "  np.save(destination_file, np.array(png_imgs), allow_pickle=True)\n",
        "\n",
        "\n",
        "def show_patches(filename):\n",
        "  patches = np.load(filename, allow_pickle=True)\n",
        "  for patch in patches:\n",
        "    print(patch.shape)\n",
        "    cv2_imshow(patch)"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Create patch file from PNGs\n",
        "#@markdown Requires PNG files with alpha channel.\n",
        "\n",
        "OUTPUT_PATCH_FILE_NAME = \"my_patches.npy\"  # @param{type:\"string\"}\n",
        "target_patch_path = f\"/content/{OUTPUT_PATCH_FILE_NAME}\"\n",
        "png_path = \"/content/pngs\"\n",
        "SHOW_PATCHES = True  #@param{type: \"boolean\"}\n",
        "\n",
        "print(\"Select PNG files to be converted:\")\n",
        "upload_files(png_path)\n",
        "print(\"Converting images.\")\n",
        "convert_pngs(png_path, target_patch_path)\n",
        "if SHOW_PATCHES:\n",
        "  show_patches(target_patch_path)\n",
        "files.download(target_patch_path)"
      ],
      "metadata": {
        "cellView": "form",
        "id": "AsCWlm82vSUp"
      },
      "execution_count": null,
      "outputs": []
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [
        "aGK__cS0HJD6",
        "XlVNTb_1NelG",
        "wG5OSH6zG-Or",
        "YbGkH7v-_--H",
        "g35H20mf61r5",
        "ZHAkxmLIPC7v"
      ],
      "machine_shape": "hm",
      "private_outputs": true,
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}