(release/2.5) Remove duplicate explain() method and consolidate XAI functionality into predict() (#4493)

* Refactor XAI utilities and remove deprecated explain method

Author: eugene123tw

docs/source/guide/explanation/additional_features/xai.rst

@@ -16,7 +16,7 @@ It looks like a heatmap, where warm-colored areas represent the areas with main
   These images are taken from `D-RISE paper <https://arxiv.org/abs/2006.03204>`_.
 
 
-We can generate saliency maps for a certain model that was trained in OpenVINO™ Training Extensions, using ``otx explain`` command line. Learn more about its usage in  :doc:`../../tutorials/base/explain` tutorial.
+We can generate saliency maps for a certain model that was trained in OpenVINO™ Training Extensions, using ``otx predict --explain True`` command line. Learn more about its usage in  :doc:`../../tutorials/base/explain` tutorial.
 
 *********************************
 XAI algorithms for classification
@@ -109,15 +109,22 @@ For instance segmentation networks the following algorithm is used to generate s
 
         .. code-block:: python
 
-            engine.explain(
-              checkpoint="<checkpoint-path>", # .pth or .xml weights of the model
+            engine.predict(
+              checkpoint="checkpoint.pth", # Use .pth when instantiating the engine with OTXEngine
               datamodule=OTXDataModule(), # The data module to use for predictions
-              dump=True # Wherether to save saliency map images or not
+              explain=True # Enable explainability features
               )
+            
+            engine.predict(
+              checkpoint="exported_model.xml", # Use .xml when instantiating the engine with OVEngine
+              datamodule=OTXDataModule(), # The data module to use for predictions
+              explain=True # Enable explainability features
+            )
 
     .. tab-item:: CLI
 
         .. code-block:: bash
 
-            (otx) ...$ otx explain ... --checkpoint <checkpoint-path> # .pth or .xml weights of the model
-                                       --data_root <dataset_path> # Path to data folder or single image 
+            (otx) ...$ otx predict ... --checkpoint <checkpoint-path> # .pth or .xml weights of the model
+                                       --data_root <dataset_path> # Path to data folder or single image
+                                       --explain True # Enable explainability features 

docs/source/guide/get_started/api_tutorial.rst

@@ -12,7 +12,7 @@ To be specific, this tutorial uses as an example of the ATSS model trained throu
 
 For visualization we use images from WGISD dataset from the :doc:`object detection tutorial <how_to_train/detection>` together with trained model.
 
-1. Activate the virtual environment
+1. Activate the virtual environment 
 created in the previous step.
 
 .. code-block:: shell
@@ -21,7 +21,7 @@ created in the previous step.
   # or by this line, if you created an environment, using tox
   . venv/otx/bin/activate
 
-2. ``otx explain`` command returns saliency maps,
+2. ``otx predict`` with the ``--explain True`` parameter returns saliency maps, 
 which are heatmaps with red-colored areas indicating focus. Here's an example how to generate saliency maps from trained checkpoint:
 
 .. tab-set::
@@ -30,31 +30,33 @@ which are heatmaps with red-colored areas indicating focus. Here's an example ho
 
         .. code-block:: shell
 
-            (otx) ...$ otx explain --work_dir otx-workspace \
+            (otx) ...$ otx predict --work_dir otx-workspace \
+                                   --explain True \
                                    --explain_config.postprocess True # Resizes and applies colormap to the saliency map
 
     .. tab-item:: CLI (with config)
 
         .. code-block:: shell
 
-            (otx) ...$ otx explain --config  src/otx/recipe/detection/atss_mobilenetv2.yaml \
+            (otx) ...$ otx predict --config src/otx/recipe/detection/atss_mobilenetv2.yaml \
                                    --data_root data/wgisd \
-                                   --checkpoint otx-workspace/20240312_051135/checkpoints/epoch_033.ckpt \
+                                   --checkpoint otx-workspace/.latest/train/best_checkpoint.ckpt \
+                                   --explain True \
                                    --explain_config.postprocess True # Resizes and applies colormap to the saliency map
 
     .. tab-item:: API
 
         .. code-block:: python
 
-            engine.explain(
-                checkpoint="<checkpoint-path>",
+            engine.predict(
+                checkpoint="otx-workspace/.latest/train/best_checkpoint.ckpt",
                 datamodule=OTXDataModule(...), # The data module to use for predictions
+                explain=True,
                 explain_config=ExplainConfig(postprocess=True), # Resizes and applies colormap to the saliency map
-                dump=True # Wherether to save saliency map images or not
               )
 
-3. The generated saliency maps will appear in  ``otx-workspace/.latest/explain/saliency_maps`` folder.
-It will contain a pair of generated images with saliency maps for each image used for the explanation process:
+3. The generated saliency maps will appear in  ``otx-workspace/.latest/explain/saliency_maps`` folder. 
+It will contain a pair of generated images with saliency maps for each image used for the explanation process: 
 
 - saliency map - where red color means more attention of the model
 - overlay - where the saliency map is combined with the original image:
@@ -64,7 +66,31 @@ It will contain a pair of generated images with saliency maps for each image use
 
 |
 
-4. We can parametrize the explanation process by specifying
+4. To use the exported OpenVINO IR model for explanation, PyTorch weights should be converted to OpenVINO IR model with additional outputs ``saliency_map`` and ``feature_map``.
+To do that we should use ``otx export --explain True`` parameter during export.
+
+.. tab-set::
+
+    .. tab-item:: CLI
+
+        .. code-block:: shell
+
+            (otx) ...$ otx export ... --explain True
+            (otx) ...$ otx predict ... --checkpoint otx-workspace/20240312_052847/exported_model.xml --explain True
+
+    .. tab-item:: API
+
+        .. code-block:: python
+            # Use .pth when instantiating the engine with OTXEngine
+            engine = OTXEngine(model="checkpoint.pth", ...)
+            engine.export(..., explain=True)
+            engine.predict(..., explain=True)
+
+  
+            engine = OVEngine(model="exported_model.xml", ...)
+            engine.predict(..., explain=True)
+
+5. We can parametrize the explanation process by specifying 
 the following parameters in ``ExplainConfig``:
 
 - ``target_explain_group`` - for which target saliency maps will be generated:
@@ -84,21 +110,23 @@ the following parameters in ``ExplainConfig``:
 
         .. code-block:: shell
 
-            (otx) ...$ otx explain ... --explain_config.postprocess True
+            (otx) ...$ otx predict ... --explain True \
+                                       --explain_config.postprocess True \
                                        --explain_config.target_explain_group PREDICTIONS
 
     .. tab-item:: API
 
         .. code-block:: python
 
-            engine.explain(...,
+            engine.predict(...,
+                           explain=True,
                            explain_config=ExplainConfig(
                              postprocess=True,
                              target_explain_group=TargetExplainGroup.PREDICTIONS
                            )
               )
 
-5. The explanation algorithm is chosen automatically
+6. The explanation algorithm is chosen automatically 
 based on the used model:
 
 - ``Recipro-CAM`` - for CNN classification models

docs/source/guide/get_started/cli_commands.rst

@@ -421,10 +421,7 @@ def predict(
                 ...     --checkpoint <CKPT_PATH, str>
                 ```
         """
-        from otx.backend.native.models.utils.xai_utils import (
-            process_saliency_maps_in_pred_entity,
-            set_crop_padded_map_flag,
-        )
+        from otx.backend.native.models.utils.xai_utils import process_saliency_maps_in_pred_entity
 
         model = self.model
 
@@ -462,7 +459,6 @@ def predict(
         if explain:
             if explain_config is None:
                 explain_config = ExplainConfig()
-            explain_config = set_crop_padded_map_flag(explain_config, datamodule)
 
             predict_result = process_saliency_maps_in_pred_entity(predict_result, explain_config, datamodule.label_info)
 
@@ -548,91 +544,6 @@ def export(
         self.model.explain_mode = False
         return exported_model_path
 
-    def explain(
-        self,
-        checkpoint: PathLike | None = None,
-        datamodule: EVAL_DATALOADERS | OTXDataModule | None = None,
-        explain_config: ExplainConfig | None = None,
-        **kwargs,
-    ) -> list | None:
-        r"""Run XAI using the specified model and data (test subset).
-
-        Args:
-            checkpoint (PathLike | None, optional): The path to the checkpoint file to load the model from.
-            datamodule (EVAL_DATALOADERS | OTXDataModule | None, optional): The data module to use for predictions.
-            explain_config (ExplainConfig | None, optional): Config used to handle saliency maps.
-            **kwargs: Additional keyword arguments for pl.Trainer configuration.
-
-        Returns:
-            list: Saliency maps.
-
-        Example:
-            >>> engine.explain(
-            ...     datamodule=OTXDataModule(),
-            ...     checkpoint=<checkpoint/path>,
-            ...     explain_config=ExplainConfig(),
-            ... )
-
-        CLI Usage:
-            1. To run XAI with the torch model in work_dir, run
-                ```shell
-                >>> otx explain \
-                ...     --work_dir <WORK_DIR_PATH, str>
-                ```
-            2. To run XAI using the specified model (torch or IR), run
-                ```shell
-                >>> otx explain \
-                ...     --work_dir <WORK_DIR_PATH, str> \
-                ...     --checkpoint <CKPT_PATH, str>
-                ```
-            3. To run XAI using the configuration, run
-                ```shell
-                >>> otx explain \
-                ...     --config <CONFIG_PATH> --data_root <DATASET_PATH, str> \
-                ...     --checkpoint <CKPT_PATH, str>
-                ```
-        """
-        from otx.backend.native.models.utils.xai_utils import (
-            process_saliency_maps_in_pred_entity,
-            set_crop_padded_map_flag,
-        )
-
-        model = self.model
-
-        checkpoint = checkpoint if checkpoint is not None else self.checkpoint
-        datamodule = datamodule if datamodule is not None else self.datamodule
-
-        if checkpoint is not None:
-            ckpt = self._load_model_checkpoint(checkpoint, map_location="cpu")
-            model.load_state_dict(ckpt)
-
-        if model.label_info != self.datamodule.label_info:
-            msg = (
-                "To launch a explain pipeline, the label information should be same "
-                "between the training and testing datasets. "
-                "Please check whether you use the same dataset: "
-                f"model.label_info={model.label_info}, "
-                f"datamodule.label_info={self.datamodule.label_info}"
-            )
-            raise ValueError(msg)
-
-        model.explain_mode = True
-
-        self._build_trainer(**kwargs)
-
-        predict_result = self.trainer.predict(
-            model=model,
-            datamodule=datamodule,
-        )
-
-        if explain_config is None:
-            explain_config = ExplainConfig()
-        explain_config = set_crop_padded_map_flag(explain_config, datamodule)
-
-        predict_result = process_saliency_maps_in_pred_entity(predict_result, explain_config, datamodule.label_info)
-        model.explain_mode = False
-        return predict_result
-
     def benchmark(
         self,
         checkpoint: PathLike | None = None,

docs/source/guide/tutorials/base/explain.rst

@@ -211,7 +211,7 @@ def _customize_outputs(
                 scores=scores,
                 bboxes=bboxes,
                 labels=labels,
-                saliency_map=[saliency_map.detach().to(torch.float32) for saliency_map in outputs["saliency_map"]],
+                saliency_map=outputs["saliency_map"],
                 feature_vector=[
                     feature_vector.detach().unsqueeze(0).to(torch.float32)
                     for feature_vector in outputs["feature_vector"]

src/otx/backend/native/engine.py

@@ -26,7 +26,6 @@
 if TYPE_CHECKING:
     from torch import LongTensor, Tensor
 
-    from otx.data.module import OTXDataModule
 
 ProcessedSaliencyMaps = list[dict[str, np.ndarray | torch.Tensor]]
 
@@ -37,13 +36,10 @@ def process_saliency_maps_in_pred_entity(
     label_info: LabelInfoTypes,
 ) -> list[OTXPredBatch]:
     """Process saliency maps in PredEntity."""
-
-    def _process(
-        predict_result_per_batch: OTXPredBatch,
-        label_info: LabelInfoTypes,
-    ) -> OTXPredBatch:
-        if predict_result_per_batch.saliency_map is None:  # skip empty saliency maps
-            return predict_result_per_batch
+    processed_predict_result = []
+    for predict_result_per_batch in predict_result:
+        if predict_result_per_batch.saliency_map is None or len(predict_result_per_batch.saliency_map) == 0:
+            continue
 
         # Extract batch data with proper type handling
         labels = predict_result_per_batch.labels if predict_result_per_batch.labels is not None else []
@@ -60,6 +56,7 @@ def _process(
         # Add additional conf threshold for saving maps with predicted classes,
         # since predictions can have less than 0.05 confidence
         conf_thr = explain_config.predicted_maps_conf_thr
+        keep_ratio = imgs_info[0].keep_ratio  # type: ignore[union-attr, index]
 
         pred_labels = []
         for labels, scores in zip(predict_result_per_batch.labels, predict_result_per_batch.scores):  # type: ignore[union-attr, arg-type]
@@ -81,11 +78,12 @@ def _process(
             ori_img_shapes,
             image_shape,
             paddings,
+            keep_ratio,
         )
         predict_result_per_batch.saliency_map = processed_saliency_maps
-        return predict_result_per_batch
+        processed_predict_result.append(predict_result_per_batch)
 
-    return [_process(predict_result_per_batch, label_info) for predict_result_per_batch in predict_result]
+    return processed_predict_result
 
 
 def process_saliency_maps(
@@ -95,6 +93,7 @@ def process_saliency_maps(
     ori_img_shapes: list,
     image_shape: tuple[int, int],
     paddings: list[tuple[int, int, int, int]],
+    keep_ratio: bool,
 ) -> ProcessedSaliencyMaps:
     """Perform saliency map convertion to dict and post-processing."""
     if explain_config.target_explain_group == TargetExplainGroup.ALL:
@@ -107,7 +106,7 @@ def process_saliency_maps(
         msg = f"Target explain group {explain_config.target_explain_group} is not supported."
         raise ValueError(msg)
 
-    if explain_config.crop_padded_map:
+    if keep_ratio:
         processed_saliency_maps = _crop_padded_map(processed_saliency_maps, image_shape, paddings)
 
     if explain_config.postprocess:
@@ -221,12 +220,3 @@ def _convert_labels_from_hcls_format(
                 pred_labels.append(label_info.label_to_idx[label_str])
 
     return pred_labels
-
-
-def set_crop_padded_map_flag(explain_config: ExplainConfig, datamodule: OTXDataModule) -> ExplainConfig:
-    """If resize with keep_ratio = True was used, set crop_padded_map flag to True."""
-    for transform in datamodule.test_subset.transforms:
-        tranf_name = transform["class_path"].split(".")[-1]
-        if tranf_name == "Resize" and transform["init_args"].get("keep_ratio", False):
-            explain_config.crop_padded_map = True
-    return explain_config

src/otx/backend/native/models/detection/base.py

@@ -1,10 +1,11 @@
-# Copyright (C) 2024 Intel Corporation
+# Copyright (C) 2024-2025 Intel Corporation
 # SPDX-License-Identifier: Apache-2.0
 
 """Algorithms for calculcalating XAI branch for Explainable AI."""
 
 from __future__ import annotations
 
+import warnings
 from typing import TYPE_CHECKING, Callable
 
 import torch
@@ -267,6 +268,14 @@ def __init__(
         self._num_classes = num_classes
         self._num_anchors = num_anchors
         # Should be switched off for tiling
+        if num_classes == 1 and use_cls_softmax:
+            # softmax would result in all 1.0 values if there's only 1 class
+            warnings.warn(
+                "use_cls_softmax is automatically disabled when num_classes=1 to prevent degenerate softmax behavior",
+                UserWarning,
+                stacklevel=2,
+            )
+            use_cls_softmax = False
         self.use_cls_softmax = use_cls_softmax
 
     def func(