training_siamese_network.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n",
    "import random\n",
    "random.seed(42)\n",
    "\n",
    "import sklearn\n",
    "import numpy as np\n",
    "from PIL import Image\n",
    "from keras import backend as K\n",
    "from keras.callbacks import ModelCheckpoint, ReduceLROnPlateau\n",
    "from keras.layers import Input, Flatten, Dense, Dropout, Lambda\n",
    "from keras.models import Model, load_model\n",
    "from keras.optimizers import RMSprop\n",
    "from keras.applications.inception_resnet_v2 import preprocess_input, InceptionResNetV2\n",
    "\n",
    "from utilities import convert_image_to_square_rgb, preprocess_image_inception_keras, contrastive_loss"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Training Data\n",
    "Define the training data for training the siamese similarity model. Training data will be created from a JSON file labeling items and their images."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Change these values to the location of your images and label file, which should be mounted to the Docker image\n",
    "IMAGE_BASEPATH = '/data/images'\n",
    "LABELED_IMAGES_FILENAME = '/data/labeled_images.json'\n",
    "BATCH_SIZE = 32"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Don't change these unless you know what you are doing\n",
    "IMAGE_SHAPE = (299, 299, 3)\n",
    "IMAGE_SIZE = IMAGE_SHAPE[:2]\n",
    "EVAL_PERCENT = 0.3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "CACHED_IMAGES = {}\n",
    "\n",
    "def get_and_cache_image(filename, image_dir=IMAGE_BASEPATH):\n",
    "    \"\"\"Loads am image into memory, transforms it to a RGB matrix\n",
    "    and caches it for fast retrieval\n",
    "    \n",
    "    Args:\n",
    "        filename (str): Relative filepath from image_dir of image\n",
    "        image_dir (str): Directory containing all images\n",
    "        \n",
    "    Returns:\n",
    "        ndarray\n",
    "    \"\"\"\n",
    "    if image_dir:\n",
    "        filename = os.path.join(image_dir, filename)\n",
    "    if filename in CACHED_IMAGES:\n",
    "        return CACHED_IMAGES[filename]\n",
    "    else:\n",
    "        image = Image.open(filename)\n",
    "        rgb_matrix = convert_image_to_square_rgb(image, IMAGE_SIZE)\n",
    "        CACHED_IMAGES[filename] = rgb_matrix\n",
    "        return rgb_matrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def batch_generator(X_pairs, y_labels, batch_size):\n",
    "    \"\"\"\"Returns a batch of images for the model to consume\n",
    "    \n",
    "    Args:\n",
    "        X_pairs (tuple[str, str]): List of tuples with two image filepaths\n",
    "        y_labels (List[int]): List of labels for pairs, 1 for similar and 0 for dissimilar\n",
    "    \n",
    "    Yields:\n",
    "        tuple[(ndarray, ndarray), ndarray]\n",
    "    \"\"\"\n",
    "    total_pairs = len(y_labels)\n",
    "    while True:\n",
    "        batch = []\n",
    "        labels = []\n",
    "        for i, ((img_filename_1, img_filename_2), label) in enumerate(zip(X_pairs, y_labels)):\n",
    "            img_1 = get_and_cache_image(img_filename_1)\n",
    "            img_2 = get_and_cache_image(img_filename_2)\n",
    "            img_1 = preprocess_image_inception_keras(img_1)\n",
    "            img_2 = preprocess_image_inception_keras(img_2)\n",
    "            batch.append([img_1, img_2])\n",
    "            labels.append(label)\n",
    "            if (i + 1) % batch_size == 0 or (i + 1) == total_pairs:\n",
    "                result = np.array(batch, dtype=np.float32)\n",
    "                yield ([result[:, 0], result[:, 1]], np.array(labels, dtype=np.float32))\n",
    "                result = None\n",
    "                batch = []\n",
    "                labels = []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_pairs(candidates, shuffle=False):\n",
    "    \"\"\"Creates positives pairs and randomly samples negative pairs to train our model\"\"\"\n",
    "    pairs_pairs = []\n",
    "    labels_pairs = []\n",
    "    \n",
    "    for item in candidates:\n",
    "        image_count = len(item['images'])\n",
    "        if image_count < 2:\n",
    "            continue\n",
    "        for i, image_1 in enumerate(item['images']):\n",
    "            for image_2 in item['images'][i+1:]:\n",
    "                negative_item = random.choice(candidates)\n",
    "                while set(item['labels']) & set(negative_item['labels']):\n",
    "                    negative_item = random.choice(candidates)\n",
    "                negative_image = random.choice(negative_item['images'])\n",
    "                positive_pair = (image_1['filename'], image_2['filename'])\n",
    "                positive_pair = random.sample(positive_pair, 2)\n",
    "                negative_pair = (positive_pair[0], negative_image['filename'])\n",
    "                pairs_pairs.append([positive_pair, negative_pair])\n",
    "                labels_pairs.append([1, 0])\n",
    "    if shuffle:           \n",
    "        pairs_pairs, labels_pairs = sklearn.utils.shuffle(pairs_pairs, labels_pairs, random_state=42)\n",
    "\n",
    "    example_count = len(labels_pairs) * 2\n",
    "    pairs = [None] * example_count\n",
    "    labels = [None] * example_count\n",
    "    for i, (a, b) in enumerate(pairs_pairs):\n",
    "        pairs[i*2] = a\n",
    "        pairs[i*2+1] = b\n",
    "    for i, (a, b) in enumerate(labels_pairs):\n",
    "        labels[i*2] = a\n",
    "        labels[i*2+1] = b\n",
    "    return (pairs, labels)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(LABELED_IMAGES_FILENAME) as f:\n",
    "    labeled_images = json.load(f)\n",
    "\n",
    "candidates = list(labeled_images.values())\n",
    "pivot = int(len(candidates) * EVAL_PERCENT)\n",
    "random.shuffle(candidates)\n",
    "\n",
    "train_candidates = candidates[:pivot]\n",
    "eval_candidates = candidates[pivot:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X_train, y_train = create_pairs(train_candidates, shuffle=True)\n",
    "X_eval, y_eval = create_pairs(eval_candidates, shuffle=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Model Objectives\n",
    "The functions the models are trying to optimize in some way. `contrastive_loss` is the important loss function, but it must be defined in the utilities to properly load Keras models into memory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def euclidean_distance(vects):\n",
    "    x, y = vects\n",
    "    return K.sqrt(K.maximum(K.sum(K.square(x - y), axis=1, keepdims=True), K.epsilon()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def eucl_dist_output_shape(shapes):\n",
    "    shape1, shape2 = shapes\n",
    "    return (shape1[0], 1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Model Definition"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_base_network(input_shape, freeze_layers_until=None):\n",
    "    \"\"\"Get the base network to do the feature extract for the latent embedding\n",
    "    \n",
    "    Args:\n",
    "        input_shape (tuple): Shape of image tensor input\n",
    "        \n",
    "    Returns:\n",
    "        keras.models.Model\n",
    "    \"\"\"\n",
    "    input = Input(shape=input_shape)\n",
    "    inception = InceptionResNetV2(weights='imagenet', input_tensor=input)\n",
    "    inception.layers.pop()  # Remove classification layer\n",
    "\n",
    "    if freeze_layers_until:\n",
    "        assert freeze_layers_until in [l.name for l in inception.layers]\n",
    "        for layer in inception.layers:\n",
    "            layer.trainable = False\n",
    "            if type(layer) == 'BatchNormalization':\n",
    "                layer.momentum = 1.0\n",
    "            if layer.name == freeze_layers_until:\n",
    "                break\n",
    "\n",
    "    model = Model(inputs=[input], outputs=[inception.layers[-1].output], name='embedding_model')\n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_accuracy(y_true, y_pred):\n",
    "    \"\"\"Compute classification accuracy with a fixed threshold on distances.\"\"\"\n",
    "    pred = y_pred.ravel() < 0.5\n",
    "    return np.mean(pred == y_true)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def accuracy(y_true, y_pred):\n",
    "    \"\"\"Compute classification accuracy with a fixed threshold on distances.\"\"\"\n",
    "    return K.mean(K.equal(y_true, K.cast(y_pred < 0.5, y_true.dtype)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Grab the Inception ResNet V2 model pretrained on image net\n",
    "# Freeze layers up to mixed_6a for faster training and less overfitting\n",
    "base_network = create_base_network(IMAGE_SHAPE, 'mixed_6a')\n",
    "\n",
    "# Create two inputs for both images from pairs\n",
    "input_a = Input(shape=IMAGE_SHAPE)\n",
    "input_b = Input(shape=IMAGE_SHAPE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Connect the same base_network created above to the two image inputs\n",
    "processed_a = base_network(input_a)\n",
    "processed_b = base_network(input_b)\n",
    "\n",
    "# Get the distance between the two input images\n",
    "distance = Lambda(euclidean_distance,\n",
    "                  output_shape=eucl_dist_output_shape)([processed_a, processed_b])\n",
    "\n",
    "# Create a model which takes in a pair of images and returns the distance\n",
    "model = Model([input_a, input_b], distance)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create callback functions to call after every epoch\n",
    "ckpt_dir = 'checkpoints'\n",
    "if not os.path.exists(ckpt_dir):\n",
    "    os.mkdir(ckpt_dir)\n",
    "ckpt_pattern = os.path.join(ckpt_dir, 'weights.{epoch:03d}-{val_loss:.5f}.hdf5')\n",
    "\n",
    "callbacks = [\n",
    "    ModelCheckpoint(ckpt_pattern,\n",
    "                    monitor='val_loss',\n",
    "                    save_best_only=True),\n",
    "    ReduceLROnPlateau('loss', factor=0.5, patience=3, verbose=1, min_lr=1e-7)\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "init_learning_rate = 0.0005\n",
    "rms = RMSprop(lr=init_learning_rate)\n",
    "model.compile(loss=contrastive_loss, optimizer=rms, metrics=[accuracy])\n",
    "model.fit_generator(batch_generator(X_train, y_train, BATCH_SIZE),\n",
    "                    steps_per_epoch=len(X_train) // BATCH_SIZE,\n",
    "                    epochs=5,\n",
    "                    callbacks=callbacks,\n",
    "                    validation_data=batch_generator(X_eval, y_eval, BATCH_SIZE),\n",
    "                    validation_steps=len(y_eval) // BATCH_SIZE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Save the final model\n",
    "model.save(os.path.join(ckpt_dir, 'final_model.hdf5'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "y_pred = model.predict_generator(batch_generator(X_eval, y_eval, BATCH_SIZE),\n",
    "                                 steps=len(y_eval) // BATCH_SIZE,\n",
    "                                 verbose=1)\n",
    "te_acc = compute_accuracy(y_eval, y_pred)\n",
    "print('* Accuracy on test set: %0.2f%%' % (100 * te_acc))"
   ]
  }
 ],
 "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.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}