allow security segmentation models to be used in auto annotation (#759)

main
Ben Hoff 6 years ago committed by Nikita Manovich
parent e9e52f35cb
commit 2640e6c579

@ -31,14 +31,19 @@ class ModelLoader():
iter_inputs = iter(network.inputs)
self._input_blob_name = next(iter_inputs)
self._input_info_name = ''
self._output_blob_name = next(iter(network.outputs))
self._require_image_info = False
info_names = ('image_info', 'im_info')
# NOTE: handeling for the inclusion of `image_info` in OpenVino2019
if 'image_info' in network.inputs:
if any(s in network.inputs for s in info_names):
self._require_image_info = True
if self._input_blob_name == 'image_info':
self._input_info_name = set(network.inputs).intersection(info_names)
self._input_info_name = self._input_info_name.pop()
if self._input_blob_name in info_names:
self._input_blob_name = next(iter_inputs)
self._net = plugin.load(network=network, num_requests=2)
@ -56,7 +61,7 @@ class ModelLoader():
info[0, 1] = w
# frame number
info[0, 2] = 1
inputs['image_info'] = info
inputs[self._input_info_name] = info
results = self._net.infer(inputs)
if len(results) == 1:

Loading…
Cancel
Save