examples/updatable_models/updatable_tiny_drawing_classifier.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Updatable Tiny Drawing Classifier\n",
    "\n",
    "This notebook creates a model which can be used to train a simple drawing / sketch classifier based on user examples. The model is a pipeline composed of a drawing embedding model and a nearest neighbor classifier. \n",
    "\n",
    "The model is updatable and starts off 'empty' in that the nearest neighbor classifier has no examples or labels. Before updating with training examples it predicts 'unknown' for all input.\n",
    "\n",
    "The input to the model is a 28 x 28 grayscale drawing. The background is expected to be black (0) while the strokes of the drawing should be rendered as white (255). For example:\n",
    "\n",
    "| Drawing of a Star | Drawing of a Heart | Drawing of 5 |\n",
    "| ----------- | ----------- | ----------- |\n",
    "| ![Star Example](images/star28x28.png) | ![Heart Example](images/heart28x28.png) | ![Five Example](images/five28x28.png) |\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Embedding\n",
    "\n",
    "Let's start by getting the first part of the model.  It is the drawing embedding model which will be used as a feature extractor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tf.estimator package not installed.\n",
      "tf.estimator package not installed.\n",
      "input {\n",
      "  name: \"drawing\"\n",
      "  shortDescription: \"Input sketch image with black background and white strokes\"\n",
      "  type {\n",
      "    imageType {\n",
      "      width: 28\n",
      "      height: 28\n",
      "      colorSpace: GRAYSCALE\n",
      "    }\n",
      "  }\n",
      "}\n",
      "output {\n",
      "  name: \"embedding\"\n",
      "  shortDescription: \"Vector embedding of sketch in 128 dimensional space\"\n",
      "  type {\n",
      "    multiArrayType {\n",
      "      shape: 128\n",
      "      dataType: FLOAT32\n",
      "    }\n",
      "  }\n",
      "}\n",
      "metadata {\n",
      "  shortDescription: \"Embeds a 28 x 28 grayscale image of a sketch into 128 dimensional space. The model was created by removing the last layer of a simple convolution based neural network classifier trained on the Quick, Draw! dataset (https://github.com/googlecreativelab/quickdraw-dataset).\"\n",
      "  author: \"Core ML Tools Example\"\n",
      "  license: \"MIT\"\n",
      "}\n",
      "\n"
     ]
    }
   ],
   "source": [
    "import coremltools\n",
    "from coremltools.models import MLModel\n",
    "\n",
    "embedding_path = './models/TinyDrawingEmbedding.mlmodel'\n",
    "embedding_model = MLModel(embedding_path)\n",
    "\n",
    "embedding_spec = embedding_model.get_spec()\n",
    "print embedding_spec.description"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can see from the description above that the embedding model takes in a 28x28 grayscale image about outputs a 128 dimensional float vector."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Nearest Neighbor Classifier\n",
    "\n",
    "Now that the feature extractor is in place, let's create the second model of our pipeline model.\n",
    "It is a nearest neighbor classifier operating on the embedding."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from coremltools.models.nearest_neighbors import KNearestNeighborsClassifierBuilder\n",
    "import coremltools.models.datatypes as datatypes\n",
    "\n",
    "knn_builder = KNearestNeighborsClassifierBuilder(input_name='embedding',\n",
    "                                                 output_name='label',\n",
    "                                                 number_of_dimensions=128,\n",
    "                                                 default_class_label='unknown',\n",
    "                                                 k=3,\n",
    "                                                 weighting_scheme='inverse_distance',\n",
    "                                                 index_type='linear')\n",
    "\n",
    "knn_builder.author = 'Core ML Tools Example'\n",
    "knn_builder.license = 'MIT'\n",
    "knn_builder.description = 'Classifies 128 dimension vector based on 3 nearest neighbors'\n",
    "\n",
    "knn_spec = knn_builder.spec\n",
    "knn_spec.description.input[0].shortDescription = 'Input vector to classify'\n",
    "knn_spec.description.output[0].shortDescription = 'Predicted label. Defaults to \\'unknown\\''\n",
    "knn_spec.description.output[1].shortDescription = 'Probabilities / score for each possible label.'\n",
    "\n",
    "# print knn_spec.description\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Updatable Pipeline\n",
    "\n",
    "Last step is to create the pipeline model and insert the feature extractor and the nearest neighbor classifier. The model will be set to be updatable."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "pipeline_spec = coremltools.proto.Model_pb2.Model()\n",
    "pipeline_spec.specificationVersion = coremltools._MINIMUM_UPDATABLE_SPEC_VERSION\n",
    "pipeline_spec.isUpdatable = True\n",
    "\n",
    "# Inputs are the inputs from the embedding model\n",
    "pipeline_spec.description.input.extend(embedding_spec.description.input[:])\n",
    "\n",
    "# Outputs are the outputs from the classification model\n",
    "pipeline_spec.description.output.extend(knn_spec.description.output[:])\n",
    "pipeline_spec.description.predictedFeatureName = knn_spec.description.predictedFeatureName\n",
    "pipeline_spec.description.predictedProbabilitiesName = knn_spec.description.predictedProbabilitiesName\n",
    "\n",
    "# Training inputs\n",
    "pipeline_spec.description.trainingInput.extend([embedding_spec.description.input[0]])\n",
    "pipeline_spec.description.trainingInput[0].shortDescription = 'Example sketch'\n",
    "pipeline_spec.description.trainingInput.extend([knn_spec.description.output[0]])\n",
    "pipeline_spec.description.trainingInput[1].shortDescription = 'Associated true label of example sketch'\n",
    "\n",
    "# Provide metadata\n",
    "pipeline_spec.description.metadata.author = 'Core ML Tools'\n",
    "pipeline_spec.description.metadata.license = 'MIT'\n",
    "pipeline_spec.description.metadata.shortDescription = ('An updatable model which can be used to train a tiny 28 x 28 drawing classifier based on user examples.'\n",
    "                                                       ' It uses a drawing embedding trained on the Quick, Draw! dataset (https://github.com/googlecreativelab/quickdraw-dataset)')\n",
    "\n",
    "# Construct pipeline by adding the embedding and then the nearest neighbor classifier\n",
    "pipeline_spec.pipelineClassifier.pipeline.models.add().CopyFrom(embedding_spec)\n",
    "pipeline_spec.pipelineClassifier.pipeline.models.add().CopyFrom(knn_spec)\n",
    "\n",
    "# Save the updated spec.\n",
    "from coremltools.models import MLModel\n",
    "mlmodel = MLModel(pipeline_spec)\n",
    "\n",
    "output_path = './TinyDrawingClassifier.mlmodel'\n",
    "from coremltools.models.utils import save_spec\n",
    "mlmodel.save(output_path)"
   ]
  }
 ],
 "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.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}