Update MRCNN model export to include feature vector and saliency map (#4056)

Author: harimkang

CHANGELOG.md

@@ -90,6 +90,16 @@ All notable changes to this project will be documented in this file.
   (<https://github.com/openvinotoolkit/training_extensions/pull/4042>)
 - Fix wrong indices setting in HLabelInfo
   (<https://github.com/openvinotoolkit/training_extensions/pull/4044>)
+- Add legacy template LiteHRNet_18 template
+  (<https://github.com/openvinotoolkit/training_extensions/pull/4049>)
+- Model templates: rename model_status value 'DISCONTINUED' to 'OBSOLETE'
+  (<https://github.com/openvinotoolkit/training_extensions/pull/4051>)
+- Enable export of feature vectors for semantic segmentation task
+  (<https://github.com/openvinotoolkit/training_extensions/pull/4055>)
+- Update MRCNN model export to include feature vector and saliency map
+  (<https://github.com/openvinotoolkit/training_extensions/pull/4056>)
+- Upgrade MAPI in 2.2
+  (<https://github.com/openvinotoolkit/training_extensions/pull/4052>)
 
 ## \[v2.1.0\]
 

src/otx/__init__.py

@@ -3,7 +3,7 @@
 # Copyright (C) 2024 Intel Corporation
 # SPDX-License-Identifier: Apache-2.0
 
-__version__ = "2.2.0rc9"
+__version__ = "2.2.0rc10"
 
 import os
 from pathlib import Path

src/otx/algo/instance_segmentation/maskrcnn.py

@@ -70,7 +70,7 @@ def _exporter(self) -> OTXModelExporter:
                 "opset_version": 11,
                 "autograd_inlining": False,
             },
-            output_names=["bboxes", "labels", "masks"],
+            output_names=["bboxes", "labels", "masks", "feature_vector", "saliency_map"] if self.explain_mode else None,
         )
 
     def load_from_otx_v1_ckpt(self, state_dict: dict, add_prefix: str = "model.") -> dict:

src/otx/algo/instance_segmentation/rtmdet_inst.py

@@ -68,6 +68,7 @@ def _exporter(self) -> OTXModelExporter:
                 "opset_version": 11,
                 "autograd_inlining": False,
             },
+            # TODO(Eugene): Add XAI support for RTMDetInst
             output_names=["bboxes", "labels", "masks", "feature_vector", "saliency_map"] if self.explain_mode else None,
         )
 

src/otx/algo/instance_segmentation/two_stage.py

@@ -210,19 +210,44 @@ def export(
         self,
         batch_inputs: torch.Tensor,
         batch_img_metas: list[dict],
-    ) -> tuple[torch.Tensor, ...]:
-        """Export for two stage detectors."""
-        x = self.extract_feat(batch_inputs)
+        explain_mode: bool,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor] | dict:
+        """Export the model for ONNX/OpenVINO.
+
+        Args:
+            batch_inputs (torch.Tensor): image tensor with shape (N, C, H, W).
+            batch_img_metas (list[dict]): image information.
+            explain_mode (bool): whether to return feature vector.
 
+        Returns:
+            tuple[torch.Tensor, torch.Tensor, torch.Tensor] | dict:
+                - bboxes (torch.Tensor): bounding boxes.
+                - labels (torch.Tensor): labels.
+                - masks (torch.Tensor): masks.
+                - feature_vector (torch.Tensor, optional): feature vector.
+                - saliency_map (torch.Tensor, optional): saliency map.
+        """
+        x = self.extract_feat(batch_inputs)
         rpn_results_list = self.rpn_head.export(
             x,
             batch_img_metas,
             rescale=False,
         )
-
-        return self.roi_head.export(
+        bboxes, labels, masks = self.roi_head.export(
             x,
             rpn_results_list,
             batch_img_metas,
             rescale=False,
         )
+
+        if explain_mode:
+            feature_vector = self.feature_vector_fn(x)
+            return {
+                "bboxes": bboxes,
+                "labels": labels,
+                "masks": masks,
+                "feature_vector": feature_vector,
+                # create dummy tensor as model API supports saliency_map
+                "saliency_map": torch.zeros(1),
+            }
+        return bboxes, labels, masks

src/otx/core/model/instance_segmentation.py

@@ -252,7 +252,7 @@ def forward_for_tracing(self, inputs: Tensor) -> tuple[Tensor, ...]:
             "scale_factor": (1.0, 1.0),
         }
         meta_info_list = [meta_info] * len(inputs)
-        return self.model.export(inputs, meta_info_list)
+        return self.model.export(inputs, meta_info_list, explain_mode=self.explain_mode)
 
     @property
     def _export_parameters(self) -> TaskLevelExportParameters:

tests/integration/api/test_xai.py

@@ -110,9 +110,9 @@ def test_predict_with_explain(
     if "dino" in model_name:
         pytest.skip("DINO is not supported.")
 
-    if "instance_segmentation" in recipe:
-        # TODO(Eugene): figure out why instance segmentation model fails after decoupling.
-        pytest.skip("There's issue with instance segmentation model. Skip for now.")
+    if any(keyword in recipe for keyword in ["rtmdet_inst_tiny", "maskdino", "maskrcnn_r50_tv"]):
+        # TODO(Eugene): inst-seg models not fully support yet.
+        pytest.skip(f"There's issue with inst-seg: {recipe}. Skip for now.")
 
     if "rtmdet_tiny" in recipe:
         # TODO (sungchul): enable xai for rtmdet_tiny (CVS-142651)