create-figure-db.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import warnings; warnings.filterwarnings('ignore')\n",
    "from keras.preprocessing.image import load_img, save_img, img_to_array, array_to_img\n",
    "from keras.applications import imagenet_utils, VGG19 #Xception, InceptionV3, MobileNetV2\n",
    "from scipy.misc import imread, imsave, imresize\n",
    "from uuid import NAMESPACE_URL, uuid1, uuid3\n",
    "import matplotlib.patches as patches\n",
    "from keras.models import load_model\n",
    "from skimage.color import rgb2gray\n",
    "from traceback import format_exc\n",
    "import matplotlib.pyplot as plt\n",
    "from functools import lru_cache\n",
    "from annoy import AnnoyIndex\n",
    "from six import string_types\n",
    "from random import randint\n",
    "try:\n",
    "  import keras.backend as K\n",
    "except:\n",
    "  from tensorflow.keras import backend as K\n",
    "from copy import deepcopy\n",
    "from hashlib import sha1\n",
    "from sys import maxsize\n",
    "from math import ceil\n",
    "import numpy as np\n",
    "import skimage\n",
    "import glob, json, shears, os, shutil"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This pipeline is founded on the idea that each image is composed of 0 or more \"figures\",\n",
    "where each figure is some semantic unit within the image (e.g. a single plant, or part of a plant)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Helpers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "##\n",
    "# NN MODEL\n",
    "##\n",
    "\n",
    "# VGG16, VGG19, and ResNet take 224×224 images; InceptionV3 and Xception take 299×299 inputs\n",
    "def resize_array(arr, shape=(304,304)):\n",
    "  '''Resize an array to a new shape'''\n",
    "  im = array_to_img(arr)\n",
    "  resized = im.resize(shape)\n",
    "  return img_to_array(resized)/255.0\n",
    "\n",
    "\n",
    "def get_uuid(*args, dtype='str', max_size=maxsize, right_pad=True):\n",
    "  '''Helper method that returns a random integer'''\n",
    "  if args and isinstance(args[0], string_types):\n",
    "    if dtype == 'int': i = str(get_int_hash(args[0]))\n",
    "    elif dtype == 'str': i = str(uuid3(NAMESPACE_URL, args[0]))\n",
    "  else:\n",
    "    if dtype == 'int': i = str(randint(0, max_size))\n",
    "    elif dtype == 'str': i = str(uuid1())\n",
    "  if dtype == 'int':\n",
    "    if right_pad:\n",
    "      while len(i) < len(str(max_size)): i = i + '0'\n",
    "    return int(i)\n",
    "  return i\n",
    "\n",
    "\n",
    "def get_int_hash(s):\n",
    "  '''Given a string return an integer hash of the string's content'''\n",
    "  try:\n",
    "    return int(sha1(args[0]).hexdigest(), 16)\n",
    "  except:\n",
    "    return int(sha1(args[0].encode('utf8')).hexdigest(), 16)\n",
    "\n",
    "\n",
    "##\n",
    "# IMAGE CLEANING\n",
    "##\n",
    "\n",
    "@lru_cache(maxsize=0)\n",
    "def load_image(path, color_mode='rgb'):\n",
    "  '''Given the path to an image return a keras image object'''\n",
    "  # nb all image loading should call this method rather than load_img\n",
    "  im = img_to_array(load_img(path, color_mode=color_mode))/255.0\n",
    "  if color_mode == 'grayscale': im = img_to_grayscale(im)\n",
    "  im = resize_image(im)\n",
    "  return im\n",
    "\n",
    "\n",
    "def resize_image(im, width=1024):\n",
    "  '''Given an image array, resize that image to a target width while maintaining aspect ratio'''\n",
    "  h,w = im.shape[:2]\n",
    "  height = h/w*width\n",
    "  resized = img_to_array( array_to_img(im).resize( (int(width), int(height) )) )/255.0\n",
    "  return resized\n",
    "  \n",
    "  \n",
    "def clean_image(im):\n",
    "  '''Given an image array return an image array that has standard width and non-significant pixels removed'''\n",
    "  im = resize_image(im)\n",
    "  _im = shears.remove_dominant_colors(im, mask_size=0.05, n_colors_to_remove=2)\n",
    "  _im = shears.filter_img(_im, min_size=1200, connectivity=60) # binarized mask-like\n",
    "  # create a mask from the remaining dark pixels in _im\n",
    "  im[np.where(_im==1)] = 1\n",
    "  return im\n",
    "\n",
    "\n",
    "##\n",
    "# IMAGE PARTITIONING\n",
    "##\n",
    "\n",
    "@lru_cache(maxsize=0)\n",
    "def get_image_figures(img_path, grayscale=True, resize_figures=True, min_area=10000):\n",
    "  '''\n",
    "  Given the path to an image, return an iterable of numpy arrays, with one\n",
    "  for each \"figure\" in the image, where a figure represents a single\n",
    "  object that's figured, or represented (like a plant, or planet)\n",
    "  ''' \n",
    "  orig = resize_image(load_image(img_path)) # get image\n",
    "  gray = skimage.color.rgb2gray(orig) # grayscale\n",
    "  o = np.zeros(gray.shape) # binarize\n",
    "  o[np.where(gray>skimage.filters.threshold_otsu(gray))] = 1 # binarize\n",
    "  o = shears.filter_img(o, min_size=500, connectivity=60) # remove text\n",
    "  o = 1-o # reverse figure / ground\n",
    "  l = [] # initialize list that will hold extracted figures\n",
    "  bboxes = [] # initialize list that will hold bounding boxes\n",
    "  for i in skimage.measure.regionprops(skimage.measure.label(o, connectivity=2)): # find figures\n",
    "    if i.area >= min_area:\n",
    "      y_min, x_min, y_max, x_max = i.bbox\n",
    "      # quick gut check to make sure this patch is not abnormally shaped\n",
    "      w = x_max-x_min\n",
    "      h = y_max-y_min\n",
    "      if (w/h<0.2) or (h/w<0.2): continue\n",
    "      img = img_to_grayscale(orig) if grayscale else orig\n",
    "      # add the figure to the array of figures\n",
    "      l.append(img[y_min:y_max, x_min:x_max])\n",
    "      # get the positions of the bounding box for this figure 0:1 in w,h\n",
    "      height, width = o.shape\n",
    "      bboxes.append([y_min/height, y_max/height, x_min/width, x_max/width])\n",
    "  if resize_figures:\n",
    "    return bboxes, np.array([resize_array(i) for i in l])\n",
    "  return bboxes, np.array([i for i in l])\n",
    "\n",
    "\n",
    "@lru_cache(maxsize=0)\n",
    "def get_image_windows(img_path, color_mode='rgb'):\n",
    "  '''Given the path to an image, return an array of the image windows within that image'''\n",
    "  img = load_image(img_path, color_mode=color_mode)\n",
    "  arr = img_to_array(img)/255.0\n",
    "  windows = get_windows(arr)\n",
    "  return windows\n",
    "\n",
    "\n",
    "def get_windows(im, step=20, win=300, n_per_side=20, plot=False):\n",
    "  '''\n",
    "  Given a np array `im` with at least two dimensions, return windows of size (`win`, `win`).\n",
    "  To create the list of windows, remove text, binarize, then use that binarized result as a \n",
    "  mask to remove non-significant pixels from the original image. Finally, pass a sliding window\n",
    "  over the remainder. Only retain the subset of those windows that contain some minimum amount of\n",
    "  pixels in the mask.\n",
    "  '''\n",
    "  s = win+(step*n_per_side)\n",
    "  im = resize_array(im, shape=(s,s))\n",
    "  windows = []\n",
    "  dx = 0\n",
    "  dy = 0\n",
    "  for _ in range(im.shape[0]//step): # x axis pass\n",
    "    for _ in range(im.shape[1]//step): # y axis pass\n",
    "      w = im[dy:dy+win, dx:dx+win]\n",
    "      if w.shape == (win, win, 3):\n",
    "        windows.append(w)\n",
    "      if plot:\n",
    "        plt.imshow(w)\n",
    "        plt.title('{}-{}'.format(dx, dy))\n",
    "        plt.show()\n",
    "      dy += step\n",
    "    dx += step\n",
    "  return np.array(windows)\n",
    "\n",
    "\n",
    "@lru_cache(maxsize=0)\n",
    "def guid_to_vecs(guid):\n",
    "  '''Given a guid for a full image, get the array of vectors that corresponds to that image'''\n",
    "  vec_path = os.path.join(out_dir, '{}-vectors.npy'.format(guid))\n",
    "  try:\n",
    "    vecs = np.load(vec_path)\n",
    "    if np.any(vecs): return vecs\n",
    "  except:\n",
    "    print(' ! could not load vecs for guid', guid)\n",
    "  return []\n",
    "\n",
    "\n",
    "@lru_cache(maxsize=0)\n",
    "def guid_to_vec(guid):\n",
    "  '''Given a guid for a figure, get the vector that corresponds to that figure'''\n",
    "  m = meta_d[guid]\n",
    "  return guid_to_vecs(m['parent_guid'])[m['vec_idx']]\n",
    "\n",
    "##\n",
    "# PLOTING\n",
    "##\n",
    "\n",
    "def plot_img_grid(img_list, rows=1, size_scalar=2, labels=[]):\n",
    "  '''Plot a list of images with `rows` rows'''\n",
    "  cols = ceil( len(img_list) / rows )\n",
    "  # initialize the plot\n",
    "  fig, axes = plt.subplots(rows, cols, figsize=(cols*size_scalar, rows*size_scalar), squeeze=False)\n",
    "  for idx, i in enumerate(img_list):\n",
    "    col = idx%cols\n",
    "    row = idx//cols\n",
    "    axes[row][col].imshow(img_list[idx])\n",
    "    if len(labels) > idx:\n",
    "      axes[row][col].set_title(labels[idx], fontsize=6.5)\n",
    "  plt.show()\n",
    "\n",
    "  \n",
    "def plot_guid(guid, title=None):\n",
    "  '''Given a GUID plot the image represented by that GUID'''\n",
    "  window = guid_to_img(guid)\n",
    "  title = title if title else guid_to_title(guid)\n",
    "  plt.imshow(window)\n",
    "  plt.title(title)\n",
    "  plt.show()\n",
    "  \n",
    "  \n",
    "def guid_to_img(guid):\n",
    "  '''Given a GUID return a numpy image'''\n",
    "  if meta_d['guid'].get('parent_guid', False):\n",
    "    parent_guid = meta_d['guid']['parent_guid']\n",
    "    windows = get_image_figures(meta_d[parent_guid]['path'])\n",
    "    return windows[guid][meta_d['guid']['vec_idx']]\n",
    "  else:\n",
    "    return img_to_arry(load_image(meta_d[guid]['path']))\n",
    "  \n",
    "  \n",
    "def guid_to_title(guid):\n",
    "  '''Given a GUID return a title for the image'''\n",
    "  m = meta_d[guid]\n",
    "  if not m.get('parent_guid', False):\n",
    "    return m.get('path')\n",
    "  return '{} {} {}'.format(m['collection'], m['idx_in_collection'], m['vec_idx'])\n",
    "\n",
    "\n",
    "def img_to_grayscale(im):\n",
    "  '''Given an image with shape h,w,3, return the image with the same shape in grayscale'''\n",
    "  im = rgb2gray(im).squeeze()\n",
    "  _im = np.zeros((im.shape[0], im.shape[1], 3))\n",
    "  # replicate identical color information across all channels\n",
    "  _im[:,:,0] = im\n",
    "  _im[:,:,1] = im\n",
    "  _im[:,:,2] = im\n",
    "  return _im"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load Images\n",
    "\n",
    "Each image is a jpg file that may contain 0 or more figures. This pipeline starts by loading all image files to be processed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "images = glob.glob('utils/voynichese/images/*.jpg')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "images += glob.glob('data/*/images/*.jpg') # load other collections... "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# to expedite research efforts, it's helpful to limit the number of images to be processed.\n",
    "# to do so, set the following conditional to True\n",
    "# the np.random.seed makes the random image selection reproducible\n",
    "if False:\n",
    "  np.random.seed(0)\n",
    "  np.random.shuffle(images)\n",
    "  images = images[:20]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create a 2D list (list of lists) wherein each sublist contains all figures for the ith image in `images`\n",
    "# create a 2D list of bounding boxes as well, where each sublist contains the y_min, y_max, x_min, x_max\n",
    "# coordinates in 0:1 space for the ith image in `images`\n",
    "figures = []\n",
    "bboxes = []\n",
    "for i in images:\n",
    "  try:\n",
    "    bbox, figure = get_image_figures(i)\n",
    "    bboxes.append(bbox)\n",
    "    figures.append(figure)\n",
    "  except:\n",
    "    print(' ! could not process', i)\n",
    "    bboxes.append([])\n",
    "    figures.append([])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot a histogram of the number of figures in each image\n",
    "plt.hist([len(i) for i in figures])\n",
    "plt.title('Counts of Figures Per Input Image')\n",
    "plt.xlabel('Number of Figures')\n",
    "plt.ylabel('Count')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To allow us to map figures back to their parent images in what follows, the next codeblock flattens `figures` to a dictionary that has keys in the form of `a-b`, where `a` denotes index position of a figure's parent image among all images, and `b` denotes the index of the figure within all figures in that image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig_map = {} # d[figure index] = [[image index among all images, figure index within image], [], ...]\n",
    "path_map = {}\n",
    "k = 0\n",
    "for idx, i in enumerate(figures):\n",
    "  for jdx, j in enumerate(i):\n",
    "    fig_map[k] = '{}-{}'.format(idx, jdx)\n",
    "    path_map[k] = images[idx]\n",
    "    k += 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# flatten `figures` to an array with shape (n_figures, resized_h, resized_w, 3)\n",
    "# the final dimension indicates each image has 3 color channels, while resized_h and\n",
    "# resized_w are dictated by the resize_array() function called by get_image_figures() above\n",
    "imgs_color = np.array([j for i in figures for j in i])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Save each figure to disk\n",
    "\n",
    "Each image contains multiple figures -- save each to disk"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create the directories in which output images and resized thumbs will be stored\n",
    "for i in ['images', 'thumbs']:\n",
    "  out_dir = os.path.join('output', i)\n",
    "  if not os.path.exists(out_dir):\n",
    "    os.makedirs(out_dir)\n",
    "\n",
    "# store a map from the index position of each figure among all figures to the\n",
    "# path of the image file in which that figure is illustrated\n",
    "fig_idx_to_path = {}\n",
    "\n",
    "# save the full-sized version of each figure to disk\n",
    "# his entails drawing a black box around the given figure within its parent image\n",
    "# some figures will contain the entirety of their parent image...\n",
    "for idx, i in enumerate(sorted(fig_map.keys())):\n",
    "  img_idx, fig_idx = [int(j) for j in fig_map[i].split('-')]\n",
    "  img_path = images[img_idx]\n",
    "  bn = os.path.basename(img_path)\n",
    "  bbox = bboxes[img_idx][fig_idx]\n",
    "  # read the image\n",
    "  im = imread(img_path)\n",
    "  # get image h, w\n",
    "  scale_h, scale_w = im.shape[:2]\n",
    "  # create a plot that will use the image to fill the entire frame\n",
    "  fig = plt.figure(frameon=False)\n",
    "  fig.set_size_inches(scale_w/100,scale_h/100)\n",
    "  # make the axes fill the figure\n",
    "  ax = plt.Axes(fig, [0., 0., 1., 1.])\n",
    "  ax.set_axis_off()\n",
    "  fig.add_axes(ax)\n",
    "  # add the image to the plot\n",
    "  ax.imshow(im)\n",
    "  # determine the bounding box size for this figure\n",
    "  y1, y2, x1, x2 = bbox\n",
    "  w = (x2-x1) * scale_w\n",
    "  h = (y2-y1) * scale_h\n",
    "  x = x1 * scale_w\n",
    "  y = y1 * scale_h\n",
    "  # draw a rectangle indicating figure size on parent image\n",
    "  rect = patches.Rectangle((x,y),w,h,linewidth=1,edgecolor='black',facecolor='none')\n",
    "  # add the patch to the Axes\n",
    "  ax.add_patch(rect)\n",
    "  # save the resulting image with box to disk  \n",
    "  out_file = bn.replace('.jpg', '-' + str(fig_idx) + '.jpg')\n",
    "  out_path = os.path.join('output', 'images', out_file)\n",
    "  plt.savefig(out_path)\n",
    "  if idx>0 and ((idx+1)%10)==0: print(' * saved', idx+1, 'of', len(fig_map))\n",
    "  plt.close()\n",
    "  # store the mapping from this figure's index position to the figure's representation\n",
    "  # in a file on disk\n",
    "  fig_idx_to_path[idx] = out_path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def resize_img(image_path, min_width=200, min_height=160):\n",
    "  '''\n",
    "  Given the path to an image, resize it to a minimum width and height, then\n",
    "  return that resized image\n",
    "  '''\n",
    "  img = imread(image_path)\n",
    "  shape = img.shape\n",
    "  if len(shape) == 2:\n",
    "    shape = (shape[0], shape[1], 1)\n",
    "  h, w, c = shape # height, width, color\n",
    "  if h > w:\n",
    "    height = int(h * (min_width / w))\n",
    "    width = min_width\n",
    "  else:\n",
    "    height = min_height\n",
    "    width = int(w * (min_height / h))\n",
    "  resized = imresize(img, (height, width))\n",
    "  return resized\n",
    "\n",
    "# resize each image we just created in './images' to a smaller size within './thumbs'\n",
    "for i in glob.glob(os.path.join('output', 'images', '*')):\n",
    "  bn = os.path.basename(i) \n",
    "  imsave(os.path.join('output', 'thumbs', bn), resize_img(i))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Clear up memory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import gc\n",
    "\n",
    "# to free up some memory, we can delete objects that are no longer used,\n",
    "# then evoke the garbage collector `gc`\n",
    "del figures\n",
    "\n",
    "# one can also clear the cache on lru decorated functions to free up memory\n",
    "load_image.cache_clear()\n",
    "get_image_figures.cache_clear()\n",
    "get_image_windows.cache_clear()\n",
    "\n",
    "gc.collect()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Image Vectorization\n",
    "\n",
    "Next let's transform each image into a vector. This notebook contains three methods we can use to achieve this goal:\n",
    "\n",
    "1. We can flatten each input image into a vector\n",
    "2. We can sample from a hidden layer in a pretrained network\n",
    "3. We can build our own network and vectorize images with that network"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# identify the vectorization method to use. Valid options are 'raw', 'pretrained', or 'autoencoder'\n",
    "vectorization_method = 'pretrained'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Raw Image Vectorization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# simply use the raw image objects as image vectors\n",
    "# this will mean using pixel-level deltas between images\n",
    "# to identify similar images\n",
    "# (this is the least sophisticated approach to image vectorization, but\n",
    "# it clarifies what is happening when we use more complex methods to vectorize\n",
    "# images below)\n",
    "if vectorization_method == 'raw':\n",
    "  z = X = np.array([i.flatten() for i in imgs_color])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Vectorize with Pretrained Network"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if vectorization_method == 'pretrained':\n",
    "  # specify the pretrained model to use and identify the size of its internal layers\n",
    "  # note that different models are optimized for different input image shapes\n",
    "  #model = Xception()\n",
    "  #model = MobileNetV2()\n",
    "  model = VGG19()\n",
    "  \n",
    "  # indicate the size of each internal hidden layer in the selected network\n",
    "  for idx, i in enumerate(model.layers):\n",
    "    shape = i.output_shape[1:]\n",
    "    vals = np.prod(shape)\n",
    "    print(idx, vals)\n",
    "    \n",
    "  # VGG16, VGG19, and ResNet take 224×224 images; InceptionV3 and Xception take 299×299 inputs\n",
    "  X = np.array([resize_array(i, shape=(224,224)) for i in imgs_color])\n",
    "  # only xception requires img preprocessing\n",
    "  if model.name == 'xception': X = imagenet_utils.preprocess_input(X)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if vectorization_method == 'pretrained':\n",
    "  # select the index of the layer to use when vectorizing images\n",
    "  # nb: different layers have different shapes! The layer from which one\n",
    "  # samples has a huge influence on the kinds of image similarities detected\n",
    "  # by the model\n",
    "  sample_layer_index = 6\n",
    "\n",
    "  # create a vectorization function that extracts the output from the `sample_layer_index` layer\n",
    "  vectorize = K.function([model.input], [model.layers[sample_layer_index].output])\n",
    "\n",
    "  # to process the entire dataframe in one shot, set the following to True\n",
    "  # and flip the boolean immediately below to False\n",
    "  if False:\n",
    "    z = vectorize(X)\n",
    "    z = np.array([i.flatten() for i in z])\n",
    "\n",
    "  # to process the dataframe figures one by one and log the progress as you go,\n",
    "  # set the following to True and the boolean above to False\n",
    "  if True:\n",
    "    vecs = []\n",
    "    for idx, i in enumerate(X):\n",
    "      query_img = np.expand_dims(i, 0)\n",
    "      vec = vectorize([query_img])[0]\n",
    "      vecs.append(vec)\n",
    "      print(' * vectorized', idx+1, 'of', len(X), 'frames')\n",
    "    z = np.array([i.flatten() for i in vecs])\n",
    "    \n",
    "  # scale X 0:1 so images can be visualized below\n",
    "  X = rgb2gray(X)\n",
    "  X = X-np.min(X)\n",
    "  X /= np.max(X)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# optionally, to avoid having to recompute z, X, images, and fig_map, one can cache\n",
    "# these data objects by setting the following boolean to True\n",
    "if vectorization_method == 'pretrained' and False:\n",
    "  np.save('pretrained-z',z)\n",
    "  np.save('pretrained-X',X)\n",
    "  np.save('pretrained-images',images)\n",
    "  json.dump(fig_map, open('pretrained-fig-map.json', 'w'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Vectorize with a Custom CNN"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Here we'll compose a custom Convolutional Autoencoder to vectorize our input images\n",
    "from keras.models import Model\n",
    "from keras.layers import Input, Reshape, Dense, Flatten, Conv2D, MaxPooling2D, UpSampling2D\n",
    "\n",
    "conv1 = (6,6)\n",
    "conv2 = (5,5)\n",
    "conv3 = (4,4)\n",
    "conv4 = (4,4)\n",
    "\n",
    "pool1 = (2,2)\n",
    "pool2 = (2,2)\n",
    "pool3 = (2,2)\n",
    "pool4 = (2,2)\n",
    "\n",
    "class Autoencoder:\n",
    "  def __init__(self, img_shape=(304, 304, 1)):\n",
    "    if not img_shape: raise Exception('Please provide img_shape (height, width) in px')\n",
    "\n",
    "    # create the encoder\n",
    "    i = h = Input(img_shape) # the encoder takes as input images    \n",
    "    h = Conv2D(16, conv1, activation='relu', padding='same')(i)\n",
    "    h = MaxPooling2D(pool1, padding='same')(h)\n",
    "    h = Conv2D(8, conv2, activation='relu', padding='same')(h)\n",
    "    h = MaxPooling2D(pool2, padding='same')(h)\n",
    "    h = Conv2D(8, conv3, activation='relu', padding='same')(h)\n",
    "    h = MaxPooling2D(pool3, padding='same')(h)\n",
    "    h = Conv2D(4, conv4, activation='relu', padding='same')(h)\n",
    "    h = MaxPooling2D(pool4, padding='same')(h)\n",
    "    self.encoder = Model(inputs=[i], outputs=[h])\n",
    "\n",
    "    # create the decoder\n",
    "    i = h = Input((19, 19, 4))\n",
    "    h = Conv2D(4, conv4, activation='relu', padding='same')(h)\n",
    "    h = UpSampling2D(pool4)(h)\n",
    "    h = Conv2D(8, conv3, activation='relu', padding='same')(h)\n",
    "    h = UpSampling2D(pool3)(h)\n",
    "    h = Conv2D(8, conv2, activation='relu', padding='same')(h)\n",
    "    h = UpSampling2D(pool2)(h)\n",
    "    h = Conv2D(16, conv1, activation='relu', padding='same')(h)\n",
    "    h = UpSampling2D(pool1)(h)\n",
    "    h = Conv2D(1, (3, 3), activation='sigmoid', padding='same')(h)\n",
    "    self.decoder = Model(inputs=[i], outputs=[h])\n",
    "\n",
    "    # combine the encoder and decoder into a full autoencoder\n",
    "    i = Input(img_shape) # take as input image vectors\n",
    "    z = self.encoder(i) # push observations into latent space\n",
    "    o = self.decoder(z) # project from latent space to feature space\n",
    "    self.model = Model(inputs=[i], outputs=[o])\n",
    "    self.model.compile(loss='mse', optimizer='adam')\n",
    "\n",
    "autoencoder = Autoencoder()\n",
    "\n",
    "# inspect the model\n",
    "autoencoder.encoder.summary()\n",
    "autoencoder.decoder.summary()\n",
    "autoencoder.model.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if vectorization_method == 'autoencoder':\n",
    "  # this autoencoder expects a single color channel (i.e. grayscale inputs)\n",
    "  X = np.expand_dims(np.array([rgb2gray(i) for i in imgs_color]), 3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if vectorization_method == 'autoencoder':\n",
    "  # train the autoencoder - for production purposes run more epochs, ~100 is a decent start.\n",
    "  # once the loss term ceases to fall much over many iterations,\n",
    "  # you can lower the learning rate and train more...\n",
    "  autoencoder.model.fit(X, X, batch_size=256, epochs=100)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = autoencoder.encoder\n",
    "\n",
    "if vectorization_method == 'autoencoder':\n",
    "  # optionally, save or load a model -- this is useful if you run lots of training\n",
    "  if True:\n",
    "    autoencoder.model.save('voynich.hdf5')\n",
    "  if False:\n",
    "    model = load_model('voynich.hdf5')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if vectorization_method == 'autoencoder':\n",
    "  # test the autoencoder's reconstruction of sample inputs\n",
    "  # the closer the outputs are to the input, the better the model\n",
    "  # has learned the data's features\n",
    "\n",
    "  # select the index of the figure to test\n",
    "  sample_idx = 131\n",
    "\n",
    "  # plot the figure\n",
    "  plt.imshow(X[sample_idx].squeeze())\n",
    "  plt.show()\n",
    "\n",
    "  # plot the model's reconstruction of the figure\n",
    "  o = model.predict(np.expand_dims(X[sample_idx], 0))\n",
    "  plt.imshow(o.squeeze())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# encode each figure - i.e. push each into the latent space\n",
    "if vectorization_method == 'autoencoder':\n",
    "  z = autoencoder.encoder.predict(X)\n",
    "  z = np.array([i.flatten() for i in z])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# this is another opportunity to cache model outputs to save compute time in future runs\n",
    "if vectorization_method == 'autoencoder':\n",
    "  np.save('autoencoder-z',z)\n",
    "  np.save('autoencoder-X',X)\n",
    "  np.save('autoencoder-images',images)\n",
    "  json.dump(fig_map, open('autoencoder-fig-map.json', 'w'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load data artifacts from cache"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# optionally, one can load saved data objects that were saved above\n",
    "vectorization_method = 'pretrained'\n",
    "\n",
    "if False:\n",
    "  z = np.load(vectorization_method + '-z.npy')\n",
    "  X = np.load(vectorization_method + '-X.npy')\n",
    "  images = np.load(vectorization_method + '-images.npy')\n",
    "  fig_map = json.load(open(vectorization_method + '-fig-map.json'))\n",
    "  fig_map = {int(k): v for k,v in fig_map.items()}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Metadata Prep\n",
    "\n",
    "In what follows, we'll create the MongoDB database tables for the Neural Neighbors application. This database uses two tables. The records in those tables take the forms illustrated below:\n",
    "\n",
    "<b>metadata</b>:\n",
    "```\n",
    "\"permalink\" : \"http://cdi.uvm.edu/.../datastream/JPG/view\",\n",
    "\"collection\" : \"UVM Italian Herbal\",\n",
    "\"date\" : 1550\"\",\n",
    "\"title\" : \"http://cdi.uvm.edu/.../datastream/TN/view\",\n",
    "\"location\" : \"UVM\",\n",
    "\"artist\" : \"Artist's name would go here\",\n",
    "\"has_knn\" : 1, # bool showing whether image has knn from another collection\n",
    "\"image\" : \"015ff675-de31-11e9-a0eb-f45c89b66fa9.jpg\",\n",
    "\"text\" : \"Metadata for this page would go here\"\n",
    "```\n",
    "\n",
    "<b>knn</b>:\n",
    "```\n",
    "\"image\" : \"015ff675-de31-11e9-a0eb-f45c89b66fa9.jpg\",\n",
    "\"knn\": [list of metadata objects, each augmented with a \"similarity\" float value]\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# let's load all extant metadata to prepare to save data in this format.\n",
    "img_to_meta = {}\n",
    "for i in images:\n",
    "  bn = os.path.basename(i)\n",
    "  meta_path = i.replace('/images/', '/metadata/').replace('.jpg', '.json')\n",
    "  j = json.load(open(meta_path))\n",
    "  img_to_meta[ bn ] = j"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from copy import deepcopy\n",
    "\n",
    "# map each figure to the metadata for its parent image\n",
    "fig_to_meta = {} # d[figure image basename] = {metadata for figure}\n",
    "for i in glob.glob('output/images/*.jpg'):\n",
    "  bn = os.path.basename(i)\n",
    "  parent_bn = '-'.join(bn.split('-')[:-1]) + '.jpg'\n",
    "  fig_to_meta[bn] = deepcopy(img_to_meta[parent_bn])\n",
    "  fig_to_meta[bn].update({\n",
    "    'image': bn,\n",
    "    'title': bn,\n",
    "  })"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Identify KNN & Store in DB\n",
    "\n",
    "Whichever vectorization method we use, let's now identify the k-nearest neighbors for each of our figures."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# z.shape = (number of figures, dimensions in each figure)\n",
    "z.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Build a KNN index\n",
    "from annoy import AnnoyIndex\n",
    "\n",
    "# if True will build a new KNN index, else we'll load from cache\n",
    "build = True\n",
    "\n",
    "# increasing n_trees takes more time and memory but produces better results\n",
    "n_trees = 1000\n",
    "\n",
    "if build:\n",
    "  t = AnnoyIndex(len(z[0].flatten()), 'angular')\n",
    "  for idx, i in enumerate(z):\n",
    "    t.add_item(idx, i.flatten())\n",
    "  _ = t.build(n_trees)\n",
    "  t.save('voynich.ann')\n",
    "  \n",
    "else:\n",
    "  t = AnnoyIndex(len(z[0].flatten()), 'angular')\n",
    "  model = t.load('voynich.ann')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# number of nearest neighbors to find (nb: the first retrieved knn should be the query image itself)\n",
    "k = 100\n",
    "\n",
    "# bool indicating whether to plot the knn for each image\n",
    "plot = False\n",
    "\n",
    "# minimum similarity two images must have to be called a match for out of collection knn identification\n",
    "min_sim = 0.4\n",
    "\n",
    "# boolean indicating whether to retain out of collection matches in knn\n",
    "\n",
    "# d[img_basename] = [[img_basename, sim], [img_basename, sim], ...]\n",
    "knn_map = {}\n",
    "\n",
    "for i in fig_idx_to_path:\n",
    "  bn = os.path.basename(fig_idx_to_path[i])\n",
    "  # store knn mapping data\n",
    "  knn, dist = t.get_nns_by_item(i, k, include_distances=True)\n",
    "  # convert distances to similarities\n",
    "  sims = [round(1-i, 4) for i in dist]\n",
    "  # get the figure indices of the matching knn\n",
    "  bns = [os.path.basename(fig_idx_to_path[j]) for j in knn]\n",
    "  # determine whether there are any out of collection matches with sufficiently high similarity\n",
    "  collections = [fig_to_meta[j]['collection'] for j in bns]\n",
    "  q_collection = fig_to_meta[bn]['collection']\n",
    "  # only retain out of collection matches above a certain threshold in the knn\n",
    "  filtered_knn = []\n",
    "  filtered_sims = []\n",
    "  for idx, i in enumerate(knn):\n",
    "    if sims[idx]>min_sim and collections[idx] != q_collection:\n",
    "      filtered_knn.append(knn[idx])\n",
    "      filtered_sims.append(sims[idx])\n",
    "  ooc_matches = 1 if any(filtered_knn) else 0\n",
    "  # update the fig_to_meta map to indicate whether this img has ooc matches\n",
    "  fig_to_meta[bn].update({'has_ooc_math': ooc_matches})\n",
    "  # update the fig_to_meta map to indicate whether this image has knn\n",
    "  fig_to_meta[bn].update({'has_knn': ooc_matches})\n",
    "  # augment the knn map to indicate nearest matches and sims for `i`\n",
    "  knn_map[bn] = list(zip(filtered_knn[:20], filtered_sims[:20]))\n",
    "  # plot the knn if requested\n",
    "  if plot:\n",
    "    # plot the knn for `i`\n",
    "    img_list = [X[j].squeeze() for j in knn] # show the most similar image to the input img\n",
    "    # curate titles for each match\n",
    "    titles = []\n",
    "    img_idx, fig_idx = zip(*[[int(j) for j in fig_map[k].split('-')] for k in knn])  \n",
    "    # display the figure offset, similarity, and image path for all matches\n",
    "    print('\\n'.join(['{} {} {}'.format(knn[idx], sims[idx], images[i]) for idx, i in enumerate(img_idx)]))\n",
    "    plot_img_grid(img_list, labels=sims)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Seed Database"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pymongo import MongoClient\n",
    "\n",
    "# when using the cluster, save to local disk, else save to mongo\n",
    "using_cluster = True\n",
    "\n",
    "if not using_cluster:\n",
    "  # identify the name of the database in which to save the resulting data\n",
    "  db = MongoClient().neuralneighbors\n",
    "  # remove all records from the extant db tables\n",
    "  print(db.knn.remove({}))\n",
    "  print(db.metadata.remove({}))\n",
    "\n",
    "# save each figure that has matches to the database\n",
    "saved = 0\n",
    "for i in knn_map:\n",
    "  if not knn_map[i]: continue\n",
    "  bns, sims = zip(*knn_map[i])\n",
    "  bns = [os.path.basename(fig_idx_to_path[j]) for j in bns]\n",
    "  knn_vals = [fig_to_meta[j] for j in bns]\n",
    "  for jdx, j in enumerate(knn_vals):\n",
    "    j['similarity'] = sims[jdx]\n",
    "  # save to mongo if not using cluster\n",
    "  knn_data = {'image': i, 'knn': knn_vals}\n",
    "  metadata_data = fig_to_meta[i]\n",
    "  if not using_cluster:\n",
    "    db.knn.insert(knn_data)\n",
    "    db.metadata.insert(metadata_data)\n",
    "  else:\n",
    "    if not os.path.exists('knn'): os.makedirs('knn')\n",
    "    if not os.path.exists('metadata'): os.makedirs('metadata')\n",
    "    with open(os.path.join('knn', i), 'w') as out: json.dump(knn_data, out)\n",
    "    with open(os.path.join('metadata', i), 'w') as out: json.dump(metadata_data, out)\n",
    "  saved += 1\n",
    "\n",
    "  # indicate the total number of matches found\n",
    "  print(' * found', saved, 'matches')\n",
    "  \n",
    "if not using_cluster:\n",
    "  # create indices for the tables to expedite query times\n",
    "  db.knn.create_index('image')\n",
    "  db.metadata.create_index('image')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# once the data is saved to local disk on cluster, we can use the following\n",
    "# to port the data into a mongo db (e.g. on Cuda / Heroku)\n",
    "\n",
    "db = MongoClient().neuralneighbors\n",
    "# remove all records from the extant db tables\n",
    "print(db.knn.remove({}))\n",
    "print(db.metadata.remove({}))\n",
    "\n",
    "for i in glob.glob(os.path.join('knn', '*')): db.knn.insert(json.load(open(i)))\n",
    "for i in glob.glob(os.path.join('metadata', '*')): db.metadta.insert(json.load(open(i)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# now just move ./output to the location of the neural-neighbors root directory\n",
    "# and run the Docker / Heroku commands to deploy the application!"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.5.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}