Merge pull request #465 from bhushan23/master

Adding tf converter options into tf coremltools path

Author: bhushan23

coremltools/converters/nnssa/coreml/ssa_converter.py

@@ -1,5 +1,7 @@
 import numpy as np
 
+from six import string_types as _string_types
+
 from coremltools.models import datatypes
 from coremltools.proto import NeuralNetwork_pb2
 from coremltools.models.neural_network import NeuralNetworkBuilder
@@ -18,7 +20,6 @@
 
 DEBUG = False
 
-
 def _is_scalar(type_):
     if type_ is None:
         return False
@@ -31,10 +32,19 @@ def ssa_convert(ssa,
                 top_func='main',
                 inputs=None,
                 outputs=None,
+                image_input_names=None,
+                is_bgr=False,
+                red_bias=0.0,
+                green_bias=0.0,
+                blue_bias=0.0,
+                gray_bias=0.0,
+                image_scale=1.0,
+                class_labels=None,
+                predicted_feature_name=None,
+                predicted_probabilities_output='',
                 add_custom_layers=False,
                 custom_conversion_functions={},
-                custom_shape_functions={}
-                ):
+                custom_shape_functions={}):
     """
     Convert NNSSA into CoreML spec.
     ssa : NetworkEnsemble
@@ -109,33 +119,82 @@ def ssa_convert(ssa,
     for f in list(ssa.functions.values()):
         check_connections(f.graph)
 
+    # Set classifier flag
+    is_classifier = class_labels is not None
+    neural_network_type = 'classifier' if is_classifier else None
+
     converter = SSAConverter(ssa,
                              top_func=top_func,
                              inputs=inputs,
                              outputs=outputs,
+                             neural_network_type=neural_network_type,
                              add_custom_layers=add_custom_layers,
                              custom_conversion_functions=custom_conversion_functions,
                              custom_shape_functions=custom_shape_functions)
     converter.convert()
+
+    builder = converter._get_builder(func=top_func)
+    # Add image input identifier
+    if image_input_names is not None and isinstance(
+        image_input_names, _string_types):
+        image_input_names = [image_input_names]
+
+    # Add classifier classes (if applicable)
+    if is_classifier:
+        classes_in = class_labels
+        if isinstance(classes_in, _string_types):
+            import os
+            if not os.path.isfile(classes_in):
+                raise ValueError("Path to class labels (%s) does not exist." % \
+                    classes_in)
+                with open(classes_in, 'r') as f:
+                    classes = f.read()
+                classes = classes.splitlines()
+            elif type(classes_in) is list: # list[int or str]
+                classes = classes_in
+            else:
+                raise ValueError('Class labels must be a list of integers / strings,'\
+                    ' or a file path')
+
+            if predicted_feature_name is not None:
+                builder.set_class_labels(
+                    classes, predicted_feature_name=predicted_feature_name,
+                    prediction_blob=predicted_probabilities_output)
+            else:
+                builder.set_class_labels(classes)
+
+    image_format = ssa.get_image_format()
+    # Set pre-processing parameters
+    builder.set_pre_processing_parameters(image_input_names=image_input_names,
+                                          is_bgr=is_bgr,
+                                          red_bias=red_bias,
+                                          green_bias=green_bias,
+                                          blue_bias=blue_bias,
+                                          gray_bias=gray_bias,
+                                          image_scale=image_scale,
+                                          image_format=image_format)
+
     mlmodel_spec = converter.get_spec()
 
+    # MLModel passes
     mlmodel_passes = [remove_disconnected_constants]
     for p in mlmodel_passes:
         p(mlmodel_spec)
 
+    
     if DEBUG:
         coremltools.models.utils.save_spec(mlmodel_spec, '/tmp/model_from_spec.mlmodel')
 
     return mlmodel_spec
 
 
 class SSAConverter(object):
-
     def __init__(self,
                  net_ensemble, # type: NetworkEnsemble
                  top_func='main', # type: str
                  inputs=None, # type: List[str]
                  outputs=None, # type: List[str]
+                 neural_network_type=None, # type: str
                  add_custom_layers=False,  # type: bool
                  custom_conversion_functions={},  # type: Dict[Text, Any]
                  custom_shape_functions={} # type: Dict[Text, Any]
@@ -213,10 +272,10 @@ def __init__(self,
         else:
             top_output_features = list(zip(top_output_names, [None] * len(top_output_names)))
 
-        self.top_builder = NeuralNetworkBuilder(
-            input_features=top_input_features,
-            output_features=top_output_features,
-            disable_rank5_shape_mapping=True)
+        self.top_builder = NeuralNetworkBuilder(input_features=top_input_features,
+                                                output_features=top_output_features,
+                                                disable_rank5_shape_mapping=True,
+                                                mode=neural_network_type)
 
         self.spec = self.top_builder.spec
 
@@ -574,6 +633,7 @@ def _convert_transpose(self, node):
             raise ValueError('[SSAConverter] Cannot handle dynamic Transpose')
         dim = list(dim)
         builder = self._get_builder()
+
         layer = builder.add_transpose(
             name=node.name, axes=dim, input_name=input_names[0], output_name=node.name)
 

coremltools/converters/nnssa/nnssa.py

@@ -430,3 +430,19 @@ def _find_free_name(self, prefix):
                 idx += 1
             else:
                 return name
+
+    def get_image_format(self):
+        """
+        Iterates over graph and returns input format (`NCHW` or `NHWC`)
+        if input is of type Image, otherwise `None`
+        """
+        for fn_key in list(self.functions.keys()):
+            graph = self.functions[fn_key].graph
+
+            for name in graph:
+                node = graph[name]
+                if node.attr.get('data_format', None) == 'NHWC' or node.attr.get('data_format') == 'NHWC_format_inserted':
+                    return 'NHWC'
+                elif node.attr.get('data_format', None) == 'NCHW':
+                    return 'NCHW'
+        return None

coremltools/converters/tensorflow/_tf_converter.py

@@ -7,10 +7,19 @@
 import os.path
 from ...models import MLModel
 
-
 def convert(filename,
             inputs=None,
             outputs=None,
+            image_input_names=None,
+            is_bgr=False,
+            red_bias=0.0,
+            green_bias=0.0,
+            blue_bias=0.0,
+            gray_bias=0.0,
+            image_scale=1.0,
+            class_labels=None,
+            predicted_feature_name=None,
+            predicted_probabilities_output='',
             add_custom_layers=False,  # type: bool
             custom_conversion_functions={},  # type: Dict[Text, Any]
             custom_shape_functions={}, # type: Dict[Text, Any]
@@ -20,7 +29,6 @@ def convert(filename,
 
     if not filename.endswith('.pb'):
         raise ValueError('invalid input tf_model_path format, expecting TensorFlow frozen graph (.pb) model.')
-
     # convert from TensorFlow to SSA
     try:
         from ..nnssa.frontend.tensorflow import load as frontend_load
@@ -35,10 +43,19 @@ def convert(filename,
                                   top_func='main',
                                   inputs=inputs,
                                   outputs=outputs,
+                                  image_input_names=image_input_names,
+                                  is_bgr=is_bgr,
+                                  red_bias=red_bias,
+                                  green_bias=green_bias,
+                                  blue_bias=blue_bias,
+                                  gray_bias=gray_bias,
+                                  image_scale=image_scale,
+                                  class_labels=class_labels,
+                                  predicted_feature_name=predicted_feature_name,
+                                  predicted_probabilities_output=predicted_probabilities_output,
                                   add_custom_layers=add_custom_layers,
                                   custom_conversion_functions=custom_conversion_functions,
-                                  custom_shape_functions=custom_shape_functions
-                                  )
+                                  custom_shape_functions=custom_shape_functions)
     except ImportError as err:
         raise ImportError("Backend converter not found! Error message:\n%s" % err)
 

coremltools/models/neural_network/builder.py

@@ -3204,7 +3204,8 @@ def add_crop_resize(self, name, input_names, output_name, target_height=1, targe
         return spec_layer
 
     def set_pre_processing_parameters(self, image_input_names=None, is_bgr=False,
-                                      red_bias=0.0, green_bias=0.0, blue_bias=0.0, gray_bias=0.0, image_scale=1.0):
+                                      red_bias=0.0, green_bias=0.0, blue_bias=0.0, gray_bias=0.0, image_scale=1.0,
+                                      image_format='NCHW'):
         """
         Add a pre-processing parameters layer to the neural network object.
 
@@ -3232,6 +3233,9 @@ def set_pre_processing_parameters(self, image_input_names=None, is_bgr=False,
 
         image_scale: float or dict()
             Value by which to scale the images.
+        
+        image_format: str
+            Image format, either 'NCHW' / 'NHWC'
 
         See Also
         --------
@@ -3241,6 +3245,9 @@ def set_pre_processing_parameters(self, image_input_names=None, is_bgr=False,
         if not image_input_names:
             return  # nothing to do here
 
+        if image_format != 'NCHW' and image_format != 'NHWC':
+            raise ValueError("Input image format must be either 'NCHW' or 'NHWC'. Provided {}".format(image_format))
+
         if not isinstance(is_bgr, dict):
             is_bgr = dict.fromkeys(image_input_names, is_bgr)
         if not isinstance(red_bias, dict):
@@ -3259,7 +3266,43 @@ def set_pre_processing_parameters(self, image_input_names=None, is_bgr=False,
             if input_.name in image_input_names:
                 if input_.type.WhichOneof('Type') == 'multiArrayType':
                     array_shape = tuple(input_.type.multiArrayType.shape)
-                    channels, height, width = array_shape
+
+                    if len(array_shape) == 4:
+                        input_indices = [0, 1, 2, 3] if image_format == 'NCHW' else [0, 3, 1, 2]
+                    elif len(array_shape) == 3:
+                        # Adding dummy index for 'batch' for compatibility
+                        input_indices = [0, 0, 1, 2] if image_format == 'NCHW' else [0, 2, 0, 1]
+                    else:
+                        raise ValueError("Invalid input shape. Input of rank {}, but expecting input of either rank 3 or rank 4".format(len(array_shape)))
+
+                    # Extract image shape depending on input format
+                    _, channels, height, width = [array_shape[e] for e in input_indices]
+
+                    if image_format == 'NHWC':
+                        # If input format is 'NHWC', then add transpose
+                        # after the input and replace all use of input
+                        # with output of transpose
+                        axes = [1, 2, 0]
+                        if len(array_shape) == 4:
+                            axes = [0, 2, 3, 1]
+                        input_transpose = input_.name + '_to_nhwc'
+                        transpose_layer = self.add_transpose(
+                                                name=input_transpose,
+                                                axes=axes,
+                                                input_name=input_.name,
+                                                output_name=input_transpose
+                                                )
+                        layers = spec.neuralNetwork.layers
+                        layers.insert(0, layers.pop())
+                        for layer_ in layers:
+                            for i in range(len(layer_.input)):
+                                if layer_.name == input_transpose:
+                                    continue
+                                if layer_.input[i] == input_.name:
+                                    layer_.input[i] = input_transpose
+                    
+                    # TODO: If input is not rank 3 or 4, then accordingly handle
+                    # e.g. for rank-2 input, squeeze additional dimension in case of Gray scale image
                     if channels == 1:
                         input_.type.imageType.colorSpace = _FeatureTypes_pb2.ImageFeatureType.ColorSpace.Value(
                             'GRAYSCALE')