def convert_NMSLayer(scope, operator, container):
# type: (keras2onnx.common.InterimContext, keras2onnx.common.Operator, keras2onnx.common.OnnxObjectContainer) -> None
box_transpose = scope.get_unique_variable_name(
operator.inputs[0].full_name + '_tx')
score_transpose = scope.get_unique_variable_name(
operator.inputs[1].full_name + '_tx')
apply_identity(
scope, operator.inputs[0].full_name, box_transpose, container)
apply_transpose(
scope, operator.inputs[1].full_name, score_transpose, container, perm=[1, 0])
box_batch = scope.get_unique_variable_name(
operator.inputs[0].full_name + '_btc')
score_batch = scope.get_unique_variable_name(
operator.inputs[1].full_name + '_btc')
container.add_node("Unsqueeze", box_transpose,
box_batch, op_version=operator.target_opset, axes=[0])
container.add_node("Unsqueeze", score_transpose,
score_batch, op_version=operator.target_opset, axes=[0])
layer = operator.raw_operator # type: YOLONMSLayer
max_output_size = scope.get_unique_variable_name('max_output_size')
iou_threshold = scope.get_unique_variable_name('iou_threshold')
score_threshold = scope.get_unique_variable_name('layer.score_threshold')
container.add_initializer(max_output_size, onnx_proto.TensorProto.INT64,
[], [layer.max_boxes])
container.add_initializer(iou_threshold, onnx_proto.TensorProto.FLOAT,
[], [layer.iou_threshold])
container.add_initializer(score_threshold, onnx_proto.TensorProto.FLOAT,
[], [layer.score_threshold])
cast_name = scope.get_unique_variable_name('casted')
nms_node = next((nd_ for nd_ in operator.nodelist if nd_.type ==
'NonMaxSuppressionV3'), operator.nodelist[0])
container.add_node("NonMaxSuppression",
[box_batch, score_batch, max_output_size,
iou_threshold, score_threshold],
cast_name,
op_version=operator.target_opset,
name=nms_node.name)
cast_batch = scope.get_unique_variable_name(
operator.output_full_names[2] + '_btc')
container.add_node("Unsqueeze", cast_name,
cast_batch, op_version=operator.target_opset, axes=[0])
apply_cast(scope, cast_batch,
operator.output_full_names[2], container, to=onnx_proto.TensorProto.INT32)
apply_identity(scope, box_batch, operator.output_full_names[0], container)
apply_identity(scope, score_batch,
operator.output_full_names[1], container)
def convert_DetectionLayer(scope, operator, container):
# type: (keras2onnx.common.InterimContext, keras2onnx.common.Operator, keras2onnx.common.OnnxObjectContainer) -> None
DETECTION_MAX_INSTANCES = 100
DETECTION_NMS_THRESHOLD = 0.3
DETECTION_MIN_CONFIDENCE = 0.7
oopb = OnnxOperatorBuilder(container, scope)
box_transpose = scope.get_unique_variable_name(operator.inputs[0].full_name + '_tx')
score_transpose = scope.get_unique_variable_name(operator.inputs[1].full_name + '_tx')
# apply_transpose(scope, operator.inputs[0].full_name, box_transpose, container, perm=[2, 0, 1])
apply_identity(scope, operator.inputs[0].full_name, box_transpose, container)
# output shape: [num_batches, spatial_dimension, 4]
score_identity = scope.get_unique_variable_name(operator.inputs[1].full_name + '_id')
apply_identity(scope, operator.inputs[1].full_name, score_identity, container)
# output shape: [num_batches, spatial_dimension, num_classes]
deltas_transpose = scope.get_unique_variable_name(operator.inputs[2].full_name + '_tx')
apply_identity(scope, operator.inputs[2].full_name, deltas_transpose, container)
image_meta = scope.get_unique_variable_name(operator.inputs[3].full_name + '_tx')
apply_identity(scope, operator.inputs[3].full_name, image_meta, container)
windows_transpose = norm_boxes_graph(scope, operator, container, oopb, image_meta)
delta_mul_output = convert_apply_box_deltas_graph(scope, operator, container, oopb, box_transpose, score_identity, deltas_transpose, windows_transpose)
sliced_score = oopb.add_node('Slice',
[score_identity,
('_start', oopb.int64, np.array([1], dtype='int64')),
('_end', oopb.int64, np.array([81], dtype='int64')),
('_axes', oopb.int64, np.array([2], dtype='int64'))
],
operator.inputs[1].full_name + '_sliced')
apply_transpose(scope, sliced_score, score_transpose, container, perm=[0, 2, 1])
# output shape: [num_batches, num_classes, spatial_dimension]
max_output_size = scope.get_unique_variable_name('max_output_size')
iou_threshold = scope.get_unique_variable_name('iou_threshold')
score_threshold = scope.get_unique_variable_name('layer.score_threshold')
container.add_initializer(max_output_size, onnx_proto.TensorProto.INT64,
[], [DETECTION_MAX_INSTANCES])
container.add_initializer(iou_threshold, onnx_proto.TensorProto.FLOAT,
[], [DETECTION_NMS_THRESHOLD])
container.add_initializer(score_threshold, onnx_proto.TensorProto.FLOAT,
[], [DETECTION_MIN_CONFIDENCE])
nms_node = next((nd_ for nd_ in operator.nodelist if nd_.type == 'NonMaxSuppressionV3'), operator.nodelist[0])
nms_output = scope.get_unique_variable_name(operator.output_full_names[0] + '_nms')
container.add_node("NonMaxSuppression",
[delta_mul_output, score_transpose, max_output_size, iou_threshold, score_threshold],
nms_output,
op_version=operator.target_opset,
name=nms_node.name)
add_init = scope.get_unique_variable_name('add')
container.add_initializer(add_init, onnx_proto.TensorProto.INT64,
[1, 3], [0, 1, 0])
nms_output_add = scope.get_unique_variable_name(operator.output_full_names[0] + '_class_add')
container.add_node("Add",
[nms_output, add_init],
nms_output_add,
op_version=operator.target_opset,
name=nms_node.name + '_class_idx_add')
starts_init = scope.get_unique_variable_name('starts')
ends_init = scope.get_unique_variable_name('ends')
axes_init = scope.get_unique_variable_name('axes')
container.add_initializer(starts_init, onnx_proto.TensorProto.INT32,
[1], [1])
container.add_initializer(ends_init, onnx_proto.TensorProto.INT32,
[1], [2])
container.add_initializer(axes_init, onnx_proto.TensorProto.INT32,
[1], [1])
class_idx_output = scope.get_unique_variable_name(operator.output_full_names[0] + '_class_idx')
container.add_node("Slice",
[nms_output_add, starts_init, ends_init, axes_init],
class_idx_output,
op_version=operator.target_opset,
name=nms_node.name+'_class_idx')
# output shape: [num_selected_indices, 1]
starts_init_2 = scope.get_unique_variable_name('starts')
ends_init_2 = scope.get_unique_variable_name('ends')
axes_init_2 = scope.get_unique_variable_name('axes')
container.add_initializer(starts_init_2, onnx_proto.TensorProto.INT32,
[1], [2])
container.add_initializer(ends_init_2, onnx_proto.TensorProto.INT32,
[1], [3])
container.add_initializer(axes_init_2, onnx_proto.TensorProto.INT32,
[1], [1])
box_idx_output = scope.get_unique_variable_name(operator.output_full_names[0] + '_box_idx')
container.add_node("Slice",
[nms_output_add, starts_init_2, ends_init_2, axes_init_2],
box_idx_output,
op_version=operator.target_opset,
name=nms_node.name + '_box_idx')
# output shape: [num_selected_indices, 1]
box_idx_squeeze = scope.get_unique_variable_name(operator.output_full_names[0] + '_box_idx_squeeze')
attrs = {'axes': [1]}
container.add_node("Squeeze",
box_idx_output,
box_idx_squeeze,
op_version=operator.target_opset,
name=nms_node.name + '_box_idx_squeeze', **attrs)
# output shape: [num_selected_indices]
starts_init_3 = scope.get_unique_variable_name('starts')
ends_init_3 = scope.get_unique_variable_name('ends')
axes_init_3 = scope.get_unique_variable_name('axes')
step_init_3 = scope.get_unique_variable_name('steps')
container.add_initializer(starts_init_3, onnx_proto.TensorProto.INT32,
[1], [2])
container.add_initializer(ends_init_3, onnx_proto.TensorProto.INT32,
[1], [0])
container.add_initializer(axes_init_3, onnx_proto.TensorProto.INT32,
[1], [1])
container.add_initializer(step_init_3, onnx_proto.TensorProto.INT32,
[1], [-1])
from keras2onnx.common.data_types import Int32TensorType, FloatTensorType
class_box_idx_output = scope.get_local_variable_or_declare_one(operator.output_full_names[0] + '_class_box_idx',
type=Int32TensorType(shape=[None, 2]))
container.add_node("Slice",
[nms_output_add, starts_init_3, ends_init_3, axes_init_3, step_init_3],
class_box_idx_output.full_name,
op_version=operator.target_opset,
name=nms_node.name + '_class_box_idx')
# output shape: [num_selected_indices, 2]
box_squeeze = scope.get_unique_variable_name(operator.output_full_names[0] + '_box_squeeze')
attrs = {'axes': [0]}
container.add_node("Squeeze",
delta_mul_output,
box_squeeze,
op_version=operator.target_opset,
name=nms_node.name + '_box_squeeze', **attrs)
# output shape: [spatial_dimension, 4]
score_squeeze = scope.get_local_variable_or_declare_one(operator.output_full_names[0] + '_score_squeeze',
type=FloatTensorType(shape=[None]))
attrs = {'axes': [0]}
container.add_node("Squeeze",
score_identity,
score_squeeze.full_name,
op_version=operator.target_opset,
name=nms_node.name + '_score_squeeze', **attrs)
# output shape: [spatial_dimension, num_classes]
box_gather = scope.get_unique_variable_name(operator.output_full_names[0] + '_box_gather')
attrs = {'axis': 0}
container.add_node("Gather",
[box_squeeze, box_idx_squeeze],
box_gather,
op_version=operator.target_opset,
name=nms_node.name + '_box_gather', **attrs)
# output shape: [num_selected_indices, 4]
score_gather = scope.get_unique_variable_name(operator.output_full_names[0] + '_score_gather')
container.add_node("GatherND",
[score_squeeze.full_name, class_box_idx_output.full_name],
score_gather,
op_version=operator.target_opset, op_domain='com.microsoft',
name=nms_node.name + '_score_gather')
# output shape: [num_selected_indices]
score_gather_unsqueeze = scope.get_unique_variable_name(operator.output_full_names[0] + '_score_gather_unsqueeze')
attrs = {'axes': [1]}
container.add_node("Unsqueeze",
score_gather,
score_gather_unsqueeze,
op_version=operator.target_opset,
name=nms_node.name + '_score_gather_unsqueeze', **attrs)
# output shape: [num_selected_indices, 1]
top_k_var = scope.get_unique_variable_name('topK')
container.add_initializer(top_k_var, onnx_proto.TensorProto.FLOAT,
[1], [100.0])
score_gather_shape = oopb.add_node('Shape',
[score_gather],
operator.inputs[1].full_name + '_score_gather_shape')
attrs = {'to': 1}
scope_gather_float = oopb.add_node('Cast',
[score_gather_shape],
operator.inputs[1].full_name + '_scope_gather_float', **attrs)
top_k_min = oopb.add_node('Min',
[scope_gather_float, top_k_var],
operator.inputs[1].full_name + '_top_k_min')
attrs = {'to': 7}
top_k_min_int = oopb.add_node('Cast',
[top_k_min],
operator.inputs[1].full_name + '_top_k_min_int', **attrs)
score_top_k_output_val = scope.get_unique_variable_name(operator.output_full_names[0] + '_score_top_k_output_val')
# output shape: [num_top_K]
score_top_k_output_idx = scope.get_unique_variable_name(operator.output_full_names[0] + '_score_top_k_output_idx')
# output shape: [num_top_K]
attrs = {'axis': 0}
container.add_node('TopK', [score_gather, top_k_min_int], [score_top_k_output_val, score_top_k_output_idx],
op_version=operator.target_opset,
name=nms_node.name + '_topK', **attrs)
class_idx_cast = scope.get_unique_variable_name(operator.output_full_names[0] + '_class_idx_cast')
attrs = {'to': 1}
container.add_node('Cast', class_idx_output, class_idx_cast, op_version=operator.target_opset,
name=nms_node.name+'_class_idx_cast', **attrs)
# output shape: [num_selected_indices, 1]
concat_var = scope.get_unique_variable_name(operator.output_full_names[0] + '_concat_var')
concat_node = next((nd_ for nd_ in operator.nodelist if nd_.type == 'Concat'), operator.nodelist[0])
attrs = {'axis': 1}
container.add_node("Concat",
[box_gather, class_idx_cast, score_gather_unsqueeze],
concat_var,
op_version=operator.target_opset,
name=concat_node.name, **attrs)
# output shape: [num_selected_indices, 6]
all_gather = scope.get_unique_variable_name(operator.output_full_names[0] + '_all_gather')
attrs = {'axis': 0}
container.add_node("Gather",
[concat_var, score_top_k_output_idx],
all_gather,
op_version=operator.target_opset,
name=nms_node.name + '_all_gather', **attrs)
# output shape: [num_top_K, 6]
padded_result = oopb.add_node('Pad',
[all_gather,
np.array([0, 0, DETECTION_MAX_INSTANCES, 0],
dtype=np.int64)],
nms_node.name + '_padded_result',
op_domain='com.microsoft')
detection_final = oopb.add_node('Slice',
[padded_result,
('_start', oopb.int64, np.array([0], dtype='int64')),
('_end', oopb.int64, np.array([DETECTION_MAX_INSTANCES], dtype='int64')),
('_axes', oopb.int64, np.array([0], dtype='int64'))
],
nms_node.name + '_detection_final'
)
attrs = {'axes': [0]}
container.add_node("Unsqueeze",
detection_final,
operator.output_full_names[0],
op_version=operator.target_opset,
name=nms_node.name + '_concat_unsqueeze', **attrs)