[Segmentation] Enable feature vector output for MPA Segmentation

external/mmsegmentation/segmentation_tasks/apis/segmentation/model_wrappers/blur.py

@@ -12,9 +12,11 @@
 # See the License for the specific language governing permissions
 # and limitations under the License.
 
+from typing import Any, Dict, Optional, Union, Iterable
+import warnings
+
 import cv2
 import numpy as np
-from typing import Any, Dict, Optional
 
 from openvino.model_zoo.model_api.models import SegmentationModel
 from openvino.model_zoo.model_api.models.types import NumericalValue
@@ -23,6 +25,16 @@
 from ote_sdk.utils.segmentation_utils import create_hard_prediction_from_soft_prediction
 
 
+@check_input_parameters_type()
+def get_actmap(
+    features: Union[np.ndarray, Iterable, int, float], output_res: Union[tuple, list]
+):
+    am = cv2.resize(features, output_res)
+    am = cv2.applyColorMap(am, cv2.COLORMAP_JET)
+    am = cv2.cvtColor(am, cv2.COLOR_BGR2RGB)
+    return am
+
+
 class BlurSegmentation(SegmentationModel):
     __model__ = 'blur_segmentation'
 
@@ -60,17 +72,24 @@ def _get_outputs(self):
     def postprocess(self, outputs: Dict[str, np.ndarray], metadata: Dict[str, Any]):
         predictions = outputs[self.output_blob_name].squeeze()
         soft_prediction = np.transpose(predictions, axes=(1, 2, 0))
-        feature_vector = outputs.get('repr_vector', None)  # Optional output
 
         hard_prediction = create_hard_prediction_from_soft_prediction(
             soft_prediction=soft_prediction,
             soft_threshold=self.soft_threshold,
             blur_strength=self.blur_strength
         )
         hard_prediction = cv2.resize(hard_prediction, metadata['original_shape'][1::-1], 0, 0, interpolation=cv2.INTER_NEAREST)
-        soft_prediction = cv2.resize(soft_prediction, metadata['original_shape'][1::-1], 0, 0, interpolation=cv2.INTER_NEAREST)
-
-        metadata['soft_predictions'] = soft_prediction
-        metadata['feature_vector'] = feature_vector
+
+        if 'feature_vector' not in outputs or 'saliency_map' not in outputs:
+            warnings.warn('Could not find Feature Vector and Saliency Map in OpenVINO output. '
+                'Please rerun OpenVINO export or retrain the model.')
+            metadata["saliency_map"] = None
+            metadata["feature_vector"] = None
+        else:
+            metadata["saliency_map"] = get_actmap(
+                outputs["saliency_map"][0],
+                (metadata["original_shape"][1], metadata["original_shape"][0]),
+            )
+            metadata["feature_vector"] = outputs["feature_vector"] 
 
         return hard_prediction

external/mmsegmentation/segmentation_tasks/apis/segmentation/openvino_task.py

@@ -105,12 +105,11 @@ def pre_process(self, image: np.ndarray) -> Tuple[Dict[str, np.ndarray], Dict[st
     @check_input_parameters_type()
     def post_process(self, prediction: Dict[str, np.ndarray], metadata: Dict[str, Any]) -> AnnotationSceneEntity:
         hard_prediction = self.model.postprocess(prediction, metadata)
-        soft_prediction = metadata['soft_predictions']
         feature_vector = metadata['feature_vector']
-
+        saliency_map = metadata['saliency_map']
         predicted_scene = self.converter.convert_to_annotation(hard_prediction, metadata)
 
-        return predicted_scene, soft_prediction, feature_vector
+        return predicted_scene, feature_vector, saliency_map
 
     @check_input_parameters_type()
     def forward(self, inputs: Dict[str, np.ndarray]) -> Dict[str, np.ndarray]:
@@ -165,38 +164,25 @@ def infer(self,
               inference_parameters: Optional[InferenceParameters] = None) -> DatasetEntity:
         if inference_parameters is not None:
             update_progress_callback = inference_parameters.update_progress
-            save_mask_visualization = not inference_parameters.is_evaluation
+            dump_saliency_map = not inference_parameters.is_evaluation
         else:
             update_progress_callback = default_progress_callback
-            save_mask_visualization = True
+            dump_saliency_map = True
 
         dataset_size = len(dataset)
         for i, dataset_item in enumerate(dataset, 1):
-            predicted_scene, soft_prediction, feature_vector = self.inferencer.predict(dataset_item.numpy)
+            predicted_scene, feature_vector, saliency_map = self.inferencer.predict(dataset_item.numpy)
             dataset_item.append_annotations(predicted_scene.annotations)
 
             if feature_vector is not None:
                 feature_vector_media = TensorEntity(name="representation_vector", numpy=feature_vector.reshape(-1))
                 dataset_item.append_metadata_item(feature_vector_media, model=self.model)
 
-            if save_mask_visualization:
-                for label_index, label in self._label_dictionary.items():
-                    if label_index == 0:
-                        continue
-
-                    if len(soft_prediction.shape) == 3:
-                        current_label_soft_prediction = soft_prediction[:, :, label_index]
-                    else:
-                        current_label_soft_prediction = soft_prediction
-
-                    class_act_map = get_activation_map(current_label_soft_prediction)
-                    result_media = ResultMediaEntity(name=f'{label.name}',
-                                                     type='Soft Prediction',
-                                                     label=label,
-                                                     annotation_scene=dataset_item.annotation_scene,
-                                                     roi=dataset_item.roi,
-                                                     numpy=class_act_map)
-                    dataset_item.append_metadata_item(result_media, model=self.model)
+            if dump_saliency_map and saliency_map is not None:
+                saliency_map_media = ResultMediaEntity(name="saliency_map", type="Saliency map",
+                                                annotation_scene=dataset_item.annotation_scene, 
+                                                numpy=saliency_map, roi=dataset_item.roi)
+                dataset_item.append_metadata_item(saliency_map_media, model=self.model)
 
             update_progress_callback(int(i / dataset_size * 100))
 

external/model-preparation-algorithm/mpa_tasks/apis/segmentation/task.py

@@ -16,6 +16,7 @@
 from mpa import MPAConstants
 from mpa_tasks.apis import BaseTask, TrainType
 from mpa_tasks.apis.segmentation import SegmentationConfig
+from mpa_tasks.utils.data_utils import get_actmap
 from mpa.utils.config_utils import MPAConfig
 from mpa.utils.logger import get_logger
 from ote_sdk.configuration import cfg_helper
@@ -71,6 +72,8 @@ def infer(self,
               inference_parameters: Optional[InferenceParameters] = None
               ) -> DatasetEntity:
         logger.info('infer()')
+        dump_features = True
+        dump_saliency_map = not inference_parameters.is_evaluation if inference_parameters else True        
 
         if inference_parameters is not None:
             update_progress_callback = inference_parameters.update_progress
@@ -84,15 +87,13 @@ def infer(self,
         stage_module = 'SegInferrer'
         self._data_cfg = self._init_test_data_cfg(dataset)
         self._label_dictionary = dict(enumerate(self._labels, 1))
-        results = self._run_task(stage_module, mode='train', dataset=dataset)
+        results = self._run_task(stage_module, mode='train', dataset=dataset, dump_features=dump_features, 
+                                  dump_saliency_map=dump_saliency_map)
         logger.debug(f'result of run_task {stage_module} module = {results}')
         predictions = results['outputs']
-        # TODO: feature maps should be came from the inference results
-        featuremaps = [None for _ in range(len(predictions))]
-        for i in range(len(dataset)):
-            result, featuremap, dataset_item = predictions[i], featuremaps[i], dataset[i]
-            self._add_predictions_to_dataset_item(result, featuremap, dataset_item,
-                                                  save_mask_visualization=not is_evaluation)
+        prediction_results = zip(predictions['eval_predictions'], predictions['feature_vectors'], 
+                                  predictions['saliency_maps'])
+        self._add_predictions_to_dataset(prediction_results, dataset, dump_saliency_map=not is_evaluation)
         return dataset
 
     def evaluate(self,
@@ -208,45 +209,33 @@ def _init_test_data_cfg(self, dataset: DatasetEntity):
         )
         return data_cfg
 
-    def _add_predictions_to_dataset_item(self, prediction, feature_vector, dataset_item, save_mask_visualization):
-        soft_prediction = np.transpose(prediction, axes=(1, 2, 0))
-        hard_prediction = create_hard_prediction_from_soft_prediction(
-            soft_prediction=soft_prediction,
-            soft_threshold=self._hyperparams.postprocessing.soft_threshold,
-            blur_strength=self._hyperparams.postprocessing.blur_strength,
-        )
-        annotations = create_annotation_from_segmentation_map(
-            hard_prediction=hard_prediction,
-            soft_prediction=soft_prediction,
-            label_map=self._label_dictionary,
-        )
-        dataset_item.append_annotations(annotations=annotations)
-
-        if feature_vector is not None:
-            active_score = TensorEntity(name="representation_vector", numpy=feature_vector)
-            dataset_item.append_metadata_item(active_score, model=self._task_environment.model)
-
-        if save_mask_visualization:
-            for label_index, label in self._label_dictionary.items():
-                if label_index == 0:
-                    continue
-
-                if len(soft_prediction.shape) == 3:
-                    current_label_soft_prediction = soft_prediction[:, :, label_index]
-                else:
-                    current_label_soft_prediction = soft_prediction
-                min_soft_score = np.min(current_label_soft_prediction)
-                max_soft_score = np.max(current_label_soft_prediction)
-                factor = 255.0 / (max_soft_score - min_soft_score + 1e-12)
-                result_media_numpy = (factor * (current_label_soft_prediction - min_soft_score)).astype(np.uint8)
-
-                result_media = ResultMediaEntity(name=f'{label.name}',
-                                                 type='Soft Prediction',
-                                                 label=label,
-                                                 annotation_scene=dataset_item.annotation_scene,
-                                                 roi=dataset_item.roi,
-                                                 numpy=result_media_numpy)
-                dataset_item.append_metadata_item(result_media, model=self._task_environment.model)
+    def _add_predictions_to_dataset(self, prediction_results, dataset, dump_saliency_map):
+        """ Loop over dataset again to assign predictions. Convert from MMSegmentation format to OTE format. """
+
+        for dataset_item, (prediction, feature_vector, saliency_map) in zip(dataset, prediction_results):
+            soft_prediction = np.transpose(prediction[0], axes=(1, 2, 0))
+            hard_prediction = create_hard_prediction_from_soft_prediction(
+                soft_prediction=soft_prediction,
+                soft_threshold=self._hyperparams.postprocessing.soft_threshold,
+                blur_strength=self._hyperparams.postprocessing.blur_strength,
+            )
+            annotations = create_annotation_from_segmentation_map(
+                hard_prediction=hard_prediction,
+                soft_prediction=soft_prediction,
+                label_map=self._label_dictionary,
+            )
+            dataset_item.append_annotations(annotations=annotations)
+
+            if feature_vector is not None:
+                active_score = TensorEntity(name="representation_vector", numpy=feature_vector)
+                dataset_item.append_metadata_item(active_score, model=self._task_environment.model)
+
+            if dump_saliency_map and saliency_map is not None:
+                saliency_map = get_actmap(saliency_map, (dataset_item.width, dataset_item.height) )
+                saliency_map_media = ResultMediaEntity(name="saliency_map", type="Saliency map",
+                                                annotation_scene=dataset_item.annotation_scene, 
+                                                numpy=saliency_map, roi=dataset_item.roi)
+                dataset_item.append_metadata_item(saliency_map_media, model=self._task_environment.model)
 
     @staticmethod
     def _patch_datasets(config: MPAConfig, domain=Domain.SEGMENTATION):