notebooks/MorphingEmoji.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "d7d6c646-8608-4d2b-ad90-8ea4b3745f4c",
   "metadata": {},
   "source": [
    "# Morphing Emojis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "a3e435ff-e7ec-47f2-9476-679e5699c4f5",
   "metadata": {},
   "outputs": [],
   "source": [
    "from controllable_nca.experiments.morphing_image.trainer import MorphingImageNCATrainer\n",
    "from controllable_nca.experiments.morphing_image.emoji_dataset import EmojiDataset\n",
    "from controllable_nca.nca import ControllableNCA\n",
    "\n",
    "import torch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "c90c54c0-ac1d-4641-b9d9-ba1d36586434",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import torch\n",
    "from einops import rearrange\n",
    "\n",
    "from controllable_nca.dataset import NCADataset\n",
    "from controllable_nca.utils import load_emoji, rgb\n",
    "\n",
    "\n",
    "class EmojiDataset(NCADataset):\n",
    "    # EMOJI = '🦎😀💥'\n",
    "    EMOJI = \"🦎😀💥👁🐠🦋🐞🕸🥨🎄\"\n",
    "    # EMOJI = \"🦎😀👁🕸🥨🎄\"\n",
    "\n",
    "    digits = [\n",
    "        \"0030\",  # 0\n",
    "        \"0031\",  # 1\n",
    "        \"0032\",  # 2\n",
    "        \"0033\",  # 3\n",
    "        \"0034\",  # 4\n",
    "        \"0035\",  # 5\n",
    "        \"0036\",  # 6\n",
    "        \"0037\",  # 7\n",
    "        \"0038\",  # 8\n",
    "        \"0039\",  # 9\n",
    "    ]\n",
    "\n",
    "    def __init__(self, image_size=64, thumbnail_size=32, use_one_hot: bool = False):\n",
    "        emojis = torch.stack(\n",
    "            [load_emoji(e, image_size, thumbnail_size) for e in EmojiDataset.EMOJI],\n",
    "            dim=0,\n",
    "        )\n",
    "        self.emojis = emojis\n",
    "        self.num_samples = len(self)\n",
    "        self._target_size = self.emojis.size()[-3:]\n",
    "\n",
    "    def num_goals(self):\n",
    "        return self.emojis.size(0)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        if isinstance(idx, int):\n",
    "            return self.emojis[idx : idx + 1].clone(), idx\n",
    "        return self.emojis[idx].clone(), idx\n",
    "\n",
    "    def __len__(self):\n",
    "        return self.emojis.size(0)\n",
    "\n",
    "    def target_size(self):\n",
    "        if self._target_size is not None:\n",
    "            return self._target_size\n",
    "        self._target_size = self.emojis.size()[-3:]\n",
    "        return self._target_size\n",
    "\n",
    "    def to(self, device: torch.device):\n",
    "        self.emojis = self.emojis.to(device)\n",
    "\n",
    "    def visualize(self, idx=0):\n",
    "        self.plot_img(self.emojis[idx : idx + 1])\n",
    "\n",
    "    def plot_img(self, img):\n",
    "        with torch.no_grad():\n",
    "            rgb_image = rgb(img, False).squeeze().detach().cpu().numpy()\n",
    "        rgb_image = rearrange(rgb_image, \"c w h -> w h c\")\n",
    "        _ = plt.imshow(rgb_image)\n",
    "        plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "d588342d-77bb-4c2d-939c-9c9f6c127eee",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = EmojiDataset(image_size=64, thumbnail_size=40)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "5ef4e59a-658a-4781-9bf4-ce7f485eed82",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAPsAAAD7CAYAAACscuKmAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjUuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/YYfK9AAAACXBIWXMAAAsTAAALEwEAmpwYAAAdjUlEQVR4nO2dfXxV5ZXvf+skJ4EEJIEcwpu8WLkovoDeVKhoi1gd1FYdP+q19U65c7lyb69t7avV9mOv0zvt1U/rVGemV0utlplxilq1WG4VEaVWbYFQAXkRCBiEmEAIBBMS8nbW/eMc9rPXmYQcknNOMnl+388nn6y919p7r/Oyzn6e/axnPaKqIIQMfSID7QAhJDcw2AnxBAY7IZ7AYCfEExjshHgCg50QT+hXsIvIQhHZKSJVInJPppwihGQe6es4u4jkAdgF4CoABwBsAPA5Vd2eOfcIIZkivx/HXgKgSlX3AoCILAdwA4Aeg72srEynTp3aj0sSQk5FdXU1Dh8+LN3p+hPsEwHsD20fADDnVAdMnToVlZWV/bgkIeRUVFRU9KjL+gM6EVkiIpUiUllfX5/tyxFCeqA/wV4D4MzQ9qTkPoOqLlXVClWtiMVi/bgcIaQ/9CfYNwCYLiLTRKQAwG0AXsyMW4SQTNPnPruqdorIlwCsApAH4AlV3ZYxzwghGaU/D+igqr8D8LsM+UIIySLMoCPEExjshHgCg50QT2CwE+IJDHZCPIHBTognMNgJ8QQGOyGewGAnxBMY7IR4AoOdEE9gsBPiCQx2QjyBwU6IJzDYCfEEBjshnsBgJ8QTGOyEeAKDnRBPYLAT4gkMdkI8gcFOiCcw2AnxBAY7IZ7AYCfEE3oNdhF5QkQOicjW0L7RIrJaRHYn/5dm101CSH9J587+SwALU/bdA2CNqk4HsCa5TQgZxPQa7Kr6BoAjKbtvALAsKS8DcGNm3SKEZJq+9tnLVbU2KdcBKM+QP4SQLNHvB3SqqgC0J72ILBGRShGprK+v7+/lCCF9pK/BflBExgNA8v+hngxVdamqVqhqRSwW6+PlCCH9pa/B/iKARUl5EYAVmXGHEJIt0hl6+xWAPwKYISIHRGQxgAcAXCUiuwF8OrlNCBnE5PdmoKqf60F1ZYZ9IYRkkV6DnQxV4iGZiZQ+wE+ZEE9gsBPiCWzGD2FUuwJZJM/o4vUvu43jvzW6yNgvO7u2UG5E/fP2AlO+444pTM2rCqdeSFr+kuzCOzshnsBgJ8QTGOyEeAL77EMY21NOmb4Q6sNH2l4xqvjBZqcrbAnkrobdxi5v+sM9X5xd9kEH7+yEeAKDnRBPYDN+SCM9yIDCNdW7IlPtYUd2BmJcnV08r8CYRY4HlcogIy4wuri6dnxE2I4fDPDOTognMNgJ8QQ244caGprgIu63PN5Uaczy4k+5jbY2o2tv+iCQI+qa4NHRNguvfdsdgVzw8deMLhIZHnYqJLNJP1Dwzk6IJzDYCfEEBjshnsA++xAjHuoei4SGv/Jtsc945PZAXlVnq4pdM9zdA9rzooH8Wuc1xm7BrHsDuTPlq5TPvvmgg3d2QjyBwU6IJ7AZP8SIRNzvd0e8M5Cjw6cYu+cP1gTyX52YYHTvRxsCuarDNeNvrm8ydi9NORHIc8psdl1n6Nr5EX7NBgO8sxPiCQx2QjyBwU6IJ7AzNcQI95Wjob7ypiO28MTiLW647UQ8bnSX1bphuqK467MfHfWRsfvs+vsC+Q/zHjK6GaPO7NYn9t8HjnSWfzpTRF4Xke0isk1E7kruHy0iq0Vkd/J/afbdJYT0lXSa8Z0AvqGqMwHMBXCniMwEcA+ANao6HcCa5DYhZJCSzlpvtQBqk3KTiOwAMBHADQDmJ82WAVgL4NtZ8ZIYNDSLrDOlCR5uuu9p+jCQ/3L9D4xdcyi7rvB4i9HdWXpbIF+zsyOQV299y9h9fV6rO/+67xvdq/Pc9SYUjw3522ns2KzPHaf1gE5EpgK4CMA6AOXJHwIAqAOQukoAIWQQkXawi8gIAM8B+Kqqmic1qqr4N+VLg+OWiEiliFTW19d3Z0IIyQFpBbuIRJEI9KdU9eQaQAdFZHxSPx7Aoe6OVdWlqlqhqhWxWKw7E0JIDui1wyQiAuAXAHao6t+FVC8CWATggeT/FVnxkJg+OgBIaEZZNGKrx/yxYXsg37beDa/VdBwxdpEutw7ceW3FRve16EVu46BbE+6c6hHG7vHpLn12S3mt0c1747uBvPzj3wjkOWXngAwM6TwdmQfgrwC8KyKbkvu+g0SQPyMiiwHsA3BrVjwkhGSEdJ7Gv4meC4ddmVl3CCHZguMeg5Rw011Sfms/6jgeyH+/xy63/ONdvwrkdWO3BfL9R8Ybu+WtowL50DA7fFe//LFAjh1zGXQ1Y22Xob7Y+Tj6RIfRXVLsuhOf/dPdgfzNGYuM3den3xjI+bDnP9V7QE4f5sYT4gkMdkI8gc34QURXqOZ7JNRs/cMBW/P985t+HMg1ea1GF21vD+SHDxQFctXhBmM3N3IskCvLJxrdxZe7jLpPHS0M5E1Fjcauttk90Z852qjw2Ij9gTynwflxz+afGbvfv/18ID9/+xNGV5DvCmJoaDkp4XJSfYJ3dkI8gcFOiCcw2AnxBPbZBxHSw5JoIyKFxq59165AHhax67RFQrPgflbvhsPmlbYbuzWXuOPOrTnD6ApCxSNfLHPDfE223iS+1rInkP+wf5jRnbXbnbMr4opbFp04YeyOjHLDbXG1Q4BcIi6z8M5OiCcw2AnxBDbjBxHhmu/hoaaLJl5g7Db89YuB/Ie9tqBEXmiJ5ac2PxfIr9T/ydhNrnRN9YOtdUZXdalr4i9Tl3n3tyfGGbsvnuvOcWCDbePXHHNVyh647v5APu9MOxHm3LFnB3Jh1HZXwjCDrv/wzk6IJzDYCfEEBjshnsA++yAlHhp3ykvpr04ePSmQrx55rdHdv21ZIL/a6obo5u63hScmxF2fem3MFra4cr0bpjsec9fOG2mHzWa87dJg/6LQFpJcev1ZgTy2yKXEtg7/prErKihCOjBdtv/wzk6IJzDYCfEENuMHKXnifod3HvvA6P5xrytY8WT1K0bXMsploV1zzA2VrbjN1o3Xr7n6dKuajhrdjd90yztHGtxsubw6u/xTSdQ1/5+4/V6jG182y20cchXL7t1rX8u3Ns0N5DvPXmj9GH9ZII8INfdPVZOP9Azv7IR4AoOdEE9gM34Aiacs3SQR1xx9ZI/Lkvvern8xdscjboJLpMvWfpPjrok7Z3/oSfqPnjR2u3duDuRZEjW68qgrLf3hTDehRVqtv1rgsuvyW1YbHT74uZPjzqeWZnuOPx17P5Df2fmPRvfj1f8QyC/f8rjzb9RYY8cn9enBOzshnsBgJ8QTGOyEeAL77Dkm3E8Pz3IDgJpGt8TyPavcEshaYotL5He42WZPjdtndEcKXTGI740vC+T5v15l7M6Ju3P886dsFttBV1IeeaHEuEiB7Q83xF1W3heqDhrdl2LuwD3t0wJ5W0uBsbuu0A37vfShHVKrqj0QyE2hWvnjUobawkUv2GfvmV7v7CIyTETWi8hmEdkmIn+T3D9NRNaJSJWIPC0iBb2dixAycKTTjG8DsEBVZwGYDWChiMwF8CCAn6jq2QCOAlicNS8JIf0mnbXeFEBzcjOa/FMACwB8Prl/GYD7ATyaeReHFuFmZmomWNmIMYF8/QiXgfbS+68bu7Y2d9zKZnuO1tBIXNkoN3HlvgdKjF1Ni5sIs2eqrR8Xbgjn1zcGctchW3u+sNgd9/LIUqNb1eD6Ahoqbf/wmP3GriLPZeX9vmaM0f1wwbcDeXqZm1iTOmSZ2h0i3ZPu+ux5yRVcDwFYDWAPgEZVPdkxOwBgYg+HE0IGAWkFu6p2qepsAJMAXAIg7UW2RWSJiFSKSGV9fX3fvCSE9JvTav+oaiOA1wF8AkCJiJzsBkwCUNPDMUtVtUJVK2KxWH98JYT0g1777CISA9Chqo0iMhzAVUg8nHsdwM0AlgNYBGBFNh0dMoQ6xF3xLqMqyHcFF5/5glv3bEvtNmO3vmZTIH/x2buNrjPUaf/dfHexLrFrwn12uxvKOiM81gag9bgrOFnQ4IbGxhbY5wP7alyRySvGHDK6F+Y629lvuq/ZV/fb9N7/dMENgbz/Kw8YXcnIkkCOw/XTlaNrfSKdcfbxAJaJSB4SLYFnVHWliGwHsFxE/hbAOwB+kUU/CSH9JJ2n8VsAXNTN/r1I9N8JIf8OYAZdjgkXWsiP9Pz2151whSHeO1ZtdP+68dlA/s8zLjW6scNd1tyP9rla8e822Gb23GK3LPOqGY1Gd9Uf3Sy4jiLn75/n26WmLq10mX1bW/OM7qt73LBf7LgbortJpxq749vcrLevTbWNwzumuGIWl5ad5xQpzfjwrLewDHBYLgzfCUI8gcFOiCewGZ8FTlVM4aPQhI6XayuN7oW6twN5VYN7An90m7X7H+UVgfzo7cuMDgfeCcR3P5gdyHNff8SY7W9vDuQHq+zXYG+r8/nYGcMD+cajk43djomuyEWj2HP8c517Ur/7KjdiMO1EsbHDE67IxZLx641qXs2rgXztaFer7q6PXWfsrh7v3g9OhOkZ3tkJ8QQGOyGewGAnxBPYZ88AqbPXwnzUcsxsf+qlLwXyu6OOG138hCv4cHHEZbxdHbN13c8fN9Nt1O0wus4GN5R1QambmzTuwvON3d4jtYH8w2abyXfFBe4e8KVzzwzkbY27jN0bda54hUTse3BJvuubT/tzVSC37LfnCI3sYX6NHb77uatniQPNbmnq69/ZaK9VNSOQfzD5FqO7fIpLBTlV4RAf8O8VE+IpDHZCPIHN+CwQHv5p62o3ur07XRO0YHK50cWHlQTyhEKXgfbpKcON3UPNuwN5Qd17Rlde6JrPzzY63QdddukmlLpCEaUltn7cP827LZAnFTmfbir8rbErDqWyfeOAfS07SlyzftPG3wfy7GN2NdnjRa6a2UNxW0+v/KjL8ls3zWUAzq8eb+zerHG16n66z3Z5TDM+VKsu4uF9zr9XTIinMNgJ8QQGOyGewD57BkhdMjg8xBMbaavzfP/yuwL5u2ttsQYUueIVL4W6+r/tsAUhiyduDeRbxRaBPNrlZqJtj4eG/T6yw3z5xw4Hct5kW7xi9EfPuI1jrpY9xPa3y/Odk5E8O2x2LOreg48vcEN7lx9sMXZjS1xa7RdH7TG6n4RG6WbVu69qfeuHxm5a1L3H37r2y0YXTl2OiN/3Nr9fPSEewWAnxBPYjM8C4eys1Brnd13+3wP53HG2SO9vtq4M5Fe2vhzIe4rtksoPfsyd8wsFNptsdI3LrpPDrolcXFtn7DpCTdrD4+wSyPfsd4Uz7g4t43Qobpv7/+ugaz7Hj9pMwWihO/+1ZW6GXe0oe3/ZrG77oSLbFWia7bLhqppdV+a8CWcZuxvOdUUuJp4xzuhMM97DrLkwfr96QjyCwU6IJ7AZnwFS657FQxNjTlVK4erpnzLbk8pccYjVVS7rLNrRbOyePDAykN9Uu8JrZ52b4HLfJJc1d8dZNpPvgjfdk/+2Y41G9w/5kwL5scYJ7tyd9nVe3e6uNe2orXH3amjzvk86P7oihcZuzhaXDTe9y64N+oMrzg7kn161BOmQ2m0Kl6CWUA3qvIjtMvgA7+yEeAKDnRBPYLAT4gnss2eA1CKHeeGeekqnfVuoBvxj1S8Z3c+3P+fO0eoy3ApbbH97Y6ubEbcjarPabh7+sUAeGXF95X9q/cDYtXS4ZZg+12Uz0h6b7Jbtm7PT9am3VzcZu3+pcB3zd8ttIcn1W93w2Mo2Vzhjx9FqYxc9tiGQI1HbZ7/vDZdh+LPNywP59su+YOxuOWtBIJ8z0hbF9HF2W0+k/U4kl21+R0RWJrenicg6EakSkadFpKC3cxBCBo7T+dm7C0C4BtKDAH6iqmcDOApgcSYdI4RklrSa8SIyCcB1AH4A4OuSaLcuAPD5pMkyAPcDeDQLPg5KwnXnTnTZlUmfPvBGIC/f/7rRra3fEsgp81uAAtc46po2NZBTh++icdesXzH9K0b36YaSQG5a9mQgL/y4zcLriLjJNDXH7TDU8++7iSstu11hiIIT1pPLD7qm+rGRRUb3+LX/O5CvibkCElt3bjF2e0t/GsgbW219uoI8d86q4e6+dH/Nc8buh3vc9gUlNrvu8pjrQtwy2tWXnzvRLl+oZrh0aNaeT/fO/jCAu4Fg0HIMgEZVPZlLeQDAxG6OI4QMEnoNdhH5DIBDqrqxN9sejl8iIpUiUllfX9+XUxBCMkA6d/Z5AK4XkWoAy5Fovj8CoEQkWPNnEoCa7g5W1aWqWqGqFbFYrDsTQkgOSGd99nsB3AsAIjIfwDdV9XYReRbAzUj8ACwCsCJ7bg4OOrvcDLD8PPfW/WbLSmO3pOrv3TEjSoxOQn3DUGl4AEA8NIQXP6M4tN/ajal2hSgufcsWfMCHbjhsZGj9tYVxOzT2dsyl3P6pzha2WHvI9eGjxe45gvwHW1TyveJQqm6nvW9MbnNpsfv2uee6Y/Jtuuz4YjfjLt5ia+B3hYp5hBN1I+32DWmLuJTYjc32/di83w3t/XyPG8rb+a23jd34ke61pabcDpXZcv15Fd9G4mFdFRJ9+F/0Yk8IGUBOK6lGVdcCWJuU9wK45FT2hJDBAzPoToOelgM+a8wUs13wm+pAjoyys9IiJSWBHM+3b78ZADvR5vY3HDF2jWXunA8d+aPR/c8Gl1FXN9b5+9SZncYukudms+V12Lrxw2KlgdxR7Ia/4lE7fFf8/v5Abi0vMbr/W7sqkP/bSDe7b1eXnR331nHXdC+otn5E810N+K6RoUzBiP0cou1u6FOabY07ibhhyoqJFwdy6TBbiMMMvQ3RZZ+HRmeEENIrDHZCPEFSCy9kk4qKCq2srMzZ9bLJqVYE/WWlm7TxvVf/j9EdaQ/VastLbS6Gtk+4pmlHzNaI65jq8pc0brP3ytpc16Ap33227Xn2CfPwXa4J3jE6pQDGuDLnUag7oSnN+Gity5vIa7JP9Fs/5nw8Q1xXoLkjZeXaqHvvCg422vPvdz6iMNTlSf3OhjaH59mlsm48/7pA/uHV3wnkWPEYYzdUMugqKipQWVnZ7QvgnZ0QT2CwE+IJDHZCPIFDb30k3E9X2D7kf6lwSx4vnLHA6F57/61A3lq73eg6Qhl6LSdc3/bxD18zdtE2Zycp3dejcVfYAg2uUOXwY7auO1qdTprtrLe8lpGhrVD3r6XN2ElHaPmnJnv+YR+4r9bxiS47TVIKVBSGMhE10mV0d1yyKJBnTbrQud5hUw8nl7oCmeeX21r8U0omoTtSn1UN1eG2MLyzE+IJDHZCPIHN+AyQOlTTFXfN0XEj7bDZ5y/8S7cRllM40uKyx5750YtG19Lqasul9CAQ7Qw1hVvcsNzsiRcauyvOnxfIj771hNEdbn4vkEVcE186bRZeKdww18yJ/9HoNuxzQ6zDOl2zu6vYFrlA6Jx5R2yNu/96/e2BfMG4megL4SFSCWXe+dBsT4V3dkI8gcFOiCcw2AnxBPbZs0B4HbHUYbmurq5U825tRxe5mWev3mELLP6/XasDuaXVzvKaVOrSVCsmzQ7Js4xde52rS9/2/JtG917UzbI7EXF96lK1qah3f+XHgXzxDDvbedOHrqDl6t1rA/ndmm3GbsRwN5vtpr/4jNGF++kdobTg8JptQEpfPEU3VApPZAK+E4R4AoOdEE9gMz7LpA7LhWvXnYpwhtesCecZXep2X1hXuTaQJ6TUpxvb7Jrr4USzwqjNtCtqaA/Z2e7K7AnndyufDuFzRiPRU1iG8G9ELW14ZyfEExjshHgCm/GDlHCGV2pp47jGU817P1/K9u6N6wL5BOz52kNP4CU84ce24vHuRvcU/5xLLzO6rtCknnADX1Lqx2m859pv4VEN0n94ZyfEExjshHgCg50QT2Cf/d8BqVlgkVP8RvdUCLOp0daer93rlkcellJIUiXUTw/1qQsL7NflSO2HPfph+tunmmHGbnnOSHd99moATQC6AHSqaoWIjAbwNICpAKoB3KqqR3s6ByFkYDmdZvwVqjpbVU+uaH8PgDWqOh3AmuQ2IWSQ0p9m/A0A5iflZUisAfftfvpD+ompjRfKQCsaYWvDX3GLq++2Z/MGo2tu+gjdUTqmzGzPufamHv2Ih64d8bBQxGAk3Tu7AnhFRDaKyJLkvnJVrU3KdQDKuz+UEDIYSPfOfpmq1ojIWACrReS9sFJVVSS1zmmC5I/DEgCYPHlyv5wlhPSdtO7sqlqT/H8IwAtILNV8UETGA0Dy/6Eejl2qqhWqWhGLxTLjNSHktOn1zi4ixQAiqtqUlK8G8H0ALwJYBOCB5P8V2XSUnD7h9NO8lOWhL/3MTd3KANDZ0f1stmhBYdrXZtGIwUc6zfhyAC8kvzj5AP5VVV8WkQ0AnhGRxQD2Abg1e24SQvpLr8GuqnsBzOpmfwOAK7PhFCEk8zCDjvwb8lOWaCJDA3asCPEEBjshnsBgJ8QTGOyEeAKDnRBPYLAT4gkMdkI8gcFOiCcw2AnxBAY7IZ7AYCfEExjshHgCg50QT2CwE+IJDHZCPIHBTognMNgJ8QQGOyGewGAnxBMY7IR4AoOdEE9gsBPiCQx2QjyBwU6IJ6QV7CJSIiK/FpH3RGSHiHxCREaLyGoR2Z38X5ptZwkhfSfdO/sjAF5W1XOQWApqB4B7AKxR1ekA1iS3CSGDlF6DXURGAfgkgF8AgKq2q2ojgBsALEuaLQNwY3ZcJIRkgnTu7NMA1AN4UkTeEZHHk0s3l6tqbdKmDonVXgkhg5R0gj0fwMUAHlXViwAcR0qTXROLeGs3x0JElohIpYhU1tfX99dfQkgfSSfYDwA4oKrrktu/RiL4D4rIeABI/j/U3cGqulRVK1S1IhaLZcJnQkgf6DXYVbUOwH4RmZHcdSWA7QBeBLAouW8RgBVZ8ZAQkhHSXZ/9ywCeEpECAHsB/DUSPxTPiMhiAPsA3JodFwkhmSCtYFfVTQAqulFdmVFvCCFZgxl0hHgCg50QT2CwE+IJDHZCPIHBTognMNgJ8QQGOyGeIIm09hxdTKQeiQScMgCHc3bh7hkMPgD0IxX6YTldP6aoard56TkN9uCiIpWq2l2Sjlc+0A/6kUs/2IwnxBMY7IR4wkAF+9IBum6YweADQD9SoR+WjPkxIH12QkjuYTOeEE/IabCLyEIR2SkiVSKSs2q0IvKEiBwSka2hfTkvhS0iZ4rI6yKyXUS2ichdA+GLiAwTkfUisjnpx98k908TkXXJz+fpZP2CrCMiecn6hisHyg8RqRaRd0Vkk4hUJvcNxHcka2XbcxbsIpIH4KcArgEwE8DnRGRmji7/SwALU/YNRCnsTgDfUNWZAOYCuDP5HuTalzYAC1R1FoDZABaKyFwADwL4iaqeDeAogMVZ9uMkdyFRnvwkA+XHFao6OzTUNRDfkeyVbVfVnPwB+ASAVaHtewHcm8PrTwWwNbS9E8D4pDwewM5c+RLyYQWAqwbSFwBFAP4MYA4SyRv53X1eWbz+pOQXeAGAlQBkgPyoBlCWsi+nnwuAUQDeR/JZWqb9yGUzfiKA/aHtA8l9A8WAlsIWkakALgKwbiB8STadNyFRKHQ1gD0AGlW1M2mSq8/nYQB3A4gnt8cMkB8K4BUR2SgiS5L7cv25ZLVsOx/Q4dSlsLOBiIwA8ByAr6rqRwPhi6p2qepsJO6slwA4J9vXTEVEPgPgkKpuzPW1u+EyVb0YiW7mnSLyybAyR59Lv8q290Yug70GwJmh7UnJfQNFWqWwM42IRJEI9KdU9fmB9AUANLG6z+tINJdLRORkXcJcfD7zAFwvItUAliPRlH9kAPyAqtYk/x8C8AISP4C5/lz6Vba9N3IZ7BsATE8+aS0AcBsS5agHipyXwhYRQWIZrR2q+ncD5YuIxESkJCkPR+K5wQ4kgv7mXPmhqveq6iRVnYrE9+E1Vb09136ISLGIjDwpA7gawFbk+HPRbJdtz/aDj5QHDdcC2IVE//C7ObzurwDUAuhA4tdzMRJ9wzUAdgN4FcDoHPhxGRJNsC0ANiX/rs21LwAuBPBO0o+tAL6X3H8WgPUAqgA8C6Awh5/RfAArB8KP5PU2J/+2nfxuDtB3ZDaAyuRn8xsApZnygxl0hHgCH9AR4gkMdkI8gcFOiCcw2AnxBAY7IZ7AYCfEExjshHgCg50QT/j/1YgvzHMEfHkAAAAASUVORK5CYII=\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "dataset.visualize(9)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "f60427a6-e940-4c95-9710-06c015887f32",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAPsAAAD7CAYAAACscuKmAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjUuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/YYfK9AAAACXBIWXMAAAsTAAALEwEAmpwYAAAfOklEQVR4nO2de5Dc1XXnv6d7ep56IWmQBBIaBWQIEBB4zHuzSDKgYNZ4t1hs7HhVLtVq48W7TiU2hn14na04hVO1dpxszFpriOUUa3AS21DYBcaywOvFPIbwsEDYEkKKJPQYCUkjaZ7dffaPbv3uuXeme1oz/RjN/X6qpub8fvf+7j39OP0753fvPVdUFYSQ6U+q0QoQQuoDjZ2QSKCxExIJNHZCIoHGTkgk0NgJiYRJGbuIrBGRX4vIdhG5t1pKEUKqj0x0nF1E0gB+A+AmAHsAvATgLlV9s3rqEUKqRdMkrr0KwHZV3QEAIvIIgNsBlDT2+fPna1dX1yS6JISUY+fOnTh06JCMVTYZYz8XwG5zvAfA1eUu6OrqQk9PzyS6JISUo7u7u2RZzR/Qich6EekRkZ7e3t5ad0cIKcFkjH0vgCXmeHHxnIeqblDVblXt7uzsnER3hJDJMBljfwnAchFZJiLNAD4G4PHqqEUIqTYTjtlVNSsinwHwFIA0gIdU9Y2qaUYIqSqTeUAHVf0xgB9XSRdCSA3hDDpCIoHGTkgk0NgJiQQaOyGRQGMnJBJo7IREAo2dkEigsRMSCTR2QiKBxk5IJNDYCYkEGjshkUBjJyQSaOyERAKNnZBIoLETEgk0dkIigcZOSCTQ2AmJBBo7IZFAYyckEmjshEQCjZ2QSKCxExIJNHZCImFcYxeRh0TkoIhsMefmisjTIrKt+P+s2qpJCJksldzZvw1gTXDuXgCbVHU5gE3FY0LIFGZcY1fVnwN4Lzh9O4CNRXkjgI9UVy1CSLWZaMy+QFX3FeX9ABZUSR9CSI2Y9AM6VVUAWqpcRNaLSI+I9PT29k62O0LIBJmosR8QkUUAUPx/sFRFVd2gqt2q2t3Z2TnB7gghk2Wixv44gLVFeS2Ax6qjDiGkVlQy9PZdAL8EcKGI7BGRdQDuB3CTiGwD8MHiMSFkCtM0XgVVvatE0eoq60IIqSHjGjuZ6gTPRjVv5JLPTYGUdeqksq7yQXtijsO+7KGYvlKctNko+M4TEgk0dkIigW78mUA+7x9br1uC32tJj12vGqTKNFhpX2VfS9hItV9A3PDOTkgk0NgJiQQaOyGRwJh9KmHjWfszXG64KjvkHeaOvuWae+81V3Byu1/v5O5E1pGjXpnY2NwOqUnG77vlbKdiR5dX1HTWxe5gzmWu3ozFqBjNOdG8ITIqtieVwDs7IZFAYyckEujGNxLjpgIAUukxq+Xfe9M/3v3DRM4d2Ow3kd2TyJlZpr0m3/XN2E8+KCs54jVqQt5WJ5541i865EKSoYEWp9/My/2uFt3qypb4CZFSLTPHVikcvuOsvIrgu0RIJNDYCYkEuvE1x/d982YxSSpw23NHtrl6W/+na+HgT7x6zXPddek2v7chdWU5O5su8HzLJBfyJrJ5tcJL1Jv+5pdl3HEqPWxOv+TX2/VCIo5se8Aryv3WukROL3eLL1Oh227DIRk7FCK8sxMSDTR2QiKBxk5IJDBmrwXeTDj/9zRlZqeNbPmmf93b30jEzKxsIg/O9tsYMLG+iB9Ie8dSYkYexlhgVoKy1bzcFaWTV+RNZ4NBi2pi+5aUn7c0tevLiTyy50euvRX/zavXNO9CV5bLemWpNL/ip+CdnZBIoLETEgn0capC4MIa1z2f7feKss99PpGb+3/mlQ21u+v6005Opfz2U6ky7rm3gMa4zKeRF6LioTd7TTi0Z6MJ4+JrPgw7nDwcxBYq7uvZlnOLeoaf/ahXb+T9f57ImaU3B3pwWO4UvLMTEgk0dkIigcZOSCQwZp8wLvbMh/nU8wOJmH3233tFzVk3PXQgEySDyLjA1xsxChI9enF5uvSKNa/eqJi9TDxf5jIPO9wWxvNe+nrXigQL/dQ8jwiHA8WU9Zv7UqZj2K/Y85lEHM7+D6+o+fx/4Q7ssFyEQ3KVbP+0REQ2i8ibIvKGiHy2eH6uiDwtItuK/8+qvbqEkIlSiRufBfDHqnoxgGsA3C0iFwO4F8AmVV0OYFPxmBAyRalkr7d9APYV5eMishXAuQBuB3BjsdpGAM8A+EJNtJyKmG2WUkHu9sH/+0eJ3DryglfWb1z3VIs/XmVzv3mjRKGrbroL08Zbl7+8G2/lSnO6hbPkzHXB0JtaHb2toMKQR0qW2VEz+3aMBC+6Sd17Kq/c45WNtC9I5Myiq0y/8SXAOK1XKCJdAK4A8AKABcUfAgDYD2BBqesIIY2nYmMXkRkA/gHAH6pqny3TwsToMadciMh6EekRkZ7e3t5JKUsImTgVGbuIZFAw9IdV9fvF0wdEZFGxfBGAg2Ndq6obVLVbVbs7OzuroTMhZAKMG7NLIUn3gwC2qupXTdHjANYCuL/4/7GaaDiVKDH1cuj1/+VVa+1/JpEHMs1eWarZxYoSDqmViNNHzfIsFdsDQb55U69szF6mzKJBgR15C7dws8NydhguUERyWrrM68rVSwcKZludnMn7ZblfumG51OpHXRuzu3yFy6xUnC5UMth4PYBPAviViLxaPPefUDDy74nIOgC7ANxZEw0JIVWhkqfxv0Dp3/rV1VWHEFIr4ptGdDqUGZ7JHXsnkXXrX3nVhm2+9nB4rZx7XqosHHork0fStu/lswx/rsu5+KUIR96Mqx6OqGnONZo3F47uyuoRDL15tcSc9+tZt36k1f9Kt6eOJ/LJ57+YyB23/G2gR5klfdOE6RmcEEJGQWMnJBLoxpdjlM/pTgy/9peJ3DbLd6YHWswsuaYwsYWR05U9jQ9/ktNmu6bUjKDQqJIdMTnqq5G3YdSaG+NaB/FE2u481Wx0HPbfj5F+e2H4NN48qfeSYfh92b7taAcADAy7r3jL0V84NXb722Y1L1lpGgxW60yTpBe8sxMSCTR2QiKBxk5IJDBmDykzkyp/fHci696nEnlwjt+EmCQU4biWnxAy6NvG6UbMtAWr6oZc+1tePOmVdZ7tnhcsvcDN3ssPVn9oycbR6Rb/dQ4POB1ff9El3eyc63/llnY5HbMn/Hhb7bBcWsc8DwDI2yHAIGllqztuUjfV7sSWDV49P2avcAvrMwze2QmJBBo7IZFAN34U1pX0fwuzO3+cyO2zTX60Zv9tTJdz1Usll4DvPTYZ9/PAYX9Lo//y3/cn8rbtQ0EbTq+7/51bZXjnXX7WsFy/XZCDigh3eBKzdVNvb6DjF/cl8ltvGR2DHPh3/9v5Tsc7fB2H+9wQmJfYI1TEhD/hNlSSdq9z2IRleuB5r17umAvR0rOX+O1Pk0UyZ67mhJDTgsZOSCTQ2AmJBMbsoyj9+5d79xl3YKas2ti1cMIMBY2K2a0cDBPZog5X8Tvf8JMADRxxH9uHrp/tlf3oOVf3m986lMg3XNfh1TtnsRuiywZTWFMlhprC/PiZVjeNdOPfHvZ1PDy2jj967oBX75sPuetuuNrXcdFCp6OdVns6iThsrJ/LuPd05ky/3uDunyZyevan/EIt/RznTOLM1ZwQclrQ2AmJBLrxIXa75cEjXlG+7zeJnLXec7hfsbe1UpkZdGE6eLu4yjT5T3tGvHrDZlHWtt3+DLoms6PUgHF9977rt3HOMpMbbygcU8OYpEL/ftC1//YOf0umYaP/tj1Ox0wQ8gyZvt/5J7+Nc5YaHQdMQahHqnTYZD1wMaEXWv2tt7L7nnMHlwZufKm45gyDd3ZCIoHGTkgk0I0vg57Y6x03w+UzGxH31oUpoVHGVS+5yyr8nZBsnrmOdj95wjv7nB6z+/1G0naGl2mvrTVQZALrYsKJa/bbM2uGr+PLe52Os2Y7ndLBDDQ7423mrDJJIuxloR6VpsW2WS+CXVxzR98q3XelUwynONPjVRBCxoXGTkgk0NgJiQTG7GXQ/n3ecabV5I1vcvFl2d2Twp/TcrO9SpTdctMsr97PnjmRyAP9XhFO9rshtuvMjLTlF7V69fImuUTlC7nC5JnuwjW3+FPSNj87to5WPwC4/lqn40WX+DrmhkroOCpZvtEwnF5n4vS8uTAfPmcZcjP58sMDXlGquQ3TgXE/ZhFpFZEXReQ1EXlDRP6keH6ZiLwgIttF5FERaR6vLUJI46jkN30IwCpVvRzACgBrROQaAF8B8DVVvQDAEQDraqYlIWTSVLLXmwI45ZNlin8KYBWAjxfPbwTwJQAPVF/FxqHDx/0TZgaWlp4kN07OsrHzzAFAyiRhsMklVq6e4dX7ii5K5Cef9nVctNB9pJ/4xNxEbmkOdjc1i18q3f4pnEGXO2l19N34+7/s5Kd+4nS0+gHAJ3/f6ZgJdey3Otr3vrIZfyFW/VxwUX6kz7Uffu6xuPEAICLp4g6uBwE8DeBtAEdV9VR6kj0Azq2JhoSQqlCRsatqTlVXAFgM4CoAF1XagYisF5EeEenp7e2dmJaEkElzWkNvqnoUwGYA1wKYI5JMI1sMYG+Jazaoareqdnd2do5VhRBSB8aN2UWkE8CIqh4VkTYAN6HwcG4zgDsAPAJgLYDHaqloQ8ie8I/TpQL105h76l0Wblpmmzfxe7AqbeUtbihu5Rp/WM5r31yXCxJUVLxNcxlsG7kgL/2qm51eq4y+o+Jrs3LODrUV2jdt2ubLvd0TfV05s+IuN1K63hlMJePsiwBsFJE0Cp7A91T1CRF5E8AjIvKnAF4B8GAN9SSETJJKnsa/DuCKMc7vQCF+J4ScAXAGXTlS/owuu83Q6CVgp480+/OQxORI8+S8v4Vw9qTJ0R7oYQ+9BXChextucTRJJPCtR06YnO+22+C6lNnbOdXir3pTq6OXb97PUa9Z+35Uou04VGV/66kH58YTEgk0dkIigW58GaTZT22M3Nj1RrmOZVxJMU5t/1tv+833uaf/zQvaEzlz9lyvXnqWe7otmSAUMG6x5+GPevJf4kn3qBNl3H1v4Y5fL2NnvNmCnO+C580qmeGDx7yykf0uB+BIn3tC3nqeP38rM8dtG6UjfvsWL1W3BqtpTMgmmXZMR3hnJyQSaOyERAKNnZBIYMxejtYF3qGd4aWtRs4Ev5l2+KvJz08++M6ORM7IKq9s5nX/JpFPbn8qkY+//v/89kd2JWK61c+1nm43CR86Wtz5thavnthtptPlMmyUju3tkJcO+bPOcicHEznf7x525Ab8Ya181jyPaPWXXLR13ZLIM357WSKfeP7zvrZN7llHKnzOkrVDgO51pnL+VtfIuI0AUs3+KkPv+UY1ph42CN7ZCYkEGjshkUA3vgypGf4QT/+we7syZjadhsNaYY40g5qFH9963k+ScMmsOYl8wwe/kciBYwqMuHxpuXdf94qGDryWyNm+ne78IT+fng655cYSKpw3oYHYr0jgwqaduytt/orGdMd5Tl64NJE7Fl7p1UstvMQc+cOI9l39xXMvur72+ENjV3c5Nz6HwAX3PieTN3DEd+NTs37HHIRhmR1zPXNn1/HOTkgk0NgJiQQaOyGRwJh9FC7GS3Wc7ZWkZnQlcjPeSeTBvB/LilmhlR/x59i2nT0vkV//5dNe2YOPuuNLL3XDUNet/ohX7+qVH07k5Zf+M69s5tKVmMr0D/nvx/Z/fCmRX3jmR17ZLzc/nsivvOyeTfzXj/rDmU1z3Hs1csKfLuvN6DWfk/YHW113XVNaaW/orXS1qQ7v7IREAo2dkEigGx+SM25msK1vutMl5pEjbiYcAjfejhnlh323taXTDQ1deMk8r+z1XjcLbc8u1/7f/NX9Xr2Hv/nnidy54Byv7Lzz35fIXef/diIv7lru1ZvXuTCRW9uDwb1Sq/bCHZPMirXDB/2hvb27nf67trvtkHft+I1X7+D+Pa69Qd8Fb2137vqsmeZ9u9J/37x7VrCqzibp0AE33NZ32M8v2H7jzSiFmhd+BnvxvLMTEgs0dkIigW58SLi7pyG99EOJnN39sCsQ/zdTm8xxMDlNc679VR/03dHHfvFuIjdl3JZDZ83z3exc3jV65L3DXtmBfZsT+flnf+b6Lb0BKzJNlc0KC737rAl5gjR53noRu61Vc4u/IKel1c2GmznT1+NIn3PJP3CZu+6CS2d79YZPjp3vrnDs2swMuGQY/TMu8eplOi8zR0Ha7WmSk453dkIigcZOSCTQ2AmJBMbsIWLis7wfcDctfH8i97e7VVJt+Te9eoM5kyQheIeHTrg2r7jc3wr40590iRy++r9dLD5nlv+bbOPSpow/myzT7GJbu8WyBNHskBmiGg6Gq0at4iuSDp5NzDSJOdLBSrG8jr0qUINEjzZvRt8JPwHG/LPcZ/G5P3Cr6sJcIXZSXip4nbkBV5g5eiiR29//Jb8Ro38+eD9S6elhJhXf2YvbNr8iIk8Uj5eJyAsisl1EHhWR5vHaIIQ0jtNx4z8LYKs5/gqAr6nqBQCOAFhXTcUIIdWlIv9ERBYD+BCALwP4IylsMboKwMeLVTYC+BKAB2qgY+OQ0J11v43Nl33GVXvpD0pVg7YHwzhmGGrwmD9edde/dPnPB80OrA985z2vXnubGcoK/Ck1oUe/yQuXC1zzxTPmJPLSWX5e+rlmOGzYjKntP9nn1dtu3OJjA/7Ckg6Tz15MOJFK+XocO+70XdDpfx3vv8/N8jvvHNfeQJ//vqXMa8vn/DYy77nhzCPalcjzLv59+JgFUNPEbQ+p9M7+FwDugRs1ngfgqKqeCm72ADh3jOsIIVOEcY1dRG4DcFBVX55IByKyXkR6RKSnt7d3/AsIITWhkjv79QA+LCI7ATyCgvv+dQBzRJIEZYsB7B3rYlXdoKrdqtrd2dk5VhVCSB2oZH/2+wDcBwAiciOAz6nqJ0Tk7wDcgcIPwFoAj9VOzQYhwTRJk3iwabFLEtH/1j/3qrVnf+7KRvygOm1iVgmm0g4dd+1/6mMujl54tv8x/fW3XQzfezgYJsq4Nt43b34if/RCP9HjFWcvSeTZzf4QYKpEbvSRnB8r7z7ppp/+cLuf+PLJd9xwZH7I3VPCrdiuusL1fc+n/ZvBknPc0N6Aeb6RGrVNtdnf7qi/mm34kHuuMOM2M8W5yd+OW81nK+HnPk2YzKSaL6DwsG47CjH8g9VRiRBSC07rsaOqPgPgmaK8A8BV5eoTQqYO03OMoVaYRAgwM8Farv1Tr1r/j93quLYZ/pDUgHER083B0J5ZETdohpd+b/Usr9qVl7qhsW9/33/oObDD5c37zPvc9kkzM77b2p93iRzCvHBaIntFODttSYvr6/O/c6tXtnzGokR+4qjbvurOW+d49dasnOnaD7od6DNbWXmuu++Q5o67PPcdh3d5Zccv/Vwityy9yRWo/5qnq+tu4dx4QiKBxk5IJEipRQ+1oLu7W3t6eurWX02xi2SCRSAjB1518qY7vLK2uWYmWJu/iCWdtvnvnMsc7ibV2ub6E987h/a5E8Pvuif6A/v9UADHnPuModLbLlkkHWjS7nLQpef6W1m1n3vMtTfPLeppavNDgWEzgy6f9du3ExjtE/d8n791U/t+l++ur/NfeWWzbvuuacTsvBvc5+QM3p3V0t3djZ6enjFfDO/shEQCjZ2QSKCxExIJHHqbKDZOD4ZxMgtWuKIbH/bKBn7+qURulwG/rNUkllS3Yi3cQXhowMS2g36ZZNxQn5zvVql1XBCEcTkz6ywbfA107PhV0kFWySZzLH5ZzgwjYtj1NfCe/0TAvjQJkmPkTZyeO2AST7zn56jvO/djiTzjlu+MqXvY23SJ0U8H3tkJiQQaOyGRQDe+GpRZMNN8ztV+1VvcTqUnf/5pr6xj8FeJnJ05x8lBrnXk3WoSCVfTDJvtjobd0J4GP+uex5wKVqeU8HCDaAUYcY1oPuigxJBuKniv7BBY7ni/V9Zy+IDr6pALTwav8WcszvrAvV6LJYnQdbfwzk5IJNDYCYkEGjshkcCYvRbYuDQcljtrmZNvf9IrG3zlLxM5t9XJHTOPevV0hpvqOtjkT3VNZexHaveO9mP7fN4k0UAZysS5NjqW8L5h9kfLD5tnAif95BLNx4+69g4d8cqGFqxO5LZ//WeJnFn4Ab8v+9pCfSOP0y28sxMSCTR2QiKBbnytCYfljMupgYvZesV/TOTc+z6eyANvPuTVy+14JJFTQ297ZS12e+dWtwIuH7j7ObvdUfCbnzeOvdjcbMFwWtpsk5Qa9leipQfdMFquzw2bHevz28ids8rpvsYfiswsuw1jEo4BTpMtlWsN7+yERAKNnZBIoBtfb4z7POo5sXFP0x0uDXTbB+7x673f5VUbefd5r2hg909d2f7nEjl3bLtXLz9stpTK+wtyWpudZtkRp9NIzneXJe1y4UnLQq8sM//6RG5ZcWMizz5vtVcvPed8lMaOJhiZbvuE4J2dkEigsRMSCTR2QiKBMftUwhums9sQB1sUmy2FM4uv88rssbepU3bYq5frP+h6GvZntSHrhsoycDq1N3d41dKtbovpVPt8TAhv9luwYs2+HynOhJssle7PvhPAcQA5AFlV7RaRuQAeBdAFYCeAO1X1SKk2CCGN5XTc+JWqukJVu4vH9wLYpKrLAWwqHhNCpiiTceNvB3BjUd6Iwh5wX5ikPiTBua3WbQfgJYbQIKu82CEqm6EimEGXnrV48iqWwwwj2kU3KQkXzPCxUb2o9J1WAD8RkZdFZH3x3AJVPZX5bz+ABVXXjhBSNSq9s9+gqntF5GwAT4vIW7ZQVVUkfLpSoPjjsB4AzjvvvEkpSwiZOBXd2VV1b/H/QQA/QGGr5gMisggAiv8Plrh2g6p2q2p3Z2dndbQmhJw24xq7iHSIyMxTMoCbAWwB8DiAtcVqawE8VislSYBI8ieS9v6QbnJ/qZT7U/X/8vnSf5pL/jTv/kbVg5q/UMd08pdKNyV/nk6M1+tKJW78AgA/KCbVbwLwf1T1SRF5CcD3RGQdgF0A7qydmoSQyTKusavqDgCXj3H+MIDVo68ghExFOIMuFiaYm40p3KYPDJoIiQQaOyGRQGMnJBJo7IREAo2dkEigsRMSCTR2QiKBxk5IJNDYCYkEGjshkUBjJyQSaOyERAKNnZBIoLETEgk0dkIigcZOSCTQ2AmJBBo7IZFAYyckEmjshEQCjZ2QSKCxExIJNHZCIoHGTkgkVGTsIjJHRP5eRN4Ska0icq2IzBWRp0VkW/H/WbVWlhAycSq9s38dwJOqehEKW0FtBXAvgE2quhzApuIxIWSKUskurrMB/C6ABwFAVYdV9SiA2wFsLFbbCOAjtVGREFINKrmzLwPQC+BvROQVEflWcevmBaq6r1hnPwq7vRJCpiiVGHsTgCsBPKCqVwA4icBlV9USm3QDIrJeRHpEpKe3t3ey+hJCJkglxr4HwB5VfaF4/PcoGP8BEVkEAMX/B8e6WFU3qGq3qnZ3dnZWQ2dCyAQY19hVdT+A3SJyYfHUagBvAngcwNriubUAHquJhoSQqlDp/uz/AcDDItIMYAeAT6HwQ/E9EVkHYBeAO2ujIiGkGlRk7Kr6KoDuMYpWV1UbQkjN4Aw6QiKBxk5IJNDYCYkEGjshkUBjJyQSaOyERAKNnZBIkMK09jp1JtKLwgSc+QAO1a3jsZkKOgDUI4R6+JyuHktVdcx56XU19qRTkR5VHWuSTlQ6UA/qUU896MYTEgk0dkIioVHGvqFB/Vqmgg4A9QihHj5V06MhMTshpP7QjSckEupq7CKyRkR+LSLbRaRu2WhF5CEROSgiW8y5uqfCFpElIrJZRN4UkTdE5LON0EVEWkXkRRF5rajHnxTPLxORF4qfz6PF/AU1R0TSxfyGTzRKDxHZKSK/EpFXRaSneK4R35GapW2vm7GLSBrAXwP4PQAXA7hLRC6uU/ffBrAmONeIVNhZAH+sqhcDuAbA3cX3oN66DAFYpaqXA1gBYI2IXAPgKwC+pqoXADgCYF2N9TjFZ1FIT36KRumxUlVXmKGuRnxHape2XVXr8gfgWgBPmeP7ANxXx/67AGwxx78GsKgoLwLw63rpYnR4DMBNjdQFQDuAfwRwNQqTN5rG+rxq2P/i4hd4FYAnAEiD9NgJYH5wrq6fC4DZAN5B8VlatfWopxt/LoDd5nhP8VyjaGgqbBHpAnAFgBcaoUvRdX4VhUShTwN4G8BRVc0Wq9Tr8/kLAPcAyBeP5zVIDwXwExF5WUTWF8/V+3Opadp2PqBD+VTYtUBEZgD4BwB/qKp9jdBFVXOqugKFO+tVAC6qdZ8hInIbgIOq+nK9+x6DG1T1ShTCzLtF5HdtYZ0+l0mlbR+Pehr7XgBLzPHi4rlGUVEq7GojIhkUDP1hVf1+I3UBAC3s7rMZBXd5joicyktYj8/negAfFpGdAB5BwZX/egP0gKruLf4/COAHKPwA1vtzmVTa9vGop7G/BGB58UlrM4CPoZCOulHUPRW2iAgK22htVdWvNkoXEekUkTlFuQ2F5wZbUTD6O+qlh6rep6qLVbULhe/Dz1T1E/XWQ0Q6RGTmKRnAzQC2oM6fi9Y6bXutH3wEDxpuBfAbFOLD/1zHfr8LYB+AERR+PdehEBtuArANwE8BzK2DHjeg4IK9DuDV4t+t9dYFwGUAXinqsQXAF4vnfwvAiwC2A/g7AC11/IxuBPBEI/Qo9vda8e+NU9/NBn1HVgDoKX42PwRwVrX04Aw6QiKBD+gIiQQaOyGRQGMnJBJo7IREAo2dkEigsRMSCTR2QiKBxk5IJPx/SUGuR+4o/IQAAAAASUVORK5CYII=\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "dataset.visualize(1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a9aeedb4-b929-465e-8713-9577c7453c8b",
   "metadata": {},
   "source": [
    "### Make NCA"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "73cbd2e5-08b9-4e56-bc61-25c05dfc6480",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([4, 64, 64])"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset.target_size()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "aafc3d69-e141-4563-95e2-887638b6cb6c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Optional, Tuple  # noqa\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from torch.nn import Embedding\n",
    "\n",
    "from controllable_nca.utils import build_conv2d_net\n",
    "\n",
    "class DeepEncoder(nn.Module):\n",
    "    def __init__(self, num_embeddings: int, out_channels: int):\n",
    "        super().__init__()\n",
    "        self.num_embeddings = num_embeddings\n",
    "        self.embedding = Embedding(num_embeddings, 32)\n",
    "        self.encoder = nn.Sequential(\n",
    "            nn.Linear(32, 32),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(32, out_channels, bias=False),\n",
    "        )\n",
    "\n",
    "    def forward(self, indices):\n",
    "        embeddings = self.encoder(self.embedding(indices))\n",
    "        return embeddings\n",
    "\n",
    "NUM_HIDDEN_CHANNELS = 16\n",
    "\n",
    "encoder = DeepEncoder(dataset.num_goals(), NUM_HIDDEN_CHANNELS)\n",
    "\n",
    "nca =  ControllableNCA(num_goals=dataset.num_goals(), use_image_encoder=False, encoder=encoder, target_shape=dataset.target_size(), living_channel_dim=3, num_hidden_channels=NUM_HIDDEN_CHANNELS, cell_fire_rate=0.5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "348e0a9e-75ee-4f1a-9c82-7119e65ab9f1",
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device('cuda')\n",
    "nca = nca.to(device)\n",
    "dataset.to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "f06a7022-06c9-4c59-bd6c-219ea60be756",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Writing to tensorboard_logs/2022-04-27 02:14:17.774250\n"
     ]
    }
   ],
   "source": [
    "trainer = MorphingImageNCATrainer(nca, dataset, nca_steps=[48, 96], lr=1e-3, num_damaged=0, damage_radius=3, device=device, pool_size=1024)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "97078f4f-c9ad-4e75-82eb-832a750a2d73",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ControllableNCA(\n",
       "  (perception_net): Conv2d(21, 63, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=21, bias=False)\n",
       "  (update_net): UpdateNet(\n",
       "    (out): Sequential(\n",
       "      (0): Conv2d(63, 64, kernel_size=(1, 1), stride=(1, 1))\n",
       "      (1): ReLU()\n",
       "      (2): Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1))\n",
       "      (3): ReLU()\n",
       "      (4): Conv2d(64, 21, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "    )\n",
       "  )\n",
       "  (encoder): DeepEncoder(\n",
       "    (embedding): Embedding(10, 32)\n",
       "    (encoder): Sequential(\n",
       "      (0): Linear(in_features=32, out_features=32, bias=True)\n",
       "      (1): ReLU()\n",
       "      (2): Linear(in_features=32, out_features=16, bias=False)\n",
       "    )\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "nca"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "191755f4-3b65-41b6-83af-edebc14ece1f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# trainer.train(batch_size=8, epochs=100000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "73f62fa2-fea2-42fb-9629-073a3ab05383",
   "metadata": {},
   "outputs": [],
   "source": [
    "nca.load(\"../saved_models/100k_linear_encoder.pt\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eca987e9-f8f5-40c0-b936-4b7b5f61a0a4",
   "metadata": {},
   "source": [
    "### Interactive Visualization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "fca43ee4-18b2-42d1-9d5e-7f1948ba4895",
   "metadata": {},
   "outputs": [],
   "source": [
    "from threading import Event, Thread\n",
    "\n",
    "import cv2\n",
    "import numpy as np\n",
    "import torch\n",
    "from einops import rearrange\n",
    "from ipycanvas import Canvas, hold_canvas  # noqa\n",
    "from ipywidgets import Button, HBox, VBox\n",
    "\n",
    "from controllable_nca.utils import create_2d_circular_mask, rgb\n",
    "\n",
    "\n",
    "def to_numpy_rgb(x, use_rgb=False):\n",
    "    return rearrange(\n",
    "        np.squeeze(rgb(x, use_rgb).detach().cpu().numpy()), \"c x y -> x y c\"\n",
    "    )\n",
    "\n",
    "\n",
    "class MorphingImageVisualizer:\n",
    "    def __init__(\n",
    "        self,\n",
    "        trainer,\n",
    "        image_size,\n",
    "        rgb: bool = False,\n",
    "        canvas_scale=5,\n",
    "        damage_radius: int = 5,\n",
    "    ):\n",
    "        self.trainer = trainer\n",
    "        self.current_state = None\n",
    "        self.current_goal = None\n",
    "\n",
    "        self.image_size = image_size\n",
    "        self.rgb = rgb\n",
    "        self.canvas_scale = canvas_scale\n",
    "        self.canvas_size = self.image_size * self.canvas_scale\n",
    "\n",
    "        self.canvas = Canvas(width=self.canvas_size, height=self.canvas_size)\n",
    "        self.canvas.on_mouse_down(self.handle_mouse_down)\n",
    "        self.stopped = Event()\n",
    "\n",
    "        self.current_goal = torch.tensor(0, device=self.trainer.device)\n",
    "\n",
    "        self.device = self.trainer.device\n",
    "        self.damage_radius = damage_radius\n",
    "        self.current_state = self.trainer.nca.generate_seed(1).to(self.device)\n",
    "\n",
    "        def button_fn(class_num):\n",
    "            def start(btn):\n",
    "                self.current_goal = torch.tensor(class_num, device=self.trainer.device)\n",
    "                if self.stopped.isSet():\n",
    "                    self.stopped.clear()\n",
    "                    Thread(target=self.loop).start()\n",
    "\n",
    "            return start\n",
    "\n",
    "        button_list = []\n",
    "        for i in range(len(self.trainer.target_dataset.EMOJI)):\n",
    "            button_list.append(Button(description=self.trainer.target_dataset.EMOJI[i]))\n",
    "            button_list[-1].on_click(button_fn(i))\n",
    "\n",
    "        self.vbox = VBox(button_list)\n",
    "\n",
    "        self.stop_btn = Button(description=\"Stop\")\n",
    "\n",
    "        def stop(btn):\n",
    "            if not self.stopped.isSet():\n",
    "                self.stopped.set()\n",
    "\n",
    "        self.stop_btn.on_click(stop)\n",
    "\n",
    "    def handle_mouse_down(self, xpos, ypos):\n",
    "        in_x = int(xpos / self.canvas_scale)\n",
    "        in_y = int(ypos / self.canvas_scale)\n",
    "\n",
    "        mask = create_2d_circular_mask(\n",
    "            self.image_size,\n",
    "            self.image_size,\n",
    "            (in_x, in_y),\n",
    "            radius=self.damage_radius,\n",
    "        )\n",
    "        self.current_state[0][:, mask] *= 0.0\n",
    "\n",
    "    def draw_image(self, rgb):\n",
    "        with hold_canvas(self.canvas):\n",
    "            rgb = np.squeeze(rearrange(rgb, \"b c w h -> b w h c\"))\n",
    "            self.canvas.clear()  # Clear the old animation step\n",
    "            self.canvas.put_image_data(\n",
    "                cv2.resize(\n",
    "                    rgb * 255.0,\n",
    "                    (self.canvas_size, self.canvas_size),\n",
    "                    interpolation=cv2.INTER_NEAREST,\n",
    "                ),\n",
    "                0,\n",
    "                0,\n",
    "            )\n",
    "\n",
    "    def loop(self):\n",
    "        with torch.no_grad():\n",
    "            self.current_state = self.trainer.nca.generate_seed(1).to(self.device)\n",
    "            while not self.stopped.wait(0.02):  # the first call is in `interval` secs\n",
    "                # update_particle_locations()\n",
    "                self.draw_image(self.trainer.to_rgb(self.current_state))\n",
    "                self.current_state = self.trainer.nca.grow(\n",
    "                    self.current_state, 1, self.current_goal\n",
    "                )\n",
    "\n",
    "    def visualize(self):\n",
    "        Thread(target=self.loop).start()\n",
    "        display(self.canvas, HBox([self.stop_btn, self.vbox]))  # noqa\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "332fb42a-c2c5-4259-be16-844b7220a1da",
   "metadata": {},
   "outputs": [],
   "source": [
    "viz = MorphingImageVisualizer(trainer, 64)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "bd22d9cf-a084-40fc-a675-5975d52963ba",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4c71f31b002c4c288c9064a8f0a83404",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Canvas(height=320, width=320)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "46c042f605414c56b058f0092fb8ddeb",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(Button(description='Stop', style=ButtonStyle()), VBox(children=(Button(description='🦎', style=B…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "viz.visualize()"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "1d64cb66d3d902aa83000daa06ca958bef94bde318911a82aee5f8df2bb8934b"
  },
  "kernelspec": {
   "display_name": "Python [conda env:pycanvas] *",
   "language": "python",
   "name": "conda-env-pycanvas-py"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}