allow multiple outputs for models (#518)

main
Ben Hoff 7 years ago committed by Nikita Manovich
parent 029d6bff3d
commit e9e1865cb7

@ -39,7 +39,11 @@ class ModelLoader():
_, _, h, w = self._input_layout _, _, h, w = self._input_layout
in_frame = image if image.shape[:-1] == (h, w) else cv2.resize(image, (w, h)) in_frame = image if image.shape[:-1] == (h, w) else cv2.resize(image, (w, h))
in_frame = in_frame.transpose((2, 0, 1)) # Change data layout from HWC to CHW in_frame = in_frame.transpose((2, 0, 1)) # Change data layout from HWC to CHW
return self._net.infer(inputs={self._input_blob_name: in_frame})[self._output_blob_name].copy() results = self._net.infer(inputs={self._input_blob_name: in_frame})
if len(results) == 1:
return results[self._output_blob_name].copy()
else:
return results.copy()
def load_label_map(labels_path): def load_label_map(labels_path):

Loading…
Cancel
Save