Refactor OpenVINO imports

openvino_xai/api/api.py

@@ -3,7 +3,7 @@
 
 from typing import List
 
-import openvino.runtime as ov
+import openvino as ov
 
 from openvino_xai.common.parameters import Method, Task
 from openvino_xai.common.utils import IdentityPreprocessFN, has_xai, logger

openvino_xai/explainer/explainer.py

@@ -2,11 +2,10 @@
 # SPDX-License-Identifier: Apache-2.0
 
 from enum import Enum
-from typing import Callable, List, Tuple
+from typing import Callable, List, Mapping, Tuple
 
 import numpy as np
-import openvino.runtime as ov
-from openvino.runtime.utils.data_helpers.wrappers import OVDict
+import openvino as ov
 
 from openvino_xai import Task
 from openvino_xai.common.parameters import Method
@@ -53,7 +52,7 @@ class Explainer:
         (assume input images are already preprocessed by user).
     :type preprocess_fn: Callable[[np.ndarray], np.ndarray]
     :param postprocess_fn: Postprocessing functions, required for black-box.
-    :type postprocess_fn: Callable[[OVDict], np.ndarray]
+    :type postprocess_fn: Callable[[Mapping], np.ndarray]
     :param explain_mode: Explain mode.
     :type explain_mode: ExplainMode
     :parameter explain_method: Explain method to use for model explanation.
@@ -71,7 +70,7 @@ def __init__(
         model: ov.Model,
         task: Task,
         preprocess_fn: Callable[[np.ndarray], np.ndarray] = IdentityPreprocessFN(),
-        postprocess_fn: Callable[[OVDict], np.ndarray] = None,
+        postprocess_fn: Callable[[Mapping], np.ndarray] = None,
         explain_mode: ExplainMode = ExplainMode.AUTO,
         explain_method: Method | None = None,
         target_layer: str | List[str] | None = None,
@@ -80,7 +79,7 @@ def __init__(
         **kwargs,
     ) -> None:
         self.model = model
-        self.compiled_model: ov.ie_api.CompiledModel | None = None
+        self.compiled_model: ov.CompiledModel | None = None
         self.task = task
 
         if isinstance(preprocess_fn, IdentityPreprocessFN):
@@ -228,7 +227,7 @@ def explain(
             overlay_weight,
         )
 
-    def model_forward(self, x: np.ndarray, preprocess: bool = True) -> OVDict:
+    def model_forward(self, x: np.ndarray, preprocess: bool = True) -> Mapping:
         """Forward pass of the compiled model."""
         return self.method.model_forward(x, preprocess)
 

openvino_xai/explainer/utils.py

@@ -2,11 +2,10 @@
 # SPDX-License-Identifier: Apache-2.0
 from enum import Enum
 from functools import partial
-from typing import Any, Callable, List, Tuple
+from typing import Any, Callable, List, Mapping, Tuple
 
 import cv2
 import numpy as np
-from openvino.runtime.utils.data_helpers.wrappers import OVDict
 
 
 def convert_targets_to_numpy(targets):
@@ -118,7 +117,7 @@ def get_preprocess_fn(
     )
 
 
-def postprocess_fn(x: OVDict, logit_name="logits") -> np.ndarray:
+def postprocess_fn(x: Mapping, logit_name="logits") -> np.ndarray:
     """Postprocess function."""
     return x.get(logit_name, x[0])  # Models from OVC has no output names at times
 

openvino_xai/methods/base.py

@@ -2,11 +2,10 @@
 # SPDX-License-Identifier: Apache-2.0
 
 from abc import ABC, abstractmethod
-from typing import Callable
+from typing import Callable, Mapping
 
 import numpy as np
 import openvino.runtime as ov
-from openvino.runtime.utils.data_helpers.wrappers import OVDict
 
 from openvino_xai.common.utils import IdentityPreprocessFN
 
@@ -33,7 +32,7 @@ def model_compiled(self) -> ov.ie_api.CompiledModel | None:
     def prepare_model(self, load_model: bool = True) -> ov.Model:
         """Model preparation steps."""
 
-    def model_forward(self, x: np.ndarray, preprocess: bool = True) -> OVDict:
+    def model_forward(self, x: np.ndarray, preprocess: bool = True) -> Mapping:
         """Forward pass of the compiled model. Applies preprocess_fn."""
         if not self._model_compiled:
             raise RuntimeError("Model is not compiled. Call prepare_model() first.")

openvino_xai/methods/black_box/rise.py

@@ -1,12 +1,11 @@
 # Copyright (C) 2023-2024 Intel Corporation
 # SPDX-License-Identifier: Apache-2.0
 
-from typing import Callable, List, Tuple
+from typing import Callable, List, Mapping, Tuple
 
 import cv2
 import numpy as np
 import openvino.runtime as ov
-from openvino.runtime.utils.data_helpers.wrappers import OVDict
 from tqdm import tqdm
 
 from openvino_xai.common.utils import IdentityPreprocessFN, scaling
@@ -19,7 +18,7 @@ class RISE(BlackBoxXAIMethod):
     :param model: OpenVINO model.
     :type model: ov.Model
     :param postprocess_fn: Preprocessing function that extract scores from IR model output.
-    :type postprocess_fn: Callable[[OVDict], np.ndarray]
+    :type postprocess_fn: Callable[[Mapping], np.ndarray]
     :param preprocess_fn: Preprocessing function, identity function by default
         (assume input images are already preprocessed by user).
     :type preprocess_fn: Callable[[np.ndarray], np.ndarray]
@@ -32,7 +31,7 @@ class RISE(BlackBoxXAIMethod):
     def __init__(
         self,
         model: ov.Model,
-        postprocess_fn: Callable[[OVDict], np.ndarray],
+        postprocess_fn: Callable[[Mapping], np.ndarray],
         preprocess_fn: Callable[[np.ndarray], np.ndarray] = IdentityPreprocessFN(),
         device_name: str = "CPU",
         prepare_model: bool = True,

openvino_xai/methods/factory.py

@@ -2,11 +2,10 @@
 # SPDX-License-Identifier: Apache-2.0
 
 from abc import ABC, abstractmethod
-from typing import Callable, List
+from typing import Callable, List, Mapping
 
 import numpy as np
 import openvino.runtime as ov
-from openvino.runtime.utils.data_helpers.wrappers import OVDict
 
 from openvino_xai.common.parameters import Method, Task
 from openvino_xai.common.utils import IdentityPreprocessFN, logger
@@ -183,7 +182,7 @@ def create_method(
         cls,
         task: Task,
         model: ov.Model,
-        postprocess_fn: Callable[[OVDict], np.ndarray],
+        postprocess_fn: Callable[[Mapping], np.ndarray],
         preprocess_fn: Callable[[np.ndarray], np.ndarray] = IdentityPreprocessFN(),
         device_name: str = "CPU",
         **kwargs,
@@ -197,7 +196,7 @@ def create_method(
     @staticmethod
     def create_classification_method(
         model: ov.Model,
-        postprocess_fn: Callable[[OVDict], np.ndarray],
+        postprocess_fn: Callable[[Mapping], np.ndarray],
         preprocess_fn: Callable[[np.ndarray], np.ndarray] = IdentityPreprocessFN(),
         device_name: str = "CPU",
         **kwargs,
@@ -207,7 +206,7 @@ def create_classification_method(
         :param model: OV IR model.
         :type model: ov.Model
         :param postprocess_fn: Preprocessing function that extract scores from IR model output.
-        :type postprocess_fn: Callable[[OVDict], np.ndarray]
+        :type postprocess_fn: Callable[[Mapping], np.ndarray]
         :param preprocess_fn: Preprocessing function, identity function by default
             (assume input images are already preprocessed by user).
         :type preprocess_fn: Callable[[np.ndarray], np.ndarray]

tests/intg/test_classification.py

@@ -6,7 +6,7 @@
 
 import cv2
 import numpy as np
-import openvino.runtime as ov
+import openvino as ov
 import pytest
 
 import openvino_xai.api.api as xai

tests/intg/test_detection.py

@@ -7,7 +7,7 @@
 import addict
 import cv2
 import numpy as np
-import openvino.runtime as ov
+import openvino as ov
 import pytest
 
 from openvino_xai.common.parameters import Method, Task

tests/perf/test_performance.py

@@ -9,7 +9,7 @@
 
 import cv2
 import numpy as np
-import openvino.runtime as ov
+import openvino as ov
 import pandas as pd
 import pytest
 

tests/unit/common/test_utils.py

@@ -3,7 +3,7 @@
 
 from pathlib import Path
 
-import openvino.runtime as ov
+import openvino as ov
 
 from openvino_xai.api.api import insert_xai
 from openvino_xai.common.parameters import Task

tests/unit/explanation/test_explainer.py

@@ -5,7 +5,7 @@
 
 import cv2
 import numpy as np
-import openvino.runtime as ov
+import openvino as ov
 import pytest
 
 from openvino_xai.api.api import insert_xai

tests/unit/insertion/test_model_parser.py

@@ -1,7 +1,7 @@
 # Copyright (C) 2023-2024 Intel Corporation
 # SPDX-License-Identifier: Apache-2.0
 
-import openvino.runtime as ov
+import openvino as ov
 import pytest
 
 from openvino_xai.common.utils import retrieve_otx_model

tests/unit/methods/black_box/test_black_box_method.py

@@ -5,7 +5,7 @@
 
 import cv2
 import numpy as np
-import openvino.runtime as ov
+import openvino as ov
 import pytest
 
 from openvino_xai.common.utils import retrieve_otx_model

tests/unit/methods/white_box/test_create_method.py

@@ -3,7 +3,7 @@
 
 from pathlib import Path
 
-import openvino.runtime as ov
+import openvino as ov
 import pytest
 
 from openvino_xai.common.parameters import Method, Task