Image similarity export to CoreML (#242)

* Image similarity export to CoreML provided reference data

* Added the ability to export data from ml_data to the model.

* Unittests and better documentation

* Export metadata

* Simplify export

src/unity/python/turicreate/test/test_image_similarity.py

@@ -6,18 +6,18 @@
 from __future__ import print_function as _
 from __future__ import division as _
 from __future__ import absolute_import as _
-import sys
 import unittest
 import pytest
 import turicreate as tc
 from turicreate.toolkits import _image_feature_extractor
+from turicreate.toolkits._internal_utils import _mac_ver
 import tempfile
 from . import util as test_util
 import coremltools
-import numpy
-import platform
+import numpy as np
 from turicreate.toolkits._main import ToolkitError as _ToolkitError
 
+
 def _get_data():
     from PIL import Image as _PIL_Image
     import random
@@ -166,6 +166,33 @@ def test_repr(self):
         self.assertEqual(type(str(model)), str)
         self.assertEqual(type(model.__repr__()), str)
 
+    def test_export_coreml(self):
+        """
+        Check the export_coreml() function.
+        """
+
+        # Save the model as a CoreML model file
+        filename = tempfile.mkstemp('ImageSimilarity.mlmodel')[1]
+        self.model.export_coreml(filename)
+
+        # Load the model back from the CoreML model file
+        coreml_model = coremltools.models.MLModel(filename)
+
+        # Get model distances for comparison
+        tc_ret = self.model.query(self.sf[:1], k=self.sf.num_rows())
+
+        if _mac_ver() >= (10, 13):
+            from PIL import Image as _PIL_Image
+
+            ref_img = self.sf[0]['awesome_image'].pixel_data
+            pil_img = _PIL_Image.fromarray(ref_img)
+            coreml_ret = coreml_model.predict({'awesome_image': pil_img}, useCPUOnly=True)
+
+            # Compare distances
+            coreml_distances = np.array(sorted(coreml_ret['distance']))
+            tc_distances = tc_ret['distance'].to_numpy()
+            np.testing.assert_array_almost_equal(tc_distances, coreml_distances, decimal=2)
+
     def test_save_and_load(self):
         with test_util.TempDirectory() as filename:
 
@@ -182,6 +209,8 @@ def test_save_and_load(self):
             print("Summary passed")
             self.test__list_fields()
             print("List fields passed")
+            self.test_export_coreml()
+            print("Export coreml passed")
 
 
 @unittest.skipIf(tc.util._num_available_gpus() == 0, 'Requires GPU')

src/unity/python/turicreate/toolkits/image_similarity/image_similarity.py

@@ -409,3 +409,134 @@ def similarity_graph(self, k=5, radius=None, include_self_edges=False,
         +----------+----------+----------------+------+
         """
         return self.similarity_model.similarity_graph(k, radius, include_self_edges, output_type, verbose)
+
+    def export_coreml(self, filename):
+        """
+        Save the model in Core ML format.
+        The exported model calculates the distance between a query image and
+        each row of the model's stored data. It does not sort and retrieve
+        the k nearest neighbors of the query image.
+
+        See Also
+        --------
+        save
+
+        Examples
+        --------
+        # Train an image similarity model
+        >>> model = turicreate.image_similarity.create(data)
+
+        # Query the model for similar images
+        >>> similar_images = model.query(data)
+        +-------------+-----------------+---------------+------+
+        | query_label | reference_label |    distance   | rank |
+        +-------------+-----------------+---------------+------+
+        |      0      |        0        |      0.0      |  1   |
+        |      0      |        2        | 24.9664942809 |  2   |
+        |      0      |        1        | 28.4416069428 |  3   |
+        |      1      |        1        |      0.0      |  1   |
+        |      1      |        2        | 21.8715131191 |  2   |
+        |      1      |        0        | 28.4416069428 |  3   |
+        |      2      |        2        |      0.0      |  1   |
+        |      2      |        1        | 21.8715131191 |  2   |
+        |      2      |        0        | 24.9664942809 |  3   |
+        +-------------+-----------------+---------------+------+
+        [9 rows x 4 columns]
+
+        # Export the model to Core ML format
+        >>> model.export_coreml('myModel.mlmodel')
+
+        # Load the Core ML model
+        >>> import coremltools
+        >>> ml_model = coremltools.models.MLModel('myModel.mlmodel')
+
+        # Prepare the first image of reference data for consumption
+        # by the Core ML model
+        >>> import PIL
+        >>> image = tc.image_analysis.resize(
+                data['image'][0], *reversed(model.input_image_shape))
+        >>> image = PIL.Image.fromarray(image.pixel_data)
+
+        # Calculate distances using the Core ML model
+        >>> ml_model.predict(data={'image': image})
+        {'distance': array([ 0.      , 28.453125, 24.96875 ])}
+        """
+        import numpy as _np
+        import coremltools as _cmt
+        from coremltools.models import datatypes as _datatypes, neural_network as _neural_network
+        from .._mxnet_to_coreml import _mxnet_converter
+        from turicreate.toolkits import _coreml_utils
+
+        # Get the reference data from the model
+        proxy = self.similarity_model.__proxy__
+        reference_data = _np.array(_tc.extensions._nearest_neighbors._nn_get_reference_data(proxy))
+        num_examples, embedding_size = reference_data.shape
+
+        # Get the input and output names
+        input_name = self.feature_extractor.data_layer
+        output_name = 'distance'
+        input_features = [(input_name, _datatypes.Array(*(self.input_image_shape)))]
+        output_features = [(output_name, _datatypes.Array(num_examples))]
+
+        # Create a neural network
+        builder = _neural_network.NeuralNetworkBuilder(
+            input_features, output_features, mode=None)
+
+        # Convert the feature extraction network
+        mx_feature_extractor = self.feature_extractor._get_mx_module(
+            self.feature_extractor.ptModel.mxmodel,
+            self.feature_extractor.data_layer,
+            self.feature_extractor.feature_layer,
+            self.feature_extractor.context,
+            self.input_image_shape
+        )
+        batch_input_shape = (1, ) + self.input_image_shape
+        _mxnet_converter.convert(mx_feature_extractor, mode=None,
+                                 input_shape={input_name: batch_input_shape},
+                                 builder=builder, verbose=False)
+
+        # To add the nearest neighbors model we add calculation of the euclidean 
+        # distance between the newly extracted query features (denoted by the vector u)
+        # and each extracted reference feature (denoted by the rows of matrix V).
+        # Calculation of sqrt((v_i-u)^2) = sqrt(v_i^2 - 2v_i*u + u^2) ensues.
+        V = reference_data
+        v_squared = (V * V).sum(axis=1)
+
+        feature_layer = self.feature_extractor.feature_layer
+        builder.add_inner_product('v^2-2vu', W=-2 * V, b=v_squared, has_bias=True,
+                                  input_channels=embedding_size, output_channels=num_examples,
+                                  input_name=feature_layer, output_name='v^2-2vu')
+
+        builder.add_unary('element_wise-u^2', mode='power', alpha=2,
+                          input_name=feature_layer, output_name='element_wise-u^2')
+
+        # Produce a vector of length num_examples with all values equal to u^2
+        builder.add_inner_product('u^2', W=_np.ones((embedding_size, num_examples)),
+                                  b=None, has_bias=False,
+                                  input_channels=embedding_size, output_channels=num_examples,
+                                  input_name='element_wise-u^2', output_name='u^2')
+
+        builder.add_elementwise('v^2-2vu+u^2', mode='ADD',
+                                input_names=['v^2-2vu', 'u^2'],
+                                output_name='v^2-2vu+u^2')
+
+        # v^2-2vu+u^2=(v-u)^2 is non-negative but some computations on GPU may result in
+        # small negative values. Apply RELU so we don't take the square root of negative values.
+        builder.add_activation('relu', non_linearity='RELU',
+                               input_name='v^2-2vu+u^2', output_name='relu')
+        builder.add_unary('sqrt', mode='sqrt', input_name='relu', output_name=output_name)
+
+        # Finalize model
+        _mxnet_converter._set_input_output_layers(builder, [input_name], [output_name])
+        builder.set_input([input_name], [self.input_image_shape])
+        builder.set_output([output_name], [(num_examples,)])
+        _cmt.models.utils.rename_feature(builder.spec, input_name, self.feature)
+        builder.set_pre_processing_parameters(image_input_names=self.feature)
+
+        # Add metadata
+        mlmodel = _cmt.models.MLModel(builder.spec)
+        model_type = 'image similarity'
+        mlmodel.short_description = _coreml_utils._mlmodel_short_description(model_type)
+        mlmodel.input_description[self.feature] = u'Input image'
+        mlmodel.output_description[output_name] = u'Distances between the input and reference images'
+        mlmodel.save(filename)

src/unity/toolkits/nearest_neighbors/nearest_neighbors.cpp

@@ -931,6 +931,22 @@ double neighbor_candidates::get_max_dist() const {
 }
 
 
+flexible_type nearest_neighbors_model::get_reference_data() const {
+    DASSERT_EQ(num_examples, mld_ref.size());
+
+    DenseVector ref_data(metadata->num_dimensions());
+    flex_list ret(num_examples);
+    for (auto it = mld_ref.get_iterator(0, 1); !it.done(); ++it) {
+      it.fill_row_expr(ref_data);
+      ret[it.row_index()] = arma::conv_to<std::vector<double>>::from(ref_data);
+    }
+
+    return ret;
+}
+
+flexible_type _nn_get_reference_data(std::shared_ptr<nearest_neighbors_model> model) {
+  return model->get_reference_data();
+}
 
 }  // namespace nearest_neighbors
 }  // namespace turi

src/unity/toolkits/nearest_neighbors/nearest_neighbors.hpp

@@ -36,8 +36,8 @@ namespace nearest_neighbors {
 
 typedef std::tuple<std::vector<std::string>, function_closure_info, double> dist_component_type;
 
-} 
-} 
+}
+}
 
 BEGIN_OUT_OF_PLACE_SAVE(arc, turi::nearest_neighbors::dist_component_type, d) {
   std::map<std::string, turi::variant_type> data;
@@ -57,7 +57,7 @@ BEGIN_OUT_OF_PLACE_LOAD(arc, turi::nearest_neighbors::dist_component_type, d) {
   __EXTRACT(column_names);
   __EXTRACT(weight);
 #undef __EXTRACT
-  arc >> distance_info;  
+  arc >> distance_info;
   d = std::make_tuple(column_names, distance_info, weight);
 
 } END_OUT_OF_PLACE_LOAD()
@@ -554,9 +554,14 @@ class EXPORT nearest_neighbors_model : public ml_model_base {
                                    const sframe& X,
                                    const function_closure_info distance_name,
                                    const double weight);
-};
 
+  /**
+   * Get reference data as a vector of vectors
+   * \returns Reference data as a vector of vectors (in ml-data form)
+   */
+  flexible_type get_reference_data() const;
 
+};
 
 // -----------------------------------------------------------------------------
 // CANDIDATE NEIGHBORS CLASS
@@ -670,6 +675,12 @@ class neighbor_candidates {
   double get_max_dist() const;
 };
 
+/**
+ * Function to get the reference data from the NN model
+ *
+ * \param[in] model Nearest neighbour model.
+ */
+flexible_type _nn_get_reference_data(std::shared_ptr<nearest_neighbors_model> model);
 
 }  // namespace nearest_neighbors
 }  // namespace turi

src/unity/toolkits/nearest_neighbors/unity_nearest_neighbors.cpp

@@ -399,6 +399,7 @@ EXPORT std::vector<toolkit_function_specification> get_toolkit_function_registra
           similarity_graph_spec};
 
   REGISTER_FUNCTION(train, "params")
+  REGISTER_FUNCTION(_nn_get_reference_data, "model")
   return specs;
 }