03_Image_Search_Using_AlexNet_0.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Image Retrieval using AlexNet "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Goals :\n",
    "- Create AlexNet architecture using Pre-Trained Weights\n",
    "- Create Image Embeddings using Transfer Learning \n",
    "- Create Nearest Neighbour Algorithm \n",
    "- Give a query iamges, retrieve similar images "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "import numpy as np\n",
    "import os\n",
    "import os.path as op\n",
    "import random\n",
    "from scipy import ndimage\n",
    "from glob import glob\n",
    "import matplotlib.pyplot as plt\n",
    "from sklearn.neighbors import NearestNeighbors\n",
    "\n",
    "%matplotlib inline\n",
    "print('Tensorflow version : {0}'.format(tf.__version__))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### AlexNet Architecture"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![AlexNet](./images/alexnet_architecture.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " Weights from Michael Guerzhoy and Davi Frossard\n",
    " [http://www.cs.toronto.edu/~guerzhoy/tf_alexnet/]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Import Pre-Trained Variables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Path to the numpy weights\n",
    "alexnet_path = op.join(op.curdir, 'processed','bvlc_alexnet.npy')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# read weights in numpy\n",
    "variable_data = np.load(alexnet_path, encoding='bytes').item()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# variable_data is a dictionary\n",
    "type(variable_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# keys\n",
    "variable_data.keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Convolution layer - 1 weights\n",
    "conv1_preW = variable_data[\"conv1\"][0]\n",
    "conv1_preb = variable_data[\"conv1\"][1]\n",
    "print(conv1_preW.shape)\n",
    "print(conv1_preb.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Convolution layer - 2 weights\n",
    "conv2_preW = variable_data[\"conv2\"][0]\n",
    "conv2_preb = variable_data[\"conv2\"][1]\n",
    "print(conv2_preW.shape)\n",
    "print(conv2_preb.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Convolution layer-3 weights\n",
    "conv3_preW = variable_data[\"conv3\"][0]\n",
    "conv3_preb = variable_data[\"conv3\"][1]\n",
    "print(conv3_preW.shape)\n",
    "print(conv3_preb.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Convolution layer-4 weights\n",
    "conv4_preW = variable_data[\"conv4\"][0]\n",
    "conv4_preb = variable_data[\"conv4\"][1]\n",
    "print(conv4_preW.shape)\n",
    "print(conv4_preb.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Convolution layer-5 weights\n",
    "conv5_preW = variable_data[\"conv5\"][0]\n",
    "conv5_preb = variable_data[\"conv5\"][1]\n",
    "print(conv5_preW.shape)\n",
    "print(conv5_preb.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Fully contected layer - 1\n",
    "fc6_preW = variable_data[\"fc6\"][0]\n",
    "fc6_preb = variable_data[\"fc6\"][1]\n",
    "print(fc6_preW.shape)\n",
    "print(fc6_preb.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Fully connected layer -2 \n",
    "fc7_preW = variable_data[\"fc7\"][0]\n",
    "fc7_preb = variable_data[\"fc7\"][1]\n",
    "print(fc7_preW.shape)\n",
    "print(fc7_preb.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Fully connected layer - 3\n",
    "fc8_preW = variable_data[\"fc8\"][0]\n",
    "fc8_preb = variable_data[\"fc8\"][1]\n",
    "print(fc8_preW.shape)\n",
    "print(fc8_preb.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create the AlexNet Network"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "pixel_depth = 255.0\n",
    "resized_height = 227\n",
    "resized_width = 227\n",
    "num_channels = 3\n",
    "\n",
    "graph = tf.Graph()\n",
    "\n",
    "with graph.as_default():\n",
    "    x = tf.placeholder(tf.uint8, [None, None, None, num_channels],\n",
    "                       name='input')\n",
    "    \n",
    "    to_float = tf.cast(x, tf.float32)\n",
    "    resized = tf.image.resize_images(to_float, (resized_height, resized_width))\n",
    "    \n",
    "    # Convolution 1\n",
    "    with tf.name_scope('conv1') as scope:\n",
    "        kernel = tf.Variable(conv1_preW, name='weights')\n",
    "        biases = tf.Variable(conv1_preb, name='biases')\n",
    "        conv = tf.nn.conv2d(resized, kernel, [1, 4, 4, 1], padding=\"SAME\")\n",
    "        bias = tf.nn.bias_add(conv, biases)\n",
    "        conv1 = tf.nn.relu(bias, name=scope)\n",
    "\n",
    "    # Local response normalization 2\n",
    "    radius = 2\n",
    "    alpha = 2e-05\n",
    "    beta = 0.75\n",
    "    bias = 1.0\n",
    "    lrn1 = tf.nn.local_response_normalization(conv1,\n",
    "                                              depth_radius=radius,\n",
    "                                              alpha=alpha,\n",
    "                                              beta=beta,\n",
    "                                              bias=bias)\n",
    "\n",
    "    # Maxpool 1\n",
    "    pool1 = tf.nn.max_pool(lrn1,\n",
    "                           ksize=[1, 3, 3, 1],\n",
    "                           strides=[1, 2, 2, 1],\n",
    "                           padding='VALID',\n",
    "                           name='pool1')\n",
    "\n",
    "    # Convolution 2\n",
    "    with tf.name_scope('conv2') as scope:\n",
    "\n",
    "        kernel = tf.Variable(conv2_preW, name='weights')\n",
    "        biases = tf.Variable(conv2_preb, name='biases')\n",
    "\n",
    "        input_a, input_b = tf.split(axis=3,\n",
    "                                    num_or_size_splits=2,\n",
    "                                    value=pool1)\n",
    "        kernel_a, kernel_b = tf.split(axis=3,\n",
    "                                    num_or_size_splits=2,\n",
    "                                      value=kernel)\n",
    "\n",
    "        with tf.name_scope('A'):\n",
    "            conv_a = tf.nn.conv2d(input_a, kernel_a, [1, 1, 1, 1], padding=\"SAME\")        \n",
    "\n",
    "        with tf.name_scope('B'):\n",
    "            conv_b = tf.nn.conv2d(input_b, kernel_b, [1, 1, 1, 1], padding=\"SAME\")\n",
    "\n",
    "        conv = tf.concat(values=[conv_a, conv_b], axis=3)\n",
    "        bias = tf.nn.bias_add(conv, biases)\n",
    "        conv2 = tf.nn.relu(bias, name=scope)\n",
    "\n",
    "    # Local response normalization 2\n",
    "    radius = 2\n",
    "    alpha = 2e-05\n",
    "    beta = 0.75\n",
    "    bias = 1.0\n",
    "    lrn2 = tf.nn.local_response_normalization(conv2,\n",
    "                                              depth_radius=radius,\n",
    "                                              alpha=alpha,\n",
    "                                              beta=beta,\n",
    "                                              bias=bias)\n",
    "\n",
    "    # Maxpool 2\n",
    "    pool2 = tf.nn.max_pool(lrn2,\n",
    "                           ksize=[1, 3, 3, 1],\n",
    "                           strides=[1, 2, 2, 1],\n",
    "                           padding='VALID',\n",
    "                           name='pool2')\n",
    "\n",
    "    with tf.name_scope('conv3') as scope:\n",
    "        kernel = tf.Variable(conv3_preW, name='weights')\n",
    "        biases = tf.Variable(conv3_preb, name='biases')\n",
    "        conv = tf.nn.conv2d(pool2, kernel, [1, 1, 1, 1], padding=\"SAME\")\n",
    "        bias = tf.nn.bias_add(conv, biases)\n",
    "        conv3 = tf.nn.relu(bias, name=scope)\n",
    "\n",
    "\n",
    "    with tf.name_scope('conv4') as scope:\n",
    "\n",
    "        kernel = tf.Variable(conv4_preW, name='weights')\n",
    "        biases = tf.Variable(conv4_preb, name='biases')\n",
    "\n",
    "        input_a, input_b = tf.split(axis=3,\n",
    "                                    num_or_size_splits=2,\n",
    "                                    value=conv3)\n",
    "        kernel_a, kernel_b = tf.split(axis=3,\n",
    "                                    num_or_size_splits=2,\n",
    "                                      value=kernel)\n",
    "\n",
    "        with tf.name_scope('A'):\n",
    "            conv_a = tf.nn.conv2d(input_a, kernel_a, [1, 1, 1, 1], padding=\"SAME\")        \n",
    "\n",
    "        with tf.name_scope('B'):\n",
    "            conv_b = tf.nn.conv2d(input_b, kernel_b, [1, 1, 1, 1], padding=\"SAME\")\n",
    "\n",
    "        conv = tf.concat(values=[conv_a, conv_b], axis=3)\n",
    "        bias = tf.nn.bias_add(conv, biases)\n",
    "        conv4 = tf.nn.relu(bias, name=scope)\n",
    "\n",
    "\n",
    "    with tf.name_scope('conv5') as scope:\n",
    "\n",
    "        kernel = tf.Variable(conv5_preW, name='weights')\n",
    "        biases = tf.Variable(conv5_preb, name='biases')\n",
    "\n",
    "        input_a, input_b = tf.split(axis=3,\n",
    "                                    num_or_size_splits=2,\n",
    "                                    value=conv4)\n",
    "        kernel_a, kernel_b = tf.split(axis=3,\n",
    "                                      num_or_size_splits=2,\n",
    "                                      value=kernel)\n",
    "\n",
    "        with tf.name_scope('A'):\n",
    "            conv_a = tf.nn.conv2d(input_a, kernel_a, [1, 1, 1, 1], padding=\"SAME\")        \n",
    "\n",
    "        with tf.name_scope('B'):\n",
    "            conv_b = tf.nn.conv2d(input_b, kernel_b, [1, 1, 1, 1], padding=\"SAME\")\n",
    "\n",
    "        conv = tf.concat(values=[conv_a, conv_b],axis=3)\n",
    "        bias = tf.nn.bias_add(conv, biases)\n",
    "        conv5 = tf.nn.relu(bias, name=scope)\n",
    "\n",
    "\n",
    "    # Maxpool 2\n",
    "    pool5 = tf.nn.max_pool(conv5,\n",
    "                           ksize=[1, 3, 3, 1],\n",
    "                           strides=[1, 2, 2, 1],\n",
    "                           padding='VALID',\n",
    "                           name='pool5')\n",
    "\n",
    "    # Fully connected 6\n",
    "    with tf.name_scope('fc6'):\n",
    "        weights = tf.Variable(fc6_preW, name='fc6_weights')\n",
    "        bias = tf.Variable(fc6_preb, name='fc6_bias')\n",
    "        shape = tf.shape(pool5)\n",
    "        size = shape[1] * shape[2] * shape[3]\n",
    "        fc6 = tf.nn.relu_layer(tf.reshape(pool5, [-1, size]),\n",
    "                               weights, bias, name='relu')\n",
    "\n",
    "    # Fully connected 7\n",
    "    with tf.name_scope('fc7'):\n",
    "        weights = tf.Variable(fc7_preW, name='weights')\n",
    "        bias = tf.Variable(fc7_preb, name='bias')\n",
    "        fc7 = tf.nn.relu_layer(fc6, weights, bias, name='relu')\n",
    "\n",
    "    # Fully connected 8\n",
    "    with tf.name_scope('fc8'):\n",
    "        weights = tf.Variable(fc8_preW, name='weights')\n",
    "        bias = tf.Variable(fc8_preb, name='bias')\n",
    "        # fc8 = tf.matmul(fc7, weights) + bias\n",
    "        fc8 = tf.nn.xw_plus_b(fc7, weights, bias)\n",
    "\n",
    "    softmax = tf.nn.softmax(fc8)\n",
    "\n",
    "    init = tf.global_variables_initializer()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# visualize alexnet graph\n",
    "sess = tf.Session(graph=graph)\n",
    "sess.run(init)\n",
    "\n",
    "writer = tf.summary.FileWriter('tensorboard/alexnet', graph=graph)\n",
    "writer.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#!tensorboard --logdir='tensorboard/alexnet'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Export AlexNet Model ( Graph + Weights )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "with graph.as_default():\n",
    "    saver = tf.train.Saver()\n",
    "    alex_vars_path = op.join(op.curdir, 'saved_models','alex_vars')\n",
    "    save_path = saver.save(sess, alex_vars_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Import Graph and Weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "alex_vars_path = op.join(op.curdir, 'saved_models','alex_vars')\n",
    "alex_meta_file_path = op.join(op.curdir, 'saved_models','alex_vars.meta')\n",
    "graph = tf.Graph()\n",
    "with graph.as_default():    \n",
    "    importer = tf.train.import_meta_graph(alex_meta_file_path)\n",
    "\n",
    "sess = tf.Session(graph=graph)\n",
    "importer.restore(sess, alex_vars_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "graph.get_operations()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# extract fc7 output\n",
    "fc7_op = graph.get_operation_by_name('fc7/relu')\n",
    "fc7 = fc7_op.outputs[0]\n",
    "x = graph.get_operation_by_name('input').outputs[0]\n",
    "init = graph.get_operation_by_name('init')\n",
    "\n",
    "sess = tf.Session(graph=graph)\n",
    "sess.run(init)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "image_embedding_size = fc7.get_shape()[1]\n",
    "print('Image embedding size : {0}'.format(image_embedding_size))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Perform Image Embedding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# UTZap50K\n",
    "UTZap50K_directory = op.join(op.curdir, 'processed','utzap50k')\n",
    "\n",
    "all_files = [y for x in os.walk(UTZap50K_directory) for y in glob(os.path.join(x[0], '*.jpg'))]\n",
    "len(all_files)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# sample paths\n",
    "all_files[:5]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# show a sample image\n",
    "image = ndimage.imread(all_files[0])\n",
    "print('image shape : {0}'.format(image.shape))\n",
    "plt.imshow(image)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Number of images to build Nearest Neighbors Model\n",
    "random.shuffle(all_files)\n",
    "num_images = 2000 # increase the number of images to increase the image database\n",
    "neighbor_list = all_files[:num_images]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create empty array with shape : num_images * image_embedding_size\n",
    "image_embeddings = np.ndarray((num_images, image_embedding_size))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create image embedding for each image \n",
    "for i, filename in enumerate(neighbor_list):\n",
    "    image = ndimage.imread(filename)\n",
    "    features = sess.run(fc7, feed_dict={x: [image]})\n",
    "    image_embeddings[i:i+1] = features\n",
    "    if i % 250 == 0:\n",
    "        print(i)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# shape of image embeddings\n",
    "image_embeddings.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Find Similar Images "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "g = tf.Graph()\n",
    "NUM_IMAGES = len(neighbor_list)\n",
    "EMBEDDING_DIMENSION = 4096\n",
    "\n",
    "with g.as_default():\n",
    "    # provide input indices \n",
    "    x = tf.placeholder(shape=[None], dtype=tf.int32, name='x')\n",
    "    \n",
    "    # create a constant initializer\n",
    "    weights_initializer = tf.constant_initializer(image_embeddings)\n",
    "    embedding_weights = tf.get_variable(\n",
    "                            name='embedding_weights', \n",
    "                            shape=(NUM_IMAGES, EMBEDDING_DIMENSION), \n",
    "                            initializer=weights_initializer,\n",
    "                            trainable=False)\n",
    "    # emebedding Lookup \n",
    "    embedding_lookup = tf.nn.embedding_lookup(embedding_weights, x)\n",
    "    \n",
    "    # We use the cosine distance:\n",
    "    norm = tf.sqrt(tf.reduce_sum(tf.square(embedding_weights), 1, keepdims=True))\n",
    "    normalized_embeddings = embedding_weights / norm\n",
    "    valid_embeddings = tf.nn.embedding_lookup(normalized_embeddings, x)\n",
    "    similarity = tf.matmul(valid_embeddings, tf.transpose(normalized_embeddings))\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "NUM_QUERY_IMAGES = 10\n",
    "NUM_NEIGHBORS = 5\n",
    "query_indices = np.random.choice(range(len(neighbor_list)), NUM_QUERY_IMAGES, replace=False)\n",
    "\n",
    "with tf.Session(graph=g) as sess:\n",
    "    sess.run(tf.global_variables_initializer())\n",
    "    sim = sess.run(similarity, feed_dict={x : query_indices})\n",
    "    \n",
    "    f,ax = plt.subplots(NUM_QUERY_IMAGES,NUM_NEIGHBORS, sharex=True, sharey=True, figsize=(14,20))\n",
    "    \n",
    "    print('Shape of Similarity Matrix: {0}'.format(sim.shape))\n",
    "    for i,image_index in enumerate(query_indices):\n",
    "       \n",
    "        top_k = NUM_NEIGHBORS # number of nearest neighbors\n",
    "        nearest = (-sim[i, :]).argsort()[1:top_k+1]\n",
    "        log = 'Nearest to {0} :'.format(neighbor_list[image_index])\n",
    "        \n",
    "        for k in range(top_k):\n",
    "            close_image = ndimage.imread(neighbor_list[nearest[k]])\n",
    "            ax[i,k].imshow(close_image)\n",
    "    plt.show()\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Further Extension\n",
    "\n",
    "**Approximate Nearest Neighbor Search** \n",
    "\n",
    "Currently we are comparing each image with every other image that is not recommended for production scenarios. Use techniques such as Approximate Nearest Neighbor to accelerate the process.\n",
    "\n",
    "- You can explore ANNOY  package by Spotify [ https://github.com/spotify/annoy ]\n",
    "- You can also explore LSH ( Locality Senstivity Hashing ) [ https://graphics.stanford.edu/courses/cs468-06-fall/Slides/aneesh-michael.pdf ]"
   ]
  }
 ],
 "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.6.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}