4_Controller Data Prepper.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "234c630d",
   "metadata": {},
   "source": [
    "# Data Preprocessor for the Controller\n",
    "The controller receives the z produced by the encoder (from the VAE) and the z' produced by the RNN's prediction to determine an action.\n",
    "The below processor takes a collection of image data and the action performed for preprocessing. It finds z and z' by putting the image through the encoder and then putting z through the RNN.\n",
    "In the form I left it in, we did not have a good dataset of images and corresponding images. Right now, it just sets every action to 22 (no action)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4645b2a9",
   "metadata": {},
   "source": [
    "## Loading stuff into Kernal"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "5a04aeb5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# imports\n",
    "import tensorflow.compat.v1.keras.backend as K\n",
    "import tensorflow as tf\n",
    "tf.compat.v1.disable_eager_execution()\n",
    "from tensorflow import keras\n",
    "from tensorflow.keras.layers import Conv2D, Conv2DTranspose, Input, Flatten, Dense, Lambda, Reshape, MaxPooling2D, Dropout, LSTM\n",
    "from tensorflow.keras.models import Model\n",
    "from tensorflow.keras.models import save_model\n",
    "from tensorflow.keras.datasets import mnist\n",
    "from tensorflow.keras import backend as K\n",
    "from tensorflow.keras.optimizers import Adam\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import os\n",
    "import pickle\n",
    "import joblib\n",
    "# added so that cv2 gets installed in kernal\n",
    "# import sys\n",
    "# !{sys.executable} -m pip install opencv-python\n",
    "# commented the above code, it started working, idk why\n",
    "# if code not working try uncommenting the above\n",
    "import cv2\n",
    "import random\n",
    "import glob"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "57faf598",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "128\n",
      "(128, 128, 1)\n"
     ]
    }
   ],
   "source": [
    "# data (preprocessed from Data Processing Script)\n",
    "\n",
    "train_data = joblib.load(\"images/train_data.z\")\n",
    "print(train_data.shape[2])\n",
    "\n",
    "# Reshape \n",
    "img_width  = train_data.shape[1]\n",
    "img_height = train_data.shape[2]\n",
    "num_channels = 1\n",
    "x_train = train_data.reshape(train_data.shape[0], img_height, img_width, num_channels)\n",
    "\n",
    "input_shape = (img_height, img_width, num_channels)\n",
    "print(input_shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "e6fdc5f9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model: \"encoder\"\n",
      "__________________________________________________________________________________________________\n",
      " Layer (type)                   Output Shape         Param #     Connected to                     \n",
      "==================================================================================================\n",
      " encoder_input (InputLayer)     [(None, 128, 128, 1  0           []                               \n",
      "                                )]                                                                \n",
      "                                                                                                  \n",
      " conv2d (Conv2D)                (None, 128, 128, 64  640         ['encoder_input[0][0]']          \n",
      "                                )                                                                 \n",
      "                                                                                                  \n",
      " max_pooling2d (MaxPooling2D)   (None, 64, 64, 64)   0           ['conv2d[0][0]']                 \n",
      "                                                                                                  \n",
      " dropout (Dropout)              (None, 64, 64, 64)   0           ['max_pooling2d[0][0]']          \n",
      "                                                                                                  \n",
      " conv2d_1 (Conv2D)              (None, 64, 64, 128)  73856       ['dropout[0][0]']                \n",
      "                                                                                                  \n",
      " max_pooling2d_1 (MaxPooling2D)  (None, 32, 32, 128)  0          ['conv2d_1[0][0]']               \n",
      "                                                                                                  \n",
      " dropout_1 (Dropout)            (None, 32, 32, 128)  0           ['max_pooling2d_1[0][0]']        \n",
      "                                                                                                  \n",
      " conv2d_2 (Conv2D)              (None, 32, 32, 64)   73792       ['dropout_1[0][0]']              \n",
      "                                                                                                  \n",
      " max_pooling2d_2 (MaxPooling2D)  (None, 16, 16, 64)  0           ['conv2d_2[0][0]']               \n",
      "                                                                                                  \n",
      " dropout_2 (Dropout)            (None, 16, 16, 64)   0           ['max_pooling2d_2[0][0]']        \n",
      "                                                                                                  \n",
      " conv2d_3 (Conv2D)              (None, 16, 16, 32)   18464       ['dropout_2[0][0]']              \n",
      "                                                                                                  \n",
      " flatten (Flatten)              (None, 8192)         0           ['conv2d_3[0][0]']               \n",
      "                                                                                                  \n",
      " dense (Dense)                  (None, 4096)         33558528    ['flatten[0][0]']                \n",
      "                                                                                                  \n",
      " latent_mu (Dense)              (None, 2048)         8390656     ['dense[0][0]']                  \n",
      "                                                                                                  \n",
      " latent_sigma (Dense)           (None, 2048)         8390656     ['dense[0][0]']                  \n",
      "                                                                                                  \n",
      " z (Lambda)                     (None, 2048)         0           ['latent_mu[0][0]',              \n",
      "                                                                  'latent_sigma[0][0]']           \n",
      "                                                                                                  \n",
      "==================================================================================================\n",
      "Total params: 50,506,592\n",
      "Trainable params: 50,506,592\n",
      "Non-trainable params: 0\n",
      "__________________________________________________________________________________________________\n",
      "None\n",
      "Model: \"decoder\"\n",
      "_________________________________________________________________\n",
      " Layer (type)                Output Shape              Param #   \n",
      "=================================================================\n",
      " decoder_input (InputLayer)  [(None, 2048)]            0         \n",
      "                                                                 \n",
      " dense_1 (Dense)             (None, 8192)              16785408  \n",
      "                                                                 \n",
      " reshape (Reshape)           (None, 16, 16, 32)        0         \n",
      "                                                                 \n",
      " conv2d_transpose (Conv2DTra  (None, 32, 32, 32)       9248      \n",
      " nspose)                                                         \n",
      "                                                                 \n",
      " conv2d_transpose_1 (Conv2DT  (None, 32, 32, 32)       9248      \n",
      " ranspose)                                                       \n",
      "                                                                 \n",
      " conv2d_transpose_2 (Conv2DT  (None, 64, 64, 64)       18496     \n",
      " ranspose)                                                       \n",
      "                                                                 \n",
      " conv2d_transpose_3 (Conv2DT  (None, 64, 64, 64)       36928     \n",
      " ranspose)                                                       \n",
      "                                                                 \n",
      " conv2d_transpose_4 (Conv2DT  (None, 128, 128, 64)     36928     \n",
      " ranspose)                                                       \n",
      "                                                                 \n",
      " decoder_output (Conv2DTrans  (None, 128, 128, 1)      577       \n",
      " pose)                                                           \n",
      "                                                                 \n",
      "=================================================================\n",
      "Total params: 16,896,833\n",
      "Trainable params: 16,896,833\n",
      "Non-trainable params: 0\n",
      "_________________________________________________________________\n"
     ]
    }
   ],
   "source": [
    "# load the vae (have to make the architecture again, make sure the code below\n",
    "#   matches the code in the Data Prepper/VAE Trainer)\n",
    "\n",
    "\n",
    "# ====== Encoder ======\n",
    "# changing this will make the model exponentially larger or smaller\n",
    "latent_dim = 2048\n",
    "\n",
    "# the model (saved in x)\n",
    "input_img = Input(shape=input_shape, name='encoder_input')\n",
    "x = Conv2D(64, 3, padding='same', activation='relu')(input_img)\n",
    "x = MaxPooling2D((2,2), padding = 'same')(x)\n",
    "x = Dropout(0.2)(x)\n",
    "x = Conv2D(128, 3, padding='same', activation='relu')(x)\n",
    "x = MaxPooling2D((2,2), padding = 'same')(x)\n",
    "x = Dropout(0.2)(x)\n",
    "x = Conv2D(64, 3, padding='same', activation='relu')(x)\n",
    "x = MaxPooling2D((2,2), padding = 'same')(x)\n",
    "x = Dropout(0.2)(x)\n",
    "x = Conv2D(32, 3, padding='same', activation='relu')(x)\n",
    "\n",
    "conv_shape = K.int_shape(x) #Shape of conv to be provided to decoder\n",
    "#Flatten\n",
    "x = Flatten()(x)\n",
    "x = Dense(latent_dim*2, activation='relu')(x)\n",
    "\n",
    "# Two outputs, for latent mean and log variance (std. dev.)\n",
    "#  Use these to sample random variables in latent space to which inputs are mapped. \n",
    "z_mu = Dense(latent_dim, name='latent_mu')(x)   #Mean values of encoded input\n",
    "z_sigma = Dense(latent_dim, name='latent_sigma')(x)  #Std dev. (variance) of encoded input\n",
    "\n",
    "#REPARAMETERIZATION TRICK\n",
    "# Define sampling function to sample from the distribution\n",
    "# Reparameterize sample based on the process defined by Gunderson and Huang\n",
    "# into the shape of: mu + sigma squared x eps\n",
    "#This is to allow gradient descent to allow for gradient estimation accurately. \n",
    "def sample_z(args):\n",
    "    z_mu, z_sigma = args\n",
    "    eps = K.random_normal(shape=(K.shape(z_mu)[0], K.int_shape(z_mu)[1]))\n",
    "    return z_mu + K.exp(z_sigma / 2) * eps\n",
    "\n",
    "# sample vector from the latent distribution\n",
    "# z is the labda custom layer we are adding for gradient descent calculations\n",
    "  # using mu and variance (sigma)\n",
    "z = Lambda(sample_z, output_shape=(latent_dim, ), name='z')([z_mu, z_sigma])\n",
    "\n",
    "#Z (lambda layer) will be the last layer in the encoder.\n",
    "# Define and summarize encoder model.\n",
    "encoder = Model(input_img, [z_mu, z_sigma, z], name='encoder')\n",
    "print(encoder.summary())\n",
    "\n",
    "# ==== Decoder ====\n",
    "\n",
    "# decoder takes the latent vector as input\n",
    "decoder_input = Input(shape=(latent_dim, ), name='decoder_input')\n",
    "\n",
    "# Need to start with a shape that can be remapped to original image shape as\n",
    "#we want our final utput to be same shape original input.\n",
    "#So, add dense layer with dimensions that can be reshaped to desired output shape\n",
    "x = Dense(conv_shape[1]*conv_shape[2]*conv_shape[3], activation='relu')(decoder_input)\n",
    "# reshape to the shape of last conv. layer in the encoder, so we can \n",
    "x = Reshape((conv_shape[1], conv_shape[2], conv_shape[3]))(x)\n",
    "# upscale (conv2D transpose) back to original shape\n",
    "# use Conv2DTranspose to reverse the conv layers defined in the encoder\n",
    "x = Conv2DTranspose(32, 3, padding='same', activation='relu',strides=(2, 2))(x)\n",
    "x = Conv2DTranspose(32, 3, padding='same', activation='relu')(x)\n",
    "x = Conv2DTranspose(64, 3, padding='same', activation='relu',strides=(2, 2))(x)\n",
    "x = Conv2DTranspose(64, 3, padding='same', activation='relu')(x)\n",
    "x = Conv2DTranspose(64, 3, padding='same', activation='relu',strides=(2, 2))(x)\n",
    "#Can add more conv2DTranspose layers, if desired. \n",
    "#Using sigmoid activation\n",
    "x = Conv2DTranspose(num_channels, 3, padding='same', activation='sigmoid', name='decoder_output')(x)\n",
    "\n",
    "# Define and summarize decoder model\n",
    "decoder = Model(decoder_input, x, name='decoder')\n",
    "decoder.summary()\n",
    "\n",
    "# apply the decoder to the latent sample \n",
    "z_decoded = decoder(z)\n",
    "\n",
    "# ===== Loss Function =====\n",
    "\n",
    "class CustomLayer(keras.layers.Layer):\n",
    "    def vae_loss(self, x, z_decoded):\n",
    "        x = K.flatten(x)\n",
    "        z_decoded = K.flatten(z_decoded)\n",
    "        \n",
    "        # Reconstruction loss (as we used sigmoid activation we can use binarycrossentropy)\n",
    "        recon_loss = keras.metrics.binary_crossentropy(x, z_decoded)\n",
    "        \n",
    "        # KL divergence\n",
    "        kl_loss = -5e-4 * K.mean(1 + z_sigma - K.square(z_mu) - K.exp(z_sigma), axis=-1)\n",
    "        return K.mean(recon_loss + kl_loss)\n",
    "\n",
    "    # add custom loss to the class\n",
    "    def call(self, inputs):\n",
    "        x = inputs[0]\n",
    "        z_decoded = inputs[1]\n",
    "        loss = self.vae_loss(x, z_decoded)\n",
    "        self.add_loss(loss, inputs=inputs)\n",
    "        return x\n",
    "\n",
    "# apply the custom loss to the input images and the decoded latent distribution sample\n",
    "y = CustomLayer()([input_img, z_decoded])\n",
    "# y is basically the original image after encoding input img to mu, sigma, z\n",
    "# and decoding sampled z values.\n",
    "#This will be used as output for vae"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "34d38731",
   "metadata": {},
   "outputs": [],
   "source": [
    "# load encoder\n",
    "vae = Model(input_img, y, name = 'vae')\n",
    "vae.load_weights(os.getcwd() + \"\\\\models\\\\vae.h5\")\n",
    "encoder = Model(vae.input, vae.layers[15].output)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "144a8880",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[ 0.0104093  -0.7574463   0.45002717 ... -1.2421685  -0.7610112\n",
      "  -1.6004298 ]]\n",
      "(1, 2048)\n",
      "2048\n"
     ]
    }
   ],
   "source": [
    "# load rnn preprocessed data\n",
    "data = joblib.load(\"images/train_data_rnn.z\")\n",
    "train_data = np.array([np.array(p[0])for p in data])\n",
    "answers = np.array([np.array(p[0]) for p in data])\n",
    "print(train_data[0])\n",
    "print(np.shape(train_data[0]))\n",
    "z_len = np.shape(train_data[0])[-1]\n",
    "print(z_len)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "c2b8d262",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Rebuild RNN architecture (make sure this matches architecture in RNN Trainer)\n",
    "input_to_rnn = Input(shape=(1,z_len))\n",
    "\n",
    "x = LSTM(2048, return_sequences=True)(input_to_rnn)\n",
    "x = Dropout(0.2)(x)\n",
    "x = Dense(2048)(x)\n",
    "x = Dropout(0.2)(x)\n",
    "\n",
    "output = Dense(2048, activation='sigmoid')(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "0d0ad530",
   "metadata": {},
   "outputs": [],
   "source": [
    "rnn = Model(input_to_rnn, output, name = 'rnn')\n",
    "rnn.load_weights(os.getcwd() + \"\\\\models\\\\rnn.h5\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "942def50",
   "metadata": {},
   "source": [
    "## Data Processor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "0d7fd702",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train_data_controller.z already exists, if this notebook is run to completion the old data will be replaced.\n",
      "C:\\Users\\mattp\\Documents\\CoemergenceLab\\WorldModel_self\\images\n",
      "FOLDER:  2021-02-27\n",
      "FOLDER:  2021-03-01\n",
      "FOLDER:  2021-03-03\n",
      "FOLDER:  2021-03-04\n",
      "FOLDER:  2021-03-08\n",
      "FOLDER:  2021-03-09\n",
      "FOLDER:  2021-03-10\n",
      "FOLDER:  2021-06-21\n",
      "FOLDER:  2021-06-23\n",
      "FOLDER:  2021-06-25\n",
      "FOLDER:  2021-06-28\n",
      "FOLDER:  2021-06-30\n",
      "FOLDER:  2021-07-01\n",
      "FOLDER:  2021-07-02\n",
      "images processed: 754\n"
     ]
    }
   ],
   "source": [
    "# puts all images in list\n",
    "os.chdir(\"images\")\n",
    "\n",
    "# constant for sizing\n",
    "IMG_SIZE = 128\n",
    "\n",
    "if os.path.exists(\"train_data_controller.z\"):\n",
    "    print(\"train_data_controller.z already exists, if this notebook is run to completion the old data will be replaced.\")\n",
    "\n",
    "# put images here\n",
    "data = []\n",
    "\n",
    "# loop for loading images\n",
    "path = os.getcwd()\n",
    "print(path)\n",
    "count = 0\n",
    "for folder in os.listdir(path):\n",
    "    if \"train_data\" in folder:  # skips any train data files, as that should be the only non-directory item in images\n",
    "        continue\n",
    "    print(\"FOLDER: \",folder)\n",
    "    # added + \"/\" + to below to make it work\n",
    "    for filename in os.listdir(path + \"/\" + folder):\n",
    "        # changed to NEF (That's what I have the images saved as, may need to change back to JPG in future)\n",
    "        if(\".NEF\" in filename):\n",
    "            # added slash here too\n",
    "            temp_path = path + \"/\" + folder + \"/\" + filename\n",
    "            try:\n",
    "                img_array = cv2.imread(temp_path)\n",
    "                img_array = cv2.cvtColor(img_array, cv2.COLOR_BGR2GRAY)\n",
    "                img_array = cv2.resize(img_array, (IMG_SIZE, IMG_SIZE))\n",
    "                data.append(img_array)\n",
    "                count += 1\n",
    "\n",
    "            except Exception as e:\n",
    "                pass\n",
    "print(\"images processed:\", count)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "533f686f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "128\n",
      "(128, 128, 1)\n"
     ]
    }
   ],
   "source": [
    "# some reshaping so the images can be encoded\n",
    "\n",
    "train_data = data\n",
    "train_data = np.array(train_data).reshape(-1, IMG_SIZE, IMG_SIZE, 1)\n",
    "train_data = train_data/255.0\n",
    "print(train_data.shape[2])\n",
    "\n",
    "# Reshape \n",
    "img_width  = train_data.shape[1]\n",
    "img_height = train_data.shape[2]\n",
    "num_channels = 1\n",
    "x_train = train_data.reshape(train_data.shape[0], img_height, img_width, num_channels)\n",
    "\n",
    "input_shape = (img_height, img_width, num_channels)\n",
    "print(input_shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "e4b7a257",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\mattp\\anaconda3\\lib\\site-packages\\keras\\engine\\training_v1.py:2079: UserWarning: `Model.state_updates` will be removed in a future version. This property should not be used in TensorFlow 2.0, as `updates` are applied automatically.\n",
      "  updates=self.state_updates,\n"
     ]
    }
   ],
   "source": [
    "# gets the z vectors by running training_data (the images) through the encoder\n",
    "# note: definiately not the most efficient way to do this\n",
    "z_vals = []\n",
    "for img in train_data:\n",
    "    z_vals.append(encoder.predict(img[None,:,:,:]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "934a4388",
   "metadata": {},
   "outputs": [],
   "source": [
    "# gets the z' vectors by running z_vals (the latent vectors) through the RNN\n",
    "# again, efficiency could be better\n",
    "zprime_vals = []\n",
    "for z in z_vals:\n",
    "    zprime_vals.append(rnn.predict(z[None,:,:]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "c15f3baa",
   "metadata": {},
   "outputs": [],
   "source": [
    "# gets the action taken\n",
    "# this code currently just labels the images as 22 so I can make the skeleton of the controller\n",
    "#    Will need to update this to properly consider actions in the future, see readme for notes\n",
    "actions = []\n",
    "for z, zprime in zip(z_vals, zprime_vals):\n",
    "    actions.append([22])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "13623227",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[array([[-0.34742382,  1.4214687 ,  0.29589504, ...,  0.7067528 ,\n",
      "         0.725762  , -0.3477292 ]], dtype=float32), array([[[1.1670291e-03, 5.1805377e-04, 9.6488714e-02, ...,\n",
      "         2.3236513e-02, 9.7199380e-01, 1.3059974e-03]]], dtype=float32), [22]]\n"
     ]
    }
   ],
   "source": [
    "# makes list of lists for use in training, where each entry is a z, z', action tuple\n",
    "final_data = []\n",
    "for z, zprime, action in zip(z_vals, zprime_vals, actions):\n",
    "    final_data.append([z, zprime, action])\n",
    "print(final_data[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "d0e2e412",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "C:\\Users\\mattp\\Documents\\CoemergenceLab\\WorldModel_self\\images\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "['train_data_controller.z']"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# finished, now zip the three lists and save the data in the images directory\n",
    "print(os.getcwd())\n",
    "joblib.dump(final_data, \"train_data_controller.z\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "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
}