def test_register_caffe_python_extractor_by_name(self, op_mock):
op_mock.op = 'TestLayer'
name = 'myTestLayer'
register_caffe_python_extractor(op_mock, name)
self.assertIn(name, CaffePythonFrontExtractorOp.registered_ops)
def test_register_caffe_python_extractor_by_op(self, op_mock):
op_mock.op = 'TestLayer'
register_caffe_python_extractor(op_mock)
self.assertIn(op_mock.op, CaffePythonFrontExtractorOp.registered_ops)
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""
import networkx as nx
from extensions.ops.proposal import ProposalOp
from mo.front.caffe.extractor import register_caffe_python_extractor
from mo.graph.graph import Graph
from mo.ops.op import Op
class ProposalPythonExampleOp(Op):
op = 'Proposal'
def __init__(self, graph: Graph, attrs: dict):
mandatory_props = {
'type': __class__.op,
'op': __class__.op,
'post_nms_topn': 300,
'infer': ProposalOp.proposal_infer
}
super().__init__(graph, mandatory_props, attrs)
register_caffe_python_extractor(ProposalPythonExampleOp,
'rpn.proposal_layer.ProposalLayer.example')
Op.excluded_classes.append(ProposalPythonExampleOp)
