# Copyright (c) 2020, Apple Inc. All rights reserved.
#
# Use of this source code is governed by a BSD-3-clause license that can be
# found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause
from enum import Enum
from typing import Optional
import numpy as np
from coremltools.converters.mil.mil import types
from coremltools.converters.mil.mil.types.symbolic import is_symbolic
class ColorLayout(Enum):
RGB = "RGB"
BGR = "BGR"
GRAYSCALE = "G"
GRAYSCALE_FLOAT16 = "G_FLOAT16"
[docs]
class ClassifierConfig:
[docs]
def __init__(
self,
class_labels,
predicted_feature_name="classLabel",
predicted_probabilities_output=None,
):
"""
Configuration for classifier models.
Parameters
----------
class_labels: str / list of int / list of str
* If a ``list`` is provided, the ``list`` maps the index of the output of a
neural network to labels in a classifier.
* If a ``str`` is provided, the ``str`` points to a file which maps the index
to labels in a classifier.
predicted_feature_name: str
Name of the output feature for the class labels exposed in the
Core ML neural network classifier. Default: ``'classLabel'``.
predicted_probabilities_output: str
* If provided, then this is the name of the neural network blob which
generates the probabilities for each class label (typically the output
of a softmax layer).
* If not provided, then the last output layer is assumed.
"""
self.class_labels = class_labels
self.predicted_feature_name = predicted_feature_name
self.predicted_probabilities_output = predicted_probabilities_output
[docs]
class ImageType(InputType):
[docs]
def __init__(
self,
name=None,
shape=None,
scale=1.0,
bias=None,
color_layout=ColorLayout.RGB,
channel_first=None,
):
"""
Configuration class used for image inputs in Core ML.
Parameters
----------
scale: float or list of floats
The scaling factor for all values in the image channels.
bias: float or list of floats
* If ``color_layout`` is ``ct.colorlayout.GRAYSCALE`` or
``ct.colorlayout.GRAYSCALE_FLOAT16``, bias would be a ``float``.
* If ``color_layout`` is ``ct.colorlayout.RGB`` or ``ct.colorlayout.BGR``,
bias would be a list of ``float``.
color_layout: string or enumeration of type ``ct.colorlayout``
Color layout of the image. Valid values are as follows:
Enumeration (recommended):
* ``ct.colorlayout.RGB``
* ``ct.colorlayout.BGR``
* ``ct.colorlayout.GRAYSCALE``
* ``ct.colorlayout.GRAYSCALE_FLOAT16``
String values (older way to specify):
* ``'G'``: Grayscale (maps to ``ct.colorlayout.GRAYSCALE``)
* ``'RGB'``: [Red, Green, Blue] (maps to ``ct.colorlayout.BGR``)
* ``'BGR'``: [Blue, Green, Red] (maps to ``ct.colorlayout.RGB``)
channel_first: (bool) or None
Set to ``True`` if input format is channel first.
Default format:
* For TensorFlow: channel last (``channel_first=False``).
* For PyTorch: channel first (``channel_first=True``).
"""
super(ImageType, self).__init__(name, shape)
self.scale = scale
msg = "color_layout should be an enum of type ct.colorlayout, i.e. one of: " \
"{ct.colorlayout.RGB, ct.colorlayout.BGR, " \
"ct.colorlayout.GRAYSCALE, ct.colorlayout.GRAYSCALE_FLOAT16}"
if not (isinstance(color_layout, str) or isinstance(color_layout, ColorLayout)):
raise ValueError(msg)
if isinstance(color_layout, str):
if color_layout not in ("G", "RGB", "BGR"):
raise ValueError(msg)
color_layout = ColorLayout(color_layout)
self.color_layout = color_layout
if color_layout == ColorLayout.GRAYSCALE_FLOAT16:
self.dtype = types.fp16
if bias is None:
if color_layout in (ColorLayout.GRAYSCALE, ColorLayout.GRAYSCALE_FLOAT16):
self.bias = 0.0
else:
self.bias = [0.0, 0.0, 0.0]
else:
self.bias = bias
self.channel_first = channel_first
def __repr__(self):
return self.__str__()
def __str__(self):
str_repr = 'ImageType(name={}, shape={}, scale={}, bias={}, ' +\
'color_layout={}, channel_first={})'
return str_repr.format(self.name, self.shape, self.scale, self.bias,
self.color_layout, self.channel_first)
[docs]
class TensorType(InputType):
[docs]
def __init__(self, name=None, shape=None, dtype=None, default_value=None):
"""
Specify a (dense) tensor input.
Parameters
----------
name: str
Input name. Must match an input name in the model (usually the
Placeholder name for TensorFlow or the input name for PyTorch).
The ``name`` is required except for a TensorFlow model in which there is
exactly one input Placeholder.
shape: The shape of the input
- List of positive int or :py:class:`RangeDim`, or
- :py:class:`EnumeratedShapes`
For TensorFlow:
* The ``shape`` is optional. If omitted, the shape is inferred from
TensorFlow graph's Placeholder shape.
For PyTorch:
* The ``shape`` is required.
dtype: np.generic or mil.type type
For example, ``np.int32`` or ``coremltools.converters.mil.mil.types.fp32``
default_value: np.ndarray
If provided, the input is considered optional. At runtime, if the
input is not provided, ``default_value`` is used.
Limitations:
* If ``default_value`` is ``np.ndarray``, all
elements are required to have the same value.
* The ``default_value`` may not be specified if ``shape`` is
:py:class:`EnumeratedShapes`.
Examples
--------
* ``ct.TensorType(name="input", shape=(1, 2, 3))`` implies ``dtype ==
np.float32``
* ``ct.TensorType(name="input", shape=(1, 2, 3), dtype=np.int32)``
* ``ct.TensorType(name="input", shape=(1, 2, 3),
dtype=ct.converters.mil.types.fp32)``
"""
super(TensorType, self).__init__(name, shape)
if dtype is not None:
if types.is_builtin(dtype):
self.dtype = dtype
if dtype not in (
types.int8,
types.uint8,
types.fp16,
types.fp32,
types.fp64,
types.int32,
types.int64,
types.bool,
):
raise TypeError(
"dtype={} is unsupported for inputs/outputs of the model".format(dtype)
)
else:
# Assume dtype is numpy type
try:
self.dtype = types.numpy_type_to_builtin_type(dtype)
except TypeError:
raise TypeError("dtype={} is unsupported".format(dtype))
if dtype not in (np.float16, np.float32, np.float64, float,
np.int32, np.int64, int,
bool, np.bool_):
raise TypeError("dtype={} is unsupported for inputs/outputs of the model".format(dtype))
if default_value is not None:
if isinstance(shape, EnumeratedShapes):
msg = 'TensorType input {} has EnumeratedShapes and ' +\
'may not be optional'
raise ValueError(msg.format(name))
if not isinstance(default_value, np.ndarray):
msg = 'TensorType {} default_value is not np.ndarray'
raise ValueError(msg.format(name))
default_fill_val = default_value.flatten()[0]
if not np.all(default_value == default_fill_val):
msg = 'TensorType {} default_value can only have ' +\
'same entries'
raise ValueError(msg.format(name))
if not self.shape.has_symbolic and list(default_value.shape) != list(
self.shape.symbolic_shape
):
msg = "TensorType {} default_value shape {} != " + "TensorType.shape {}"
raise ValueError(msg.format(name, default_value.shape, self.shape.to_list()))
if (
self.dtype is not None
and types.numpy_type_to_builtin_type(default_value.dtype) != self.dtype
):
msg = "TensorType {} default_value dtype {} != " + "TensorType.dtype {}"
raise ValueError(msg.format(name, default_value.dtype, self.dtype.__type_info__()))
else:
self.dtype = types.numpy_type_to_builtin_type(default_value.dtype)
self.default_value = default_value
def __repr__(self):
return self.__str__()
def __str__(self):
return 'TensorType(name={}, shape={}, dtype={})'.format(self.name,
self.shape,
self.dtype)
[docs]
class StateType(InputType):
SUPPORTED_WRAPPER_TYPE = (
TensorType,
)
[docs]
def __init__(
self,
wrapped_type: type,
name: Optional[str] = None,
):
"""
Specify a model state as a wrapper of a ``TensorType``.
For example, you can use the following code to create a
state type input that wraps a fp16 tensor with shape ``(2, 3)``::
ct.StateType(
wrapped_type=ct.TensorType(
shape=(2, 3),
dtype=np.float16
),
name="state",
)
Parameters
----------
wrapped_type: coremltools.converters.mil.input_types.InputType
- The type wrapped in the state.
- Must be ``TensorType``.
Note that the ``name`` and ``default_value`` of the wrapped ``TensorType`` must not be provided.
name: str
The name of the state.
It must match the key of ``named_buffers()`` in the source TorchScript model.
"""
if not isinstance(wrapped_type, StateType.SUPPORTED_WRAPPER_TYPE):
raise ValueError(
f"StateType only supports {StateType.SUPPORTED_WRAPPER_TYPE}. Got {type(wrapped_type)}."
)
# name and default_value cannot be set
if wrapped_type.name is not None:
raise ValueError("name cannot be set in the state wrapped_type.")
if wrapped_type.default_value is not None:
raise ValueError("default_value cannot be set in the state wrapped_type.")
super(StateType, self).__init__(name, wrapped_type.shape, wrapped_type.dtype)
self.wrapped_type = wrapped_type
def __repr__(self):
return self.__str__()
def __str__(self):
return f"StateType[{self.wrapped_type}]"
[docs]
class RangeDim:
[docs]
def __init__(
self,
lower_bound: int = 1,
upper_bound: int = -1,
default: Optional[int] = None,
symbol: Optional[str] = None,
):
"""
A class for providing a range of accepted shapes.
Parameters
----------
lower_bound:
The minimum valid value for the shape.
upper_bound:
The maximum valid value for the shape.
Set to ``-1`` if there is no upper limit (only works if backend is set to "neuralnetwork").
When backend is set to "mlprogram" during conversion, -1 is not allowed. A finite
positive upper bound must be provided.
default:
The default value that is used for initiating the model, and set in the input shape
field of the model file.
If set to ``None``, ``lower_bound`` would be used as default.
symbol:
Optional symbol name for the dim. Autogenerate a symbol name if not specified.
"""
if symbol is None:
from coremltools.converters.mil.mil import get_new_symbol
self.symbol = get_new_symbol()
else:
from coremltools.converters.mil.mil import Symbol
self.symbol = Symbol(symbol)
self.lower_bound = lower_bound
self.upper_bound = upper_bound
if default is None:
self.default = lower_bound
else:
if default < lower_bound:
raise ValueError(
f"Default value {default} is less than minimum value ({lower_bound}) for range"
)
if default > upper_bound > 0:
raise ValueError(
f"Default value {default} is greater than maximum value ({upper_bound}) for range"
)
self.default = default
def __repr__(self):
return self.__str__()
def __str__(self):
return 'RangeDim(lower_bound={}, upper_bound={}, default={}, symbol="{}")'.format(
self.lower_bound, self.upper_bound, self.default, self.symbol)
[docs]
class Shape:
[docs]
def __init__(self, shape, default=None):
"""
The basic shape class to be set in :py:class:`InputType`.
Parameters
----------
shape: list of (int), symbolic values, RangeDim object
The valid shape of the input.
default: tuple of int or None
The default shape that is used for initiating the model, and set in
the metadata of the model file.
If ``None``, then ``shape`` is used.
"""
from coremltools.converters.mil.mil import get_new_symbol
if not isinstance(shape, (list, tuple)):
msg = "Shape should be list or tuple, got type {} instead"
raise ValueError(msg.format(type(shape)))
self.symbolic_shape = []
shape = list(shape)
for idx, s in enumerate(shape):
if s is None or s == -1:
msg = 'Dimension cannot be None or -1. Use ' +\
'ct.RangeDim for runtime determined dimension. ' +\
'Dim {}: {} ' +\
'See https://coremltools.readme.io/docs/flexible-inputs'
raise ValueError(msg.format(idx, s))
if isinstance(s, RangeDim):
sym = s.symbol
self.symbolic_shape.append(sym)
elif isinstance(s, (np.generic, int)) or is_symbolic(s):
self.symbolic_shape.append(s)
else:
raise ValueError(
"Unknown type {} to build symbolic shape.".format(type(s))
)
self.shape = tuple(shape)
if default is not None:
if not isinstance(default, (list, tuple)):
raise ValueError(
"Default shape should be list or tuple, got type {} instead".format(
type(default)
)
)
for idx, s in enumerate(default):
if not isinstance(
s, (np.generic, int)
) and not is_symbolic(s):
raise ValueError(
"Default shape invalid, got error at index {} which is {}".format(
idx, s
)
)
else:
default = []
for idx, s in enumerate(self.shape):
if isinstance(s, RangeDim):
default.append(s.default)
elif s is None or s == -1:
default.append(self.symbolic_shape[idx])
else:
default.append(s)
self.default = tuple(default)
def __str__(self):
return str(self.shape)
def __repr__(self):
return self.__str__()
@property
def has_symbolic(self):
return any(is_symbolic(s) for s in self.symbolic_shape)
def to_list(self, allow_symbolic=False):
if not allow_symbolic and self.has_symbolic:
return None
return self.symbolic_shape
[docs]
class EnumeratedShapes:
[docs]
def __init__(self, shapes, default=None):
"""
A shape class for setting multiple valid shapes in InputType.
Parameters
----------
shapes: list of Shape objects, or Shape-compatible lists
* The valid shapes of the inputs.
* If input provided is not a :py:class:`Shape` object,
but can be converted to a :py:class:`Shape`,
the :py:class:`Shape` object would be stored in ``shapes`` instead.
default: tuple of int or None
* The default shape that is used for initiating the model, and set in
the metadata of the model file.
* If ``None``, then the first element in ``shapes`` is used.
Examples
--------
.. sourcecode:: python
sample_shape = ct.EnumeratedShapes(
shapes=[(2, 4, 64, 64), (2, 4, 48, 48), (2, 4, 32, 32)], default=(2, 4, 64, 64)
)
my_core_ml_model = ct.convert(
my_model,
inputs=[ct.TensorType(name="sample", shape=sample_shape)],
)
"""
# lazy import to avoid circular import
from coremltools.converters.mil.mil import get_new_symbol
if not isinstance(shapes, (list, tuple)):
raise ValueError(
"EnumeratedShapes should be list or tuple of shape, got type {} instead".format(
type(shapes)
)
)
if len(shapes) < 2:
raise ValueError(
"EnumeratedShapes should be take a list or tuple with len >= 2, got {} instead".format(
len(shapes)
)
)
self.shapes = []
for idx, s in enumerate(shapes):
if isinstance(s, Shape):
self.shapes.append(s)
else:
self.shapes.append(Shape(s))
self.symbolic_shape = self.shapes[0].symbolic_shape
for shape in self.shapes:
for idx, s in enumerate(shape.symbolic_shape):
if is_symbolic(self.symbolic_shape[idx]):
continue
elif is_symbolic(s):
self.symbolic_shape[idx] = s
elif s != self.symbolic_shape[idx]:
self.symbolic_shape[idx] = get_new_symbol()
if default is not None:
if not isinstance(default, (list, tuple)):
raise ValueError(
"Default shape should be list or tuple, got type {} instead".format(
type(default)
)
)
for idx, s in enumerate(default):
if not isinstance(
s, (np.generic, int)
) and not is_symbolic(s):
raise ValueError(
"Default shape invalid, got error at index {} which is {}".format(
idx, s
)
)
else:
default = self.shapes[0].default
self.default = default
def __repr__(self):
return self.__str__()
def __str__(self):
return "EnumeratedShapes(" + str(self.shapes) + ", default=" + str(self.default) + ")"
def _get_shaping_class(shape):
"""
Returns a Shape class or EnumeratedShapes class for `shape`
where `shape` could be lists/tuple/Shape/EnumeratedShapes/etc.
"""
if isinstance(shape, (Shape, EnumeratedShapes)):
return shape
try:
enum_shape = EnumeratedShapes(shape)
return enum_shape
except ValueError:
pass
try:
shape = Shape(shape)
return shape
except ValueError:
pass
raise ValueError("Can't convert to CoreML shaping class from {}.".format(shape))