{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Example: Convert a TensorFlow 2 model to CoreML\n",
    "\n",
    "In this example, we will convert a pretrained TensorFlow 2 model created with the `tf.keras` API.  For more details on all the different ways to convert TensorFlow 2 models, see the [coremltools TensorFlow 2 documentation](https://coremltools.readme.io/docs/tensorflow-2).  For TensorFlow 1.x, see the [coremltools TensorFlow 1 documentation](https://coremltools.readme.io/docs/tensorflow-1)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create a TensorFlow Model\n",
    "\n",
    "For this example, we will load [MobileNetV2](https://keras.io/api/applications/mobilenet/#mobilenetv2-function) with pretrained weights from ImageNet.  Keras assumes the input image has a been rescaled and offset to a range of [-1, 1], which coremltools can do during the conversion process.  If more complex input preparation is required, those layers should be included in the model before conversion, or the target application should do the required data preparation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "\n",
    "input_shape = (192, 192, 3)\n",
    "\n",
    "tf_model = tf.keras.applications.MobileNetV2(\n",
    "    input_shape=input_shape,\n",
    "    include_top=True,\n",
    "    weights=\"imagenet\",\n",
    "    classes=1000,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We also want the class labels to convert the model predictions to English category names."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Download class labels (from a separate file)\n",
    "import urllib\n",
    "label_url = 'https://storage.googleapis.com/download.tensorflow.org/data/ImageNetLabels.txt'\n",
    "class_labels = urllib.request.urlopen(label_url).read().decode(\"utf-8\").splitlines()\n",
    "class_labels = class_labels[1:] # remove the first class which is background\n",
    "assert len(class_labels) == 1000"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Test the TensorFlow Model\n",
    "\n",
    "As a quick check, let's load an image of a dog ([source, CC BY-SA 4.0](https://commons.wikimedia.org/wiki/File:Dog_Breeds.jpg)) and run it through the Keras model to ensure it is working."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import PIL\n",
    "\n",
    "img = PIL.Image.open('Dog.jpg').resize(input_shape[:2])\n",
    "img"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "# prepare with appropriate dimensionality (batch size of 1), scale and offset\n",
    "img_np = (np.array(img)[np.newaxis,:,:,:]/127.0) - 1\n",
    "\n",
    "out = tf_model.predict(img_np)[0]\n",
    "top_5_indices = np.flip(np.argsort(out))[:5]\n",
    "for i in top_5_indices:\n",
    "    print(class_labels[i], out[i])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Convert to CoreML Model\n",
    "\n",
    "The convert() function in coremltools will take a TensorFlow or PyTorch model in one of many forms (in this case, we have a `tf.keras.Model` object), analyze it, and generate a CoreML model wrapped in a `ct.models.MLModel` Python object.\n",
    "\n",
    "The API has two optional arguments which we will use:\n",
    "\n",
    "- `inputs`: List of type descriptions for the input layers.  The types can be inferred from the input model in some cases, but here we want to interpret the (192, 192, 3) input tensor as an image.  This will allow us to specify a scale and offset (\"bias\") for the pixel values, as required by the model, and also describe the color channel ordering.  This color layout metadata is included in the CoreML model itself and can be used by the consuming application of the model to ensure correct pixel ordering.\n",
    "\n",
    "- `classifier_config`: Since this is a classifier model, we include the class labels here so that the user of the CoreML model has that data as well.\n",
    "\n",
    "The converter will translate the TensorFlow model into a [Model Intermediate Language](https://coremltools.readme.io/docs/model-intermediate-language) representation, optimize it for performance, and finally emit the equivalent CoreML operations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import coremltools as ct\n",
    "\n",
    "cml_model = ct.convert(tf_model,\n",
    "                     inputs=[ct.ImageType(color_layout=\"RGB\", scale=1/127.0, bias=[-1, -1, -1], )],\n",
    "                     classifier_config=ct.ClassifierConfig(class_labels)\n",
    "                    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can see the metadata describing this model by looking at the string representation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cml_model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we can write the model to disk in the CoreML format."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cml_model.save('MobileNet_v2_tf.mlmodel')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Using the CoreML Model\n",
    "\n",
    "There are two ways to use the trained CoreML model for inference:\n",
    "\n",
    "- Load the model in Python using coremltools and call the predict() method.  This requires macOS as the model is executed by the CoreML framework.  The OS will use whatever specialized hardware may be available to accelerate execution.\n",
    "\n",
    "- Load the model in a compiled macOS or iOS app using the [CoreML APIs](https://developer.apple.com/documentation/coreml/integrating_a_core_ml_model_into_your_app).\n",
    "\n",
    "If you are running this notebook on a Mac (not available on Binder), you can test the model directly"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "IS_MACOS = sys.platform == 'darwin'\n",
    "\n",
    "if IS_MACOS:\n",
    "    loaded_model = ct.models.MLModel('MobileNet_v2_tf.mlmodel')\n",
    "    prediction = loaded_model.predict({'input_1': img})\n",
    "else:\n",
    "    prediction = 'Skipping prediction on non-macOS system'\n",
    "prediction"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Testing the CoreML Model on iOS\n",
    "\n",
    "If you are running this notebook remotely and browsing on an iOS device, you can download the model you just trained and test it out with the CoreMLCompare app ([source code](https://github.com/sivu22/CoreMLCompare)) and your device camera.  (Note: the author of CoreMLCompare is not affiliated with coremltools or Apple.  It is just a handy app for loading and comparing CoreML image classifiers.)\n",
    "\n",
    "Steps:\n",
    "- [Install CoreMLCompare from the App Store](https://itunes.apple.com/us/app/coremlcompare/id1378662341) or build it from source and load on your device if you have an Apple developer account.\n",
    "- Download the CoreML model you just trained onto your device from the Jupyter file browser or by tapping this link: [MobileNet_v2_tf.mlmodel](MobileNet_v2_tf.mlmodel)\n",
    "- Open the \"Files\" app in iOS, browse to the file you downloaded, and tap on it.  Then tap the share icon and select \"CoreMLCompare\" on the share sheet that appears.  This will open the CoreMLCompare app.\n",
    "- The app will pop up a modal asking for the index of the new model.  CoreMLCompare can run up to 4 different image classifiers simultaneously, and includes two by default (MobileNet and SqueezeNet).  Select index 3 and tap \"Import\" to put your new model into an empty slot.\n",
    "- After a brief delay, the image classifications at the bottom of the image will start to update live based on what your camera is pointing at.\n",
    "\n",
    "Note that this model is using ImageNet categories which might not correspond to the types of objects you have nearby.  A pen on a white piece of paper provides a good test case.  It will likely be classified as \"Ballpoint\" or \"Fountain Pen\"."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}