Skip to content

Commit 3826bfd

Browse files
benhoffChris Lee-Messer
authored and
Chris Lee-Messer
committed
allow security segmentation models to be used in auto annotation (cvat-ai#759)
1 parent c043228 commit 3826bfd

File tree

1 file changed

+8
-3
lines changed

1 file changed

+8
-3
lines changed

cvat/apps/auto_annotation/model_loader.py

+8-3
Original file line numberDiff line numberDiff line change
@@ -31,14 +31,19 @@ def __init__(self, model, weights):
3131

3232
iter_inputs = iter(network.inputs)
3333
self._input_blob_name = next(iter_inputs)
34+
self._input_info_name = ''
3435
self._output_blob_name = next(iter(network.outputs))
3536

3637
self._require_image_info = False
3738

39+
info_names = ('image_info', 'im_info')
40+
3841
# NOTE: handeling for the inclusion of `image_info` in OpenVino2019
39-
if 'image_info' in network.inputs:
42+
if any(s in network.inputs for s in info_names):
4043
self._require_image_info = True
41-
if self._input_blob_name == 'image_info':
44+
self._input_info_name = set(network.inputs).intersection(info_names)
45+
self._input_info_name = self._input_info_name.pop()
46+
if self._input_blob_name in info_names:
4247
self._input_blob_name = next(iter_inputs)
4348

4449
self._net = plugin.load(network=network, num_requests=2)
@@ -56,7 +61,7 @@ def infer(self, image):
5661
info[0, 1] = w
5762
# frame number
5863
info[0, 2] = 1
59-
inputs['image_info'] = info
64+
inputs[self._input_info_name] = info
6065

6166
results = self._net.infer(inputs)
6267
if len(results) == 1:

0 commit comments

Comments
 (0)