From a312dd4a9f2391e6f33ff3248df8891f250ebaa0 Mon Sep 17 00:00:00 2001
From: Nikita Malinin <nikita.malinin@intel.com>
Date: Sun, 20 Feb 2022 09:44:04 +0300
Subject: [PATCH] [POT] IEEngine output data order (#10527)

* IEEngine fix for multiply-output nets

* Update docstrings and docs

* Codestyle changes

* Update docs

* Update docstring

* Pylint
---
 tools/pot/openvino/tools/pot/api/README.md       |  7 +++----
 .../3d_segmentation/3d_segmentation_sample.py    |  5 +++--
 .../pot/openvino/tools/pot/engines/ie_engine.py  | 16 +++++++++-------
 3 files changed, 15 insertions(+), 13 deletions(-)

diff --git a/tools/pot/openvino/tools/pot/api/README.md b/tools/pot/openvino/tools/pot/api/README.md
index c0a66988a09..728711284f3 100644
--- a/tools/pot/openvino/tools/pot/api/README.md
+++ b/tools/pot/openvino/tools/pot/api/README.md
@@ -222,14 +222,13 @@ This class support inference in synchronous and asynchronous modes and can be re
 with some modifications, e.g. in case of custom post-processing of inference results.
 
 The following methods can be overridden in subclasses:
-- `postprocess_output(outputs, metadata)` - processes raw model output using the image metadata obtained during 
-data loading.<br><br>
+- `postprocess_output(outputs, metadata)` - Processes model output data using the image metadata obtained during data loading.<br><br>
   *Parameters*
-  - `outputs` - raw output of the model.
+  - `outputs` - dictionary of output data per output name.
   - `metadata` - information about the data used for inference.
   
   *Return*
-  - post-processed model output
+  - list of the output data in an order expected by the accuracy metric if any is used
   
 `IEEngine` supports data returned by `DataLoader` in the format:
 ```
diff --git a/tools/pot/openvino/tools/pot/api/samples/3d_segmentation/3d_segmentation_sample.py b/tools/pot/openvino/tools/pot/api/samples/3d_segmentation/3d_segmentation_sample.py
index 1b8b8d03079..2056d7e215a 100644
--- a/tools/pot/openvino/tools/pot/api/samples/3d_segmentation/3d_segmentation_sample.py
+++ b/tools/pot/openvino/tools/pot/api/samples/3d_segmentation/3d_segmentation_sample.py
@@ -187,12 +187,13 @@ class SegmentationEngine(IEEngine):
         """
         Processes model raw output for future metric and loss calculation.
         Uses image metadata that can be passed using dataloader.
-        :param outputs: network infer result in format of numpy ndarray (batch x image shape)
+        :param outputs: network infer result in the format of dictionary numpy ndarray
+                        by layer name (batch x image shape)
         :param metadata: dictionary of image metadata
         :return: processed numpy ndarray with the same shape as the original output
         """
         processed_outputs = []
-        for output, meta in zip(outputs, metadata):
+        for output, meta in zip(outputs.values(), metadata):
             # Resize to bounding box size and extend to mask size
             low = meta['bbox'][0]
             high = meta['bbox'][1]
diff --git a/tools/pot/openvino/tools/pot/engines/ie_engine.py b/tools/pot/openvino/tools/pot/engines/ie_engine.py
index 9faffd420a0..bff5a1e0d4d 100644
--- a/tools/pot/openvino/tools/pot/engines/ie_engine.py
+++ b/tools/pot/openvino/tools/pot/engines/ie_engine.py
@@ -143,8 +143,12 @@ class IEEngine(Engine):
 
     @staticmethod
     def postprocess_output(outputs, _metadata):
-        """ Processes raw model output using the image metadata obtained during data loading """
-        return outputs
+        """ Processes model output data using the image metadata obtained during data loading
+        :param outputs: dictionary of output data per output name
+        :param _metadata: metadata obtained during data loading
+        :return: list of the output data in an order expected by the accuracy metric if any is used
+        """
+        return list(outputs.values())
 
     def _reset(self):
         """ Resets collected statistics """
@@ -182,14 +186,12 @@ class IEEngine(Engine):
                                      annotations=batch_annotations)
 
         # Postprocess network output
-        outputs = process_raw_output(predictions)
-        output = outputs[self._output_layers[0]]
-        outputs[self._output_layers[0]] = self.postprocess_output(output, batch_meta)
+        processed_outputs = process_raw_output(predictions)
+        outputs = {name: processed_outputs[name] for name in self._output_layers}
+        logits = self.postprocess_output(outputs, batch_meta)
 
         # Update metrics
         if batch_annotations:
-            # TODO: Create some kind of an order for the correct metric calculation
-            logits = [outputs[name] for name in self._output_layers]  # output_layers are in a random order
             self._update_metrics(output=logits, annotations=batch_annotations,
                                  need_metrics_per_sample=need_metrics_per_sample)