model-optimizer/extensions/front/tf/SSDToolboxDetectionOutput.py

   1 """
   2  Copyright (c) 2018-2019 Intel Corporation
   3
   4  Licensed under the Apache License, Version 2.0 (the "License");
   5  you may not use this file except in compliance with the License.
   6  You may obtain a copy of the License at
   7
   8       http://www.apache.org/licenses/LICENSE-2.0
   9
  10  Unless required by applicable law or agreed to in writing, software
  11  distributed under the License is distributed on an "AS IS" BASIS,
  12  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13  See the License for the specific language governing permissions and
  14  limitations under the License.
  15 """
  16
  17 import numpy as np
  18
  19 from extensions.front.standalone_const_eraser import StandaloneConstEraser
  20 from extensions.ops.DetectionOutput import DetectionOutput
  21 from mo.front.subgraph_matcher import SubgraphMatch
  22 from mo.front.tf.replacement import FrontReplacementFromConfigFileSubGraph
  23 from mo.graph.graph import Node, Graph
  24 from mo.ops.op import PermuteAttrs
  25 from mo.ops.output import Output
  26 from mo.ops.reshape import Reshape
  27
  28
  29 class SSDToolboxDetectionOutputReplacement(FrontReplacementFromConfigFileSubGraph):
  30     replacement_id = 'SSDToolboxDetectionOutput'
  31
  32     def run_before(self):
  33         return [StandaloneConstEraser]
  34
  35     def nodes_to_remove(self, graph: Graph, match: SubgraphMatch):
  36         return []
  37
  38     def generate_sub_graph(self, graph: Graph, match: SubgraphMatch):
  39         # IE DetectionOutput layer consumes flattened confidences and locations tensors.
  40         # That is why we add reshapes before them.
  41         locs_node = match.single_input_node(0)
  42         conf_node = match.single_input_node(1)
  43         prior_boxes_node = match.single_input_node(2)
  44
  45         locs_out_nodes = locs_node[0].out_nodes()
  46         assert len(locs_out_nodes) == 1
  47         locs_out_node = locs_out_nodes[list(locs_out_nodes.keys())[0]]
  48         assert locs_out_node.op == "OpOutput", locs_out_node.op
  49         graph.remove_node(locs_out_node.id)
  50
  51         conf_out_nodes = conf_node[0].out_nodes()
  52         assert len(conf_out_nodes) == 1
  53         conf_out_node = conf_out_nodes[list(conf_out_nodes.keys())[0]]
  54         assert conf_out_node.op == "OpOutput", conf_out_node.op
  55         graph.remove_node(conf_out_node.id)
  56
  57         # reshape operation to flatten confidence tensor
  58         reshape_loc_op = Reshape(graph, {'dim': np.array([0, -1])})
  59         reshape_loc_node = reshape_loc_op.create_node([locs_node], dict(name='DetectionOutput_Reshape_loc_'))
  60
  61         # reshape operation to flatten confidence tensor
  62         reshape_conf_op = Reshape(graph, {'dim': np.array([0, -1])})
  63         reshape_conf_node = reshape_conf_op.create_node([conf_node], dict(name='DetectionOutput_Reshape_conf_'))
  64
  65         # remove the OpOutput node after the priors node
  66         assert prior_boxes_node[0].out_node().op == "OpOutput"
  67         graph.remove_node(prior_boxes_node[0].out_node().id)
  68
  69         # reshape operation for prior boxes tensor
  70         reshape_priors_op = Reshape(graph, {'dim': np.array([1, 2, -1])})
  71         reshape_priors_node = reshape_priors_op.create_node([prior_boxes_node],
  72                                                             dict(name='DetectionOutput_Reshape_priors_'))
  73         # create Detection Output node with three inputs: locations, confidences and prior boxes
  74         detection_output_op = DetectionOutput(graph, match.custom_replacement_desc.custom_attributes)
  75         detection_output_node = detection_output_op.create_node(
  76             [reshape_loc_node, reshape_conf_node, reshape_priors_node],
  77             dict(name=detection_output_op.attrs['type'] + '_'))
  78         PermuteAttrs.set_permutation(reshape_priors_node, detection_output_node, None)
  79
  80         # create Output node to mark DetectionOutput as a graph output operation
  81         output_op = Output(graph)
  82         output_op.create_node([detection_output_node], dict(name='sink_'))
  83         return {}