# Copyright (c) 2023, 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 coremltools.converters.mil.mil import types
from coremltools.converters.mil.mil.input_type import DefaultInputs, InputSpec, TensorInputType
from coremltools.converters.mil.mil.ops.defs._op_reqs import register_op
from coremltools.converters.mil.mil.ops.defs.iOS15.tensor_operation import (
non_maximum_suppression as _nms_iOS15,
)
from coremltools.converters.mil.mil.ops.defs.iOS16.tensor_operation import topk as _topk_iOS16
from coremltools.converters.mil.mil.ops.defs.iOS17 import _IOS17_TARGET
[docs]@register_op(opset_version=_IOS17_TARGET)
class non_maximum_suppression(_nms_iOS15):
"""
Performs non-maximum suppression (NMS) on the boxes according to their intersection-over-union (IoU).
NMS iteratively removes lower-scoring boxes which have an IoU greater than ``iou_threshold`` with
another (higher-scoring) box.
The major differences between this version and the iOS 15 :py:class:`~.iOS15.tensor_operation.non_maximum_suppression`
are as follows:
- The input parameter ``score_threshold`` has been removed.
- The inputs ``boxes`` and ``scores`` are ordered with number of boxes in the last dimension.
- The fourth output containing number of boxes for each batch has been removed.
Parameters
----------
boxes: tensor<[n, 4, B], T> (Required)
* Box coordinates on which to perform NMS. The coordinates are expected in
``CENTER_SIZE_WIDTH_FIRST`` format ``(x, y, width, height)``, in which ``(x, y)`` is the center.
scores: tensor<[n, K, B], T> (Required)
* Scores for each one of the boxes. ``K`` is the number of classes.
iou_threshold: const<T> (Required)
* The intersection over union (IoU) threshold over which boxes are
suppressed. NMS remove all overlapping boxes with ``IoU > iou_threshold``.
max_boxes: const<i32> (Required)
* Maximum number of boxes to select. If the number of surviving boxes are
less, the output is padded up to this number.
per_class_suppression: const<bool> (Optional)
* Defaults to ``False``.
* If ``True``, suppression is performed independently within boxes of each class.
Returns
-------
tensor<[n, 4, max_boxes], T>
* Coordinates of selected boxes.
tensor<[n, K, max_boxes], T>
* Scores of selected boxes.
tensor<[n, max_boxes], i32>
* Indices of selected boxes.
Attributes
----------
T: fp16, fp32
"""
input_spec = InputSpec(
boxes=TensorInputType(type_domain="T"),
scores=TensorInputType(type_domain="T"),
iou_threshold=TensorInputType(const=True, type_domain="T"),
max_boxes=TensorInputType(const=True, type_domain=types.int32),
per_class_suppression=TensorInputType(const=True, optional=True, type_domain=types.bool),
)
def type_inference(self):
boxes_dtype = self.boxes.dtype
scores_dtype = self.scores.dtype
n_batch, n_score_class, _ = self.scores.shape
max_boxes = self.max_boxes.val
return (
types.tensor(boxes_dtype, (n_batch, 4, max_boxes)),
types.tensor(scores_dtype, (n_batch, n_score_class, max_boxes)),
types.tensor(types.int32, (n_batch, max_boxes)),
)
[docs]@register_op(opset_version=_IOS17_TARGET)
class topk(_topk_iOS16):
"""
A version of ``topk`` for iOS 17+. The differences between this version and the
iOS 16 :py:class:`~.iOS16.tensor_operation.topk` are:
- New data type support. The newly added data types are:
- int8, uint8, int16, unint16 for ``x`` and output.
- int8, int16 for ``k``.
- Validation restrictions on the optional ``indices`` output must be either
uint16 or int32.
- A new input parameter ``output_indices_dtype`` has been added
to set the dtype of output ``indices``.
Parameters
----------
x: <\*?, T> (Required)
* Input tensor.
k: const<K> (Optional)
* Defaults to ``1``.
* Number of values/indices to be computed along each axis.
* Set to ``-1`` to select all elements.
axis: const<i32> (Optional)
* Defaults to ``-1`` (last dimension).
* Axis to perform the operation.
ascending: const<bool> (Optional)
* Defaults to ``False``, sort in descending order.
* ``True`` to sort in ascending order.
sort: const<bool> (Optional)
* Defaults to ``True``.
* If ``True``, ``top-k`` elements are themselves sorted.
Otherwise, no particular ordering is guaranteed.
return_indices: const<bool> (Optional)
* Defaults to ``True``.
* If ``True``, returns both values and indices. Otherwise, returns only the ``top-k`` values.
output_indices_dtype: const<str> (Optional, default="int32")
* It can only be set when ``return_indices`` is ``True``.
* This parameter can take ``"int32"`` or ``"uint16"`` as values.
Returns
-------
tensor<\*?, T>
* Values of top/bottom ``k`` elements.
tensor<\*?, U>
* Only returned when ``return_indices = True``
* Indices of the top/bottom ``k`` elements along axis.
* U is int32 or uint16 determined by ``output_indices_dtype`` (int32 by default).
Attributes
----------
T: fp16, fp32, int8, int16, int32, uint8, uint16
K: int8, int16, int32
"""
input_spec = InputSpec(
x=TensorInputType(type_domain="T"),
k=TensorInputType(const=True, optional=True, type_domain="K"),
axis=TensorInputType(const=True, optional=True, type_domain=types.int32),
ascending=TensorInputType(const=True, optional=True, type_domain=types.bool),
sort=TensorInputType(const=True, optional=True, type_domain=types.bool),
return_indices=TensorInputType(const=True, optional=True, type_domain=types.bool),
output_indices_dtype=TensorInputType(const=True, optional=True, type_domain=types.str),
)
type_domains = {
"T": (
types.fp16,
types.fp32,
types.int8,
types.int16,
types.int32,
types.uint8,
types.uint16,
),
"K": (types.int8, types.int16, types.int32),
}
_ALLOWED_OUTPUT_INDICES_DTYPES = {"int32", "uint16"}
def default_inputs(self):
parent_default_inputs = super().default_inputs()
# If return_indices is not set, it is default to True.
# output_indices_dtype can only be set when return_indices = True
if self.return_indices is None or self.return_indices.val:
return parent_default_inputs + DefaultInputs(output_indices_dtype="int32")
return parent_default_inputs
def type_inference(self):
if not self.return_indices.val and self.output_indices_dtype is not None:
raise ValueError(
'In iOS17 topk op, "output_indices_dtype" can only be set when "return_indices=True".'
)
if self.return_indices.val:
if self.output_indices_dtype.val not in self._ALLOWED_OUTPUT_INDICES_DTYPES:
raise ValueError(
f'"topk" op invalid output_indices_dtype: "{self.output_indices_dtype.val}". '
f"Valid options are: {self._ALLOWED_OUTPUT_INDICES_DTYPES}"
)
value_type, indices_type = super().type_inference()
indices_type = types.tensor(
types.string_to_builtin(self.output_indices_dtype.val), indices_type.get_shape()
)
return value_type, indices_type
else:
return super().type_inference()