[ENHANCE] Parametrize saliency maps dumping in export parametrizable (#…

…1708)

* Add a `dump_features` parameter for export

* Fix pre-commit & add parameter to IExport

* Add dump_features=True to other tasks' export

* Add comments

* Fix pre-commit

* Add NotImplementedError

* Fix CLI tests

* Solve merge conflict

* Fix pre-commit

* Fix pre-commit

otx/algorithms/action/tasks/inference.py

@@ -307,8 +307,15 @@ def unload(self):
             self._delete_scratch_space()
 
     @check_input_parameters_type()
-    def export(self, export_type: ExportType, output_model: ModelEntity):
+    def export(self, export_type: ExportType, output_model: ModelEntity, dump_features: bool = True):
         """Export function of OTX Action Task."""
+        # TODO: add dumping saliency maps and representation vectors according to dump_features flag
+        if not dump_features:
+            logger.warning(
+                "Ommitting feature dumping is not implemented."
+                "The saliency maps and representation vector outputs will be dumped in the exported model."
+            )
+
         # copied from OTX inference_task.py
         logger.info("Exporting the model")
         if export_type != ExportType.OPENVINO:

otx/algorithms/anomaly/tasks/inference.py

@@ -247,16 +247,24 @@ def _export_to_onnx(self, onnx_path: str):
             opset_version=11,
         )
 
-    def export(self, export_type: ExportType, output_model: ModelEntity) -> None:
+    def export(self, export_type: ExportType, output_model: ModelEntity, dump_features: bool = True) -> None:
         """Export model to OpenVINO IR.
 
         Args:
             export_type (ExportType): Export type should be ExportType.OPENVINO
             output_model (ModelEntity): The model entity in which to write the OpenVINO IR data
+            dump_features (bool): Flag to return "feature_vector" and "saliency_map".
 
         Raises:
             Exception: If export_type is not ExportType.OPENVINO
         """
+        # TODO: add dumping saliency maps and representation vectors according to dump_features flag
+        if not dump_features:
+            logger.warning(
+                "Ommitting feature dumping is not implemented."
+                "The saliency maps and representation vector outputs will be dumped in the exported model."
+            )
+
         assert export_type == ExportType.OPENVINO, f"Incorrect export_type={export_type}"
         output_model.model_format = ModelFormat.OPENVINO
         output_model.optimization_type = ModelOptimizationType.MO

otx/algorithms/classification/adapters/openvino/model_wrappers/openvino_models.py

@@ -103,8 +103,6 @@ def postprocess(self, outputs: Dict[str, np.ndarray], meta: Dict[str, Any]):  #
     @check_input_parameters_type()
     def postprocess_aux_outputs(self, outputs: Dict[str, np.ndarray], metadata: Dict[str, Any]):
         """Post-process for auxiliary outputs."""
-        saliency_map = outputs["saliency_map"][0]
-        repr_vector = outputs["feature_vector"].reshape(-1)
         logits = outputs[self.out_layer_name].squeeze()
         if self.multilabel:
             probs = sigmoid_numpy(logits)
@@ -113,6 +111,13 @@ def postprocess_aux_outputs(self, outputs: Dict[str, np.ndarray], metadata: Dict
         else:
             probs = softmax_numpy(logits)
         act_score = float(np.max(probs) - np.min(probs))
+
+        if "saliency_map" in outputs:
+            saliency_map = outputs["saliency_map"][0]
+            repr_vector = outputs["feature_vector"].reshape(-1)
+        else:
+            saliency_map, repr_vector = None, None
+
         return probs, saliency_map, repr_vector, act_score
 
 

otx/algorithms/classification/configs/efficientnet_b0_cls_incr/deployment.py

@@ -3,7 +3,7 @@
 _base_ = ["../base/deployments/base_classification_dynamic.py"]
 
 ir_config = dict(
-    output_names=["logits", "feature_vector", "saliency_map"],
+    output_names=["logits"],
 )
 
 backend_config = dict(

otx/algorithms/classification/configs/efficientnet_v2_s_cls_incr/deployment.py

@@ -3,7 +3,7 @@
 _base_ = ["../base/deployments/base_classification_dynamic.py"]
 
 ir_config = dict(
-    output_names=["logits", "feature_vector", "saliency_map"],
+    output_names=["logits"],
 )
 
 backend_config = dict(

otx/algorithms/classification/configs/mobilenet_v3_large_1_cls_incr/deployment.py

@@ -3,7 +3,7 @@
 _base_ = ["../base/deployments/base_classification_dynamic.py"]
 
 ir_config = dict(
-    output_names=["logits", "feature_vector", "saliency_map"],
+    output_names=["logits"],
 )
 
 backend_config = dict(

otx/algorithms/classification/tasks/inference.py

@@ -215,7 +215,7 @@ def unload(self):
         self.cleanup()
 
     @check_input_parameters_type()
-    def export(self, export_type: ExportType, output_model: ModelEntity):
+    def export(self, export_type: ExportType, output_model: ModelEntity, dump_features: bool = False):
         """Export function of OTX Classification Task."""
 
         logger.info("Exporting the model")
@@ -225,7 +225,7 @@ def export(self, export_type: ExportType, output_model: ModelEntity):
         output_model.optimization_type = ModelOptimizationType.MO
 
         stage_module = "ClsExporter"
-        results = self._run_task(stage_module, mode="train", export=True)
+        results = self._run_task(stage_module, mode="train", export=True, dump_features=dump_features)
         outputs = results.get("outputs")
         logger.debug(f"results of run_task = {outputs}")
         if outputs is None:

otx/algorithms/classification/tasks/openvino.py

@@ -19,6 +19,7 @@
 import logging
 import os
 import tempfile
+import warnings
 from typing import Any, Dict, Optional, Tuple, Union
 from zipfile import ZipFile
 
@@ -78,8 +79,6 @@
     from openvino.model_zoo.model_api.adapters import OpenvinoAdapter, create_core
     from openvino.model_zoo.model_api.models import Model
 except ImportError:
-    import warnings
-
     warnings.warn("ModelAPI was not found.")
 
 logger = logging.getLogger(__name__)
@@ -239,18 +238,25 @@ def infer(
             probs_meta = TensorEntity(name="probabilities", numpy=probs.reshape(-1))
             dataset_item.append_metadata_item(probs_meta, model=self.model)
 
-            feature_vec_media = TensorEntity(name="representation_vector", numpy=repr_vector.reshape(-1))
-            dataset_item.append_metadata_item(feature_vec_media, model=self.model)
             if dump_features:
-                add_saliency_maps_to_dataset_item(
-                    dataset_item=dataset_item,
-                    saliency_map=saliency_map,
-                    model=self.model,
-                    labels=self.task_environment.get_labels(),
-                    predicted_scored_labels=item_labels,
-                    explain_predicted_classes=explain_predicted_classes,
-                    process_saliency_maps=process_saliency_maps,
-                )
+                if saliency_map is not None and repr_vector is not None:
+                    feature_vec_media = TensorEntity(name="representation_vector", numpy=repr_vector.reshape(-1))
+                    dataset_item.append_metadata_item(feature_vec_media, model=self.model)
+
+                    add_saliency_maps_to_dataset_item(
+                        dataset_item=dataset_item,
+                        saliency_map=saliency_map,
+                        model=self.model,
+                        labels=self.task_environment.get_labels(),
+                        predicted_scored_labels=item_labels,
+                        explain_predicted_classes=explain_predicted_classes,
+                        process_saliency_maps=process_saliency_maps,
+                    )
+                else:
+                    warnings.warn(
+                        "Could not find Feature Vector and Saliency Map in OpenVINO output. "
+                        "Please rerun OpenVINO export or retrain the model."
+                    )
             update_progress_callback(int(i / dataset_size * 100))
         return dataset
 

otx/algorithms/common/tasks/training_base.py

@@ -332,6 +332,12 @@ def _initialize(self, options=None):  # noqa: C901
                 assert len(self._precision) == 1
                 options["precision"] = str(self._precision[0])
 
+            options["deploy_cfg"]["dump_features"] = options["dump_features"]
+            if options["dump_features"]:
+                output_names = options["deploy_cfg"]["ir_config"]["output_names"]
+                if "feature_vector" not in output_names and "saliency_map" not in output_names:
+                    options["deploy_cfg"]["ir_config"]["output_names"] += ["feature_vector", "saliency_map"]
+
         self._initialize_post_hook(options)
 
         logger.info("initialized.")

otx/algorithms/detection/configs/detection/cspdarknet_yolox/deployment.py

@@ -3,7 +3,7 @@
 _base_ = ["../../base/deployments/base_detection_dynamic.py"]
 
 ir_config = dict(
-    output_names=["boxes", "labels", "feature_vector", "saliency_map"],
+    output_names=["boxes", "labels"],
 )
 
 backend_config = dict(

otx/algorithms/detection/configs/detection/mobilenetv2_atss/deployment.py

@@ -3,7 +3,7 @@
 _base_ = ["../../base/deployments/base_detection_dynamic.py"]
 
 ir_config = dict(
-    output_names=["boxes", "labels", "feature_vector", "saliency_map"],
+    output_names=["boxes", "labels"],
 )
 
 backend_config = dict(

otx/algorithms/detection/configs/detection/mobilenetv2_ssd/deployment.py

@@ -3,7 +3,7 @@
 _base_ = ["../../base/deployments/base_detection_dynamic.py"]
 
 ir_config = dict(
-    output_names=["boxes", "labels", "feature_vector", "saliency_map"],
+    output_names=["boxes", "labels"],
 )
 
 backend_config = dict(

otx/algorithms/detection/tasks/inference.py

@@ -238,7 +238,7 @@ def unload(self):
         self.cleanup()
 
     @check_input_parameters_type()
-    def export(self, export_type: ExportType, output_model: ModelEntity):
+    def export(self, export_type: ExportType, output_model: ModelEntity, dump_features: bool = False):
         """Export function of OTX Detection Task."""
         # copied from OTX inference_task.py
         logger.info("Exporting the model")
@@ -252,6 +252,7 @@ def export(self, export_type: ExportType, output_model: ModelEntity):
             stage_module,
             mode="train",
             export=True,
+            dump_features=dump_features,
         )
         outputs = results.get("outputs")
         logger.debug(f"results of run_task = {outputs}")

otx/algorithms/segmentation/tasks/inference.py

@@ -24,6 +24,7 @@
     patch_data_pipeline,
     patch_default_config,
     patch_runner,
+    remove_from_configs_by_type,
 )
 from otx.algorithms.common.configs import TrainType
 from otx.algorithms.common.tasks import BaseTask
@@ -87,6 +88,7 @@ class SegmentationInferenceTask(BaseTask, IInferenceTask, IExportTask, IEvaluati
     @check_input_parameters_type()
     def __init__(self, task_environment: TaskEnvironment, **kwargs):
         # self._should_stop = False
+        self.freeze = True
         self.metric = "mDice"
         self._label_dictionary = {}  # type: Dict
 
@@ -145,19 +147,26 @@ def unload(self):
         self.cleanup()
 
     @check_input_parameters_type()
-    def export(self, export_type: ExportType, output_model: ModelEntity):
+    def export(self, export_type: ExportType, output_model: ModelEntity, dump_features: bool = True):
         """Export function of OTX Segmentation Task."""
         logger.info("Exporting the model")
         if export_type != ExportType.OPENVINO:
             raise RuntimeError(f"not supported export type {export_type}")
         output_model.model_format = ModelFormat.OPENVINO
         output_model.optimization_type = ModelOptimizationType.MO
+        # TODO: add dumping saliency maps and representation vectors according to dump_features flag
+        if not dump_features:
+            logger.warning(
+                "Ommitting feature dumping is not implemented."
+                "The saliency maps and representation vector outputs will be dumped in the exported model."
+            )
 
         stage_module = "SegExporter"
         results = self._run_task(
             stage_module,
             mode="train",
             export=True,
+            dump_features=dump_features,
         )
         outputs = results.get("outputs")
         logger.debug(f"results of run_task = {outputs}")
@@ -212,6 +221,8 @@ def _init_recipe(self):
         if self._recipe_cfg.get("override_configs", None):
             self.override_configs.update(self._recipe_cfg.override_configs)
 
+        if not self.freeze:
+            remove_from_configs_by_type(self._recipe_cfg.custom_hooks, "FreezeLayers")
         logger.info(f"initialized recipe = {recipe}")
 
     def _update_stage_module(self, stage_module: str):

otx/api/usecases/tasks/interfaces/export_interface.py

@@ -20,11 +20,12 @@ class IExportTask(metaclass=abc.ABCMeta):
     """A base interface class for tasks which can export their models."""
 
     @abc.abstractmethod
-    def export(self, export_type: ExportType, output_model: ModelEntity):
+    def export(self, export_type: ExportType, output_model: ModelEntity, dump_features: bool = False):
         """This method defines the interface for export.
 
         Args:
             export_type (ExportType): The type of optimization.
             output_model (ModelEntity): The output model entity.
+            dump_features (bool): Flag to return "feature_vector" and "saliency_map".
         """
         raise NotImplementedError

otx/cli/tools/export.py

@@ -40,6 +40,11 @@ def get_args():
         "--save-model-to",
         help="Location where exported model will be stored.",
     )
+    parser.add_argument(
+        "--dump_features",
+        action="store_true",
+        help="Whether to return feature vector and saliency map for explanation purposes.",
+    )
 
     return parser.parse_args()
 
@@ -84,7 +89,7 @@ def main():
 
     exported_model = ModelEntity(None, environment.get_model_configuration())
 
-    task.export(ExportType.OPENVINO, exported_model)
+    task.export(ExportType.OPENVINO, exported_model, args.dump_features)
 
     if "save_model_to" not in args or not args.save_model_to:
         args.save_model_to = str(config_manager.workspace_root / "model-exported")

otx/mpa/modules/models/classifiers/sam_classifier.py

@@ -295,6 +295,10 @@ def sam_image_classifier__extract_feat(ctx, self, img):
     def sam_image_classifier__simple_test(ctx, self, img, img_metas):
         feat, backbone_feat = self.extract_feat(img)
         logit = self.head.simple_test(feat)
-        saliency_map = ReciproCAMHook(self).func(backbone_feat)
-        feature_vector = FeatureVectorHook.func(backbone_feat)
-        return logit, feature_vector, saliency_map
+
+        if ctx.cfg["dump_features"]:
+            saliency_map = ReciproCAMHook(self).func(backbone_feat)
+            feature_vector = FeatureVectorHook.func(backbone_feat)
+            return logit, feature_vector, saliency_map
+
+        return logit

otx/mpa/modules/models/detectors/custom_atss_detector.py

@@ -86,10 +86,14 @@ def custom_atss__simple_test(ctx, self, img, img_metas, **kwargs):
         feat = self.extract_feat(img)
         outs = self.bbox_head(feat)
         bbox_results = self.bbox_head.get_bboxes(*outs, img_metas=img_metas, cfg=self.test_cfg, **kwargs)
-        feature_vector = FeatureVectorHook.func(feat)
-        cls_scores = outs[0]
-        saliency_map = DetSaliencyMapHook(self).func(cls_scores, cls_scores_provided=True)
-        return (*bbox_results, feature_vector, saliency_map)
+
+        if ctx.cfg["dump_features"]:
+            feature_vector = FeatureVectorHook.func(feat)
+            cls_scores = outs[0]
+            saliency_map = DetSaliencyMapHook(self).func(cls_scores, cls_scores_provided=True)
+            return (*bbox_results, feature_vector, saliency_map)
+
+        return bbox_results
 
     @mark("custom_atss_forward", inputs=["input"], outputs=["dets", "labels", "feats", "saliencies"])
     def __forward_impl(ctx, self, img, img_metas, **kwargs):

otx/mpa/modules/models/detectors/custom_maskrcnn_detector.py

@@ -91,12 +91,16 @@ def load_state_dict_pre_hook(model, model_classes, chkpt_classes, chkpt_dict, pr
     def custom_mask_rcnn__simple_test(ctx, self, img, img_metas, proposals=None, **kwargs):
         assert self.with_bbox, "Bbox head must be implemented."
         x = self.extract_feat(img)
-        feature_vector = FeatureVectorHook.func(x)
-        saliency_map = ActivationMapHook.func(x[-1])
         if proposals is None:
             proposals, _ = self.rpn_head.simple_test_rpn(x, img_metas)
         out = self.roi_head.simple_test(x, proposals, img_metas, rescale=False)
-        return (*out, feature_vector, saliency_map)
+
+        if ctx.cfg["dump_features"]:
+            feature_vector = FeatureVectorHook.func(x)
+            saliency_map = ActivationMapHook.func(x[-1])
+            return (*out, feature_vector, saliency_map)
+
+        return out
 
     @mark("custom_maskrcnn_forward", inputs=["input"], outputs=["dets", "labels", "masks", "feats", "saliencies"])
     def __forward_impl(ctx, self, img, img_metas, **kwargs):

otx/mpa/modules/models/detectors/custom_single_stage_detector.py

@@ -142,10 +142,14 @@ def custom_single_stage_detector__simple_test(ctx, self, img, img_metas, **kwarg
         feat = self.extract_feat(img)
         outs = self.bbox_head(feat)
         bbox_results = self.bbox_head.get_bboxes(*outs, img_metas=img_metas, cfg=self.test_cfg, **kwargs)
-        feature_vector = FeatureVectorHook.func(feat)
-        cls_scores = outs[0]
-        saliency_map = DetSaliencyMapHook(self).func(cls_scores, cls_scores_provided=True)
-        return (*bbox_results, feature_vector, saliency_map)
+
+        if ctx.cfg["dump_features"]:
+            feature_vector = FeatureVectorHook.func(feat)
+            cls_scores = outs[0]
+            saliency_map = DetSaliencyMapHook(self).func(cls_scores, cls_scores_provided=True)
+            return (*bbox_results, feature_vector, saliency_map)
+
+        return bbox_results
 
     @mark("custom_ssd_forward", inputs=["input"], outputs=["dets", "labels", "feats", "saliencies"])
     def __forward_impl(ctx, self, img, img_metas, **kwargs):