Manipulation of metadata props (#115)

* Add image metadata to CoreML conversion

* Remove extra properties in data_types

* Add set_denotation

* Remove properties with same case insensitive name

* Add denotation automatically on images

* Check ONNX version before adding metadata

* Use StrictVersion for comparison

* Fix code style

* Move metadata and denotation to Topology/TensorType

* Change Topology default metadata_props value

* Match color_space with ONNX definitions

* Refactor case insensitive dictionaries

* Use warnings module

* Put valid metadata props in separate constant

* Avoid casefold() for compatibility with Python 2

* Don't guess the color space and nominal pixel range from CoreML

* Add tests to image metadata

* Don't set denotation in ONNX < 1.2.2

* Use unittest.skipIf instead of running empty test

* Add docstrings

* Add usage examples

* Forward tensor denotation when removing redundant variables

* Only check denotation in tensors

Author: tiagoshibata

onnxmltools/convert/common/_topology.py

@@ -9,6 +9,7 @@
 from distutils.version import StrictVersion
 from ...proto import onnx
 from ...proto import helper
+from ...utils.metadata_props import add_metadata_props
 from .data_types import *
 from ._container import ModelComponentContainer
 from . import _registration
@@ -234,7 +235,7 @@ class Topology:
 
     def __init__(self, model, default_batch_size=1, initial_types=None,
                  reserved_variable_names=None, reserved_operator_names=None, targeted_onnx=None,
-                 custom_conversion_functions=None, custom_shape_calculators=None):
+                 custom_conversion_functions=None, custom_shape_calculators=None, metadata_props=None):
         '''
         Initialize a Topology object, which is an intermediate representation of a computational graph.
 
@@ -253,6 +254,7 @@ def __init__(self, model, default_batch_size=1, initial_types=None,
         self.variable_name_set = reserved_variable_names if reserved_variable_names is not None else set()
         self.operator_name_set = reserved_operator_names if reserved_operator_names is not None else set()
         self.initial_types = initial_types if initial_types else list()
+        self.metadata_props = metadata_props if metadata_props else dict()
         self.default_batch_size = default_batch_size
         self.targeted_onnx_version = StrictVersion(targeted_onnx)
         self.custom_conversion_functions = custom_conversion_functions if custom_conversion_functions else {}
@@ -520,19 +522,29 @@ def _resolve_duplicates(self):
                             continue
                         another_operator.inputs[i] = original
 
-            # When original variable's document string is empty but duplicate's document string is not, we
-            # copy that non-empty string to the original variable to avoid information loss.
+            # When original variable's documentation string or denotation is empty but duplicate's is not, we
+            # copy that field to the original variable to avoid information loss.
             if not original.type.doc_string and duplicate.type.doc_string:
                 original.type.doc_string = duplicate.type.doc_string
 
-            # Sometime, shapes of duplicates are different. We try to replace the original variable's unknown dimensions
-            # as many as possible because we will get rid of the duplicate.
-            if isinstance(original.type, TensorType) and isinstance(duplicate.type, TensorType) and \
-                    len(original.type.shape) == len(duplicate.type.shape):
-                for i in range(len(original.type.shape)):
-                    if original.type.shape[i] != 'None':
-                        continue
-                    original.type.shape[i] = duplicate.type.shape[i]
+            if isinstance(original.type, TensorType) and isinstance(duplicate.type, TensorType):
+                if not original.type.denotation and duplicate.type.denotation:
+                    original.type.denotation = duplicate.type.denotation
+                if not original.type.channel_denotations:
+                    original.type.channel_denotations = duplicate.type.channel_denotations
+                elif duplicate.type.channel_denotations:
+                    # Merge the channel denotations if available in both the original and the duplicate
+                    for i in range(len(original.type.channel_denotations)):
+                        if original.type.channel_denotations[i]:
+                            continue
+                        original.type.channel_denotations[i] = duplicate.type.channel_denotations[i]
+                # Sometime, shapes of duplicates are different. We try to replace the original variable's unknown dimensions
+                # as many as possible because we will get rid of the duplicate.
+                if len(original.type.shape) == len(duplicate.type.shape):
+                    for i in range(len(original.type.shape)):
+                        if original.type.shape[i] != 'None':
+                            continue
+                        original.type.shape[i] = duplicate.type.shape[i]
 
             # Because we're iterating through the topology, we cannot delete any operator or variable. Otherwise,
             # the traversing function may be broken. We will delete those abandoned ones later.
@@ -735,6 +747,7 @@ def convert_topology(topology, model_name, doc_string, targeted_onnx):
         i += 1
 
     # Add extra information
+    add_metadata_props(onnx_model, topology.metadata_props)
     onnx_model.ir_version = onnx_proto.IR_VERSION
     onnx_model.producer_name = utils.get_producer()
     onnx_model.producer_version = utils.get_producer_version()

onnxmltools/convert/common/case_insensitive_dict.py

@@ -0,0 +1,48 @@
+try:
+    from collections.abc import Mapping, MutableMapping
+except ImportError:
+    from collections import Mapping, MutableMapping
+from collections import OrderedDict
+
+
+class CaseInsensitiveDict(MutableMapping):
+    def __init__(self, data=None, **kwargs):
+        self._dict = OrderedDict()
+        if data:
+            self.update(data, **kwargs)
+
+    def __setitem__(self, key, value):
+        self._dict[key.lower()] = (key, value)
+
+    def __getitem__(self, key):
+        return self._dict[key.lower()][1]
+
+    def __delitem__(self, key):
+        del self._dict[key.lower()]
+
+    def __iter__(self):
+        return (key for key, _ in self._dict.values())
+
+    def __len__(self):
+        return len(self._dict)
+
+    def lower_key_iteritems(self):
+        """Like iteritems(), but with lowercase keys."""
+        return (
+            (lower_key, keyval[1])
+            for lower_key, keyval
+            in self._dict.items()
+        )
+
+    def __eq__(self, other):
+        if isinstance(other, Mapping):
+            other = CaseInsensitiveDict(other)
+        else:
+            return NotImplemented
+        return dict(self.lower_key_iteritems()) == dict(other.lower_key_iteritems())
+
+    def copy(self):
+         return CaseInsensitiveDict(self._dict.values())
+
+    def __repr__(self):
+        return str(dict(self.items()))

onnxmltools/convert/common/data_types.py

@@ -54,8 +54,10 @@ def to_onnx_type(self):
 
 
 class TensorType(DataType):
-    def __init__(self, shape=None, doc_string=''):
+    def __init__(self, shape=None, doc_string='', denotation=None, channel_denotations=None):
         super(TensorType, self).__init__([] if not shape else shape, doc_string)
+        self.denotation = denotation
+        self.channel_denotations = channel_denotations
 
     def _get_element_onnx_type(self):
         raise NotImplementedError()
@@ -71,6 +73,13 @@ def to_onnx_type(self):
                 s.dim_param = 'None'
             else:
                 raise ValueError('Unsupported dimension type: %s' % type(d))
+        if getattr(onnx_type, 'denotation', None) is not None:
+            if self.denotation:
+                onnx_type.denotation = self.denotation
+            if self.channel_denotations:
+                for d, denotation in zip(onnx_type.tensor_type.shape.dim, self.channel_denotations):
+                    if denotation:
+                        d.denotation = denotation
         return onnx_type
 
 
@@ -83,8 +92,8 @@ def _get_element_onnx_type(self):
 
 
 class FloatTensorType(TensorType):
-    def __init__(self, shape=None, color_space=None, doc_string=''):
-        super(FloatTensorType, self).__init__(shape, doc_string)
+    def __init__(self, shape=None, color_space=None, doc_string='', denotation=None, channel_denotations=None):
+        super(FloatTensorType, self).__init__(shape, doc_string, denotation, channel_denotations)
         self.color_space = color_space
 
     def _get_element_onnx_type(self):

onnxmltools/convert/coreml/_parse.py

@@ -5,6 +5,7 @@
 # --------------------------------------------------------------------------
 
 from distutils.version import StrictVersion
+import warnings
 from ...proto import onnx
 from ..common._container import CoremlModelContainer
 from ..common._topology import Topology
@@ -59,8 +60,9 @@ def _parse_coreml_feature(feature_info, targeted_onnx_version, batch_size=1):
             raise ValueError('Unknown image format. Only gray-level, RGB, and BGR are supported')
         shape.append(raw_type.imageType.height)
         shape.append(raw_type.imageType.width)
-        color_space_map = {10: 'GRAY', 20: 'RGB', 30: 'BGR'}
-        return FloatTensorType(shape, color_space_map[color_space], doc_string=doc_string)
+        color_space_map = {10: 'Gray8', 20: 'Rgb8', 30: 'Bgr8'}
+        return FloatTensorType(shape, color_space_map[color_space], doc_string=doc_string,
+                               denotation='IMAGE', channel_denotations=['DATA_BATCH', 'DATA_CHANNEL', 'DATA_FEATURE', 'DATA_FEATURE'])
     elif type_name == 'multiArrayType':
         element_type_id = raw_type.multiArrayType.dataType
         shape = [d for d in raw_type.multiArrayType.shape]
@@ -463,8 +465,12 @@ def parse_coreml(model, initial_types=None, targeted_onnx=onnx.__version__, cust
     _parse_model(topology, scope, model)
     topology.compile()
 
-    # Use original CoreML names for model-level input(s)/output(s)
     for variable in topology.find_root_and_sink_variables():
+        color_space = getattr(variable.type, 'color_space', None)
+        if color_space:
+            if topology.metadata_props.setdefault('Image.BitmapPixelFormat', color_space) != color_space:
+                warnings.warn('Conflicting pixel formats found. In ONNX, all input/output images must use the same pixel format.')
+        # Use original CoreML names for model-level input(s)/output(s)
         if variable.raw_name not in reserved_variable_names:
             continue
         topology.rename_variable(variable.onnx_name, variable.raw_name)

onnxmltools/convert/coreml/operator_converters/neural_network/ImageScaler.py

@@ -14,11 +14,11 @@ def convert_preprocessing_scaler(scope, operator, container):
 
     attrs = {'name': operator.full_name, 'scale': params.channelScale}
     color_space = operator.inputs[0].type.color_space
-    if color_space == 'GRAY':
+    if color_space == 'Gray8':
         attrs['bias'] = [params.grayBias]
-    elif color_space == 'RGB':
+    elif color_space == 'Rgb8':
         attrs['bias'] = [params.redBias, params.greenBias, params.blueBias]
-    elif color_space == 'BGR':
+    elif color_space == 'Bgr8':
         attrs['bias'] = [params.blueBias, params.greenBias, params.redBias]
     else:
         raise ValueError('Unknown color space for tensor {}'.format(operator.inputs[0].full_name))

onnxmltools/utils/__init__.py

@@ -10,5 +10,6 @@
 from .main import set_model_version
 from .main import set_model_domain
 from .main import set_model_doc_string
+from .metadata_props import add_metadata_props, set_denotation
 from .visualize import visualize_model
 from .float16_converter import convert_float_to_float16

onnxmltools/utils/metadata_props.py

@@ -0,0 +1,84 @@
+import warnings
+from ..convert.common.case_insensitive_dict import CaseInsensitiveDict
+from ..proto import onnx, onnx_proto
+from distutils.version import StrictVersion
+
+
+KNOWN_METADATA_PROPS = CaseInsensitiveDict({
+    'Image.BitmapPixelFormat': ['gray8', 'rgb8', 'bgr8', 'rgba8', 'bgra8'],
+    'Image.ColorSpaceGamma': ['linear', 'srgb'],
+    'Image.NominalPixelRange': ['nominalrange_0_255', 'normalized_0_1', 'normalized_1_1', 'nominalrange_16_235'],
+})
+
+
+def _validate_metadata(metadata_props):
+    '''
+    Validate metadata properties and possibly show warnings or throw exceptions.
+
+    :param metadata_props: A dictionary of metadata properties, with property names and values (see :func:`~onnxmltools.utils.metadata_props.add_metadata_props` for examples)
+    '''
+    if len(CaseInsensitiveDict(metadata_props)) != len(metadata_props):
+        raise RuntimeError('Duplicate metadata props found')
+
+    for key, value in metadata_props.items():
+        valid_values = KNOWN_METADATA_PROPS.get(key)
+        if valid_values and value.lower() not in valid_values:
+            warnings.warn('Key {} has invalid value {}. Valid values are {}'.format(key, value, valid_values))
+
+
+def add_metadata_props(onnx_model, metadata_props, targeted_onnx=onnx.__version__):
+    '''
+    Add metadata properties to the model. See recommended key names at:
+    `Extensibility - Metadata <https://github.com/onnx/onnx/blob/296953db87b79c0137c5d9c1a8f26dfaa2495afc/docs/IR.md#metadata>`_ and
+    `Optional Metadata <https://github.com/onnx/onnx/blob/master/docs/IR.md#optional-metadata>`_
+
+
+    :param onnx_model: ONNX model object
+    :param metadata_props: A dictionary of metadata properties, with property names and values (example: `{ 'model_author': 'Alice', 'model_license': 'MIT' }`)
+    :param targeted_onnx: Target ONNX version
+    '''
+    if StrictVersion(targeted_onnx) < StrictVersion('1.2.1'):
+        warnings.warn('Metadata properties are not supported in targeted ONNX-%s' % targeted_onnx)
+        return
+    _validate_metadata(metadata_props)
+    new_metadata = CaseInsensitiveDict({x.key: x.value for x in onnx_model.metadata_props})
+    new_metadata.update(metadata_props)
+    del onnx_model.metadata_props[:]
+    onnx_model.metadata_props.extend(
+        onnx_proto.StringStringEntryProto(key=key, value=value)
+        for key, value in metadata_props.items()
+    )
+
+
+def set_denotation(onnx_model, input_name, denotation, dimension_denotation=None, targeted_onnx=onnx.__version__):
+    '''
+    Set input type denotation and dimension denotation.
+
+    Type denotation is a feature in ONNX 1.2.1 that let's the model specify the content of a tensor (e.g. IMAGE or AUDIO).
+    This information can be used by the backend. One example where it is useful is in images: Whenever data is bound to
+    a tensor with type denotation IMAGE, the backend can process the data (such as transforming the color space and
+    pixel format) based on model metadata properties.
+
+    :param onnx_model: ONNX model object
+    :param input_name: Name of input tensor to edit (example: `'data0'`)
+    :param denotation: Input type denotation (`documentation <https://github.com/onnx/onnx/blob/master/docs/TypeDenotation.md#type-denotation-definition>`_)
+    (example: `'IMAGE'`)
+    :param dimension_denotation: List of dimension type denotations. The length of the list must be the same of the number of dimensions in the tensor
+    (`documentation https://github.com/onnx/onnx/blob/master/docs/DimensionDenotation.md#denotation-definition>`_)
+    (example: `['DATA_BATCH', 'DATA_CHANNEL', 'DATA_FEATURE', 'DATA_FEATURE']`)
+    :param targeted_onnx: Target ONNX version
+    '''
+    if StrictVersion(targeted_onnx) < StrictVersion('1.2.1'):
+        warnings.warn('Denotation is not supported in targeted ONNX-%s' % targeted_onnx)
+        return
+    for graph_input in onnx_model.graph.input:
+        if graph_input.name == input_name:
+            graph_input.type.denotation = denotation
+            if dimension_denotation:
+                dimensions = graph_input.type.tensor_type.shape.dim
+                if len(dimension_denotation) != len(dimensions):
+                    raise RuntimeError('Wrong number of dimensions: input "{}" has {} dimensions'.format(input_name, len(dimensions)))
+                for dimension, channel_denotation in zip(dimensions, dimension_denotation):
+                    dimension.denotation = channel_denotation
+            return onnx_model
+    raise RuntimeError('Input "{}" not found'.format(input_name))

tests/coreml/test_AllNeuralNetworkConverters.py

@@ -2,7 +2,10 @@
 import numpy
 from coremltools.models.neural_network import NeuralNetworkBuilder
 from coremltools.models import datatypes
+from coremltools.proto.FeatureTypes_pb2 import ImageFeatureType
+from distutils.version import StrictVersion
 from onnxmltools import convert_coreml
+from onnxmltools.proto import onnx
 
 class TestNeuralNetworkLayerConverter(unittest.TestCase):
 
@@ -444,3 +447,26 @@ def test_bidirectional_lstm_converter(self):
         model_onnx = convert_coreml(builder.spec)
         self.assertTrue(model_onnx is not None)
 
+    def test_image_input_type_converter(self):
+        dim = (3, 15, 25)
+        inputs = [('input', datatypes.Array(*dim))]
+        outputs = [('output', datatypes.Array(*dim))]
+        builder = NeuralNetworkBuilder(inputs, outputs)
+        builder.add_elementwise(name='Identity', input_names=['input'],
+                                output_name='output', mode='ADD', alpha=0.0)
+        spec = builder.spec
+        input = spec.description.input[0]
+        input.type.imageType.height = dim[1]
+        input.type.imageType.width = dim[2]
+        for coreml_colorspace, onnx_colorspace in (('RGB', 'Rgb8'), ('BGR', 'Bgr8'), ('GRAYSCALE', 'Gray8')):
+            input.type.imageType.colorSpace = ImageFeatureType.ColorSpace.Value(coreml_colorspace)
+            model_onnx = convert_coreml(spec)
+            dims = [(d.dim_param or d.dim_value) for d in model_onnx.graph.input[0].type.tensor_type.shape.dim]
+            self.assertEqual(dims, ['None', 1 if onnx_colorspace == 'Gray8' else 3, 15, 25])
+
+            if StrictVersion(onnx.__version__) >= StrictVersion('1.2.1'):
+                metadata = {prop.key: prop.value for prop in model_onnx.metadata_props}
+                self.assertEqual(metadata, { 'Image.BitmapPixelFormat': onnx_colorspace })
+                self.assertEqual(model_onnx.graph.input[0].type.denotation, 'IMAGE')
+                channel_denotations = [d.denotation for d in model_onnx.graph.input[0].type.tensor_type.shape.dim]
+                self.assertEqual(channel_denotations, ['DATA_BATCH', 'DATA_CHANNEL', 'DATA_FEATURE', 'DATA_FEATURE'])

tests/utils/test_utils.py

@@ -1,12 +1,14 @@
 """
 Tests utilities.
 """
+from distutils.version import StrictVersion
 import filecmp
+from onnxmltools.proto import onnx
 import os
 import unittest
 
 from onnxmltools.utils import load_model, save_model, save_text
-from onnxmltools.utils import set_model_version, set_model_domain, set_model_doc_string
+from onnxmltools.utils import set_denotation, set_model_version, set_model_domain, set_model_doc_string
 
 
 class TestUtils(unittest.TestCase):
@@ -69,3 +71,13 @@ def test_set_docstring_blank(self):
         self.assertRaises(ValueError, set_model_doc_string, onnx_model.doc_string, "sample")
         set_model_doc_string(onnx_model, "", True)
         self.assertEqual(onnx_model.doc_string, "")
+
+    @unittest.skipIf(StrictVersion(onnx.__version__) < StrictVersion('1.2.1'),
+                     "not supported in this ONNX version")
+    def test_set_denotation(self):
+        this = os.path.dirname(__file__)
+        onnx_file = os.path.join(this, "models", "coreml_OneHotEncoder_BikeSharing.onnx")
+        onnx_model = load_model(onnx_file)
+        set_denotation(onnx_model, "1", "IMAGE", dimension_denotation=["DATA_FEATURE"])
+        self.assertEqual(onnx_model.graph.input[0].type.denotation, "IMAGE")
+        self.assertEqual(onnx_model.graph.input[0].type.tensor_type.shape.dim[0].denotation, "DATA_FEATURE")