Fixes for deep dive notebook (#35)

openvino_xai/explainer/utils.py

@@ -157,30 +157,31 @@ def get_score(x: np.ndarray, index: int, activation: ActivationType = Activation
     return x[index]
 
 
-def format_to_hwc(image: np.ndarray) -> np.ndarray:
-    """Format image to HWC."""
-    ori_ndim = image.ndim
-    if image.ndim == 4:
-        image = np.squeeze(image, axis=0)
+def format_to_bhwc(image: np.ndarray) -> np.ndarray:
+    """Format image to BHWC from ndim=3 or ndim=4."""
+    if image.ndim == 3:
+        image = np.expand_dims(image, axis=0)
+    if not is_bhwc_layout(image):
+        # bchw layout -> bhwc
+        image = image.transpose((0, 2, 3, 1))
+    return image
 
-    dim0, dim1, dim2 = image.shape
-    if dim0 < dim1 and dim0 < dim2:
-        image = image.transpose((1, 2, 0))
 
-    if ori_ndim:
-        return np.expand_dims(image, axis=0)
-    return image
+def is_bhwc_layout(image: np.array) -> bool:
+    """Check whether layout of image is BHWC."""
+    _, dim0, dim1, dim2 = image.shape
+    if dim0 > dim2 and dim1 > dim2:  # bhwc layout
+        return True
+    return False
 
 
 def infer_size_from_image(image: np.ndarray) -> Tuple[int, int]:
     """Estimate image size."""
-    image = format_to_hwc(image)
-
     if image.ndim == 2:
         return image.shape
-    elif image.ndim == 3:
-        h, w, _ = image.shape
-    elif image.ndim == 4:
+
+    image = format_to_bhwc(image)
+    if image.ndim == 4:
         _, h, w, _ = image.shape
     else:
         raise ValueError(f"Supports only two, three, and four dimensional image, but got {image.ndim}.")

openvino_xai/explainer/visualizer.py

@@ -15,15 +15,28 @@
     Explanation,
     Layout,
 )
-from openvino_xai.explainer.utils import format_to_hwc, infer_size_from_image
+from openvino_xai.explainer.utils import format_to_bhwc, infer_size_from_image
 
 
 def resize(saliency_map: np.ndarray, output_size: Tuple[int, int]) -> np.ndarray:
     """Resize saliency map."""
     x = saliency_map.transpose((1, 2, 0))
-    x = cv2.resize(x, output_size[::-1])
+
+    # Resize fails for tensors with 700 and more channels (targets=all classes scenario)
+    # Resizing in batches instead
+    batch_size = 500
+    channels = x.shape[-1]
+    resized_batches = []
+    for start_idx in range(0, channels, batch_size):
+        end_idx = min(start_idx + batch_size, channels)
+        batch = x[:, :, start_idx:end_idx]
+        resized_batch = cv2.resize(batch, output_size[::-1])
+        resized_batches.append(resized_batch)
+    x = np.concatenate(resized_batches, axis=-1)
+
     if x.ndim == 2:
         return np.expand_dims(x, axis=0)
+
     return x.transpose((2, 0, 1))
 
 
@@ -112,7 +125,7 @@ def visualize(
         :type overlay_weight: float
         """
         if original_input_image is not None:
-            original_input_image = format_to_hwc(original_input_image)
+            original_input_image = format_to_bhwc(original_input_image)
 
         saliency_map_dict = explanation.saliency_map
         class_idx_to_return = list(saliency_map_dict.keys())

openvino_xai/methods/black_box/rise.py

@@ -103,8 +103,9 @@ def _run_synchronous_explanation(
         prob: float,
         seed: int,
     ) -> np.ndarray:
-        _, _, height, width = data_preprocessed.shape
-        input_size = height, width
+        from openvino_xai.explainer.utils import is_bhwc_layout
+
+        input_size = data_preprocessed.shape[1:3] if is_bhwc_layout(data_preprocessed) else data_preprocessed.shape[2:4]
 
         forward_output = self.model_forward(data_preprocessed, preprocess=False)
         logits = self.postprocess_fn(forward_output)
@@ -121,7 +122,10 @@ def _run_synchronous_explanation(
         for _ in tqdm(range(0, num_masks), desc="Explaining in synchronous mode"):
             mask = self._generate_mask(input_size, num_cells, prob, rand_generator)
             # Add channel dimensions for masks
-            masked = mask * data_preprocessed
+            if is_bhwc_layout(data_preprocessed):
+                masked = np.expand_dims(mask, 2) * data_preprocessed
+            else:
+                masked = mask * data_preprocessed
 
             forward_output = self.model_forward(masked, preprocess=False)
             raw_scores = self.postprocess_fn(forward_output)

tests/unit/explanation/test_explanation_utils.py

@@ -8,6 +8,7 @@
     ActivationType,
     get_explain_target_indices,
     get_score,
+    is_bhwc_layout,
 )
 
 VOC_NAMES = [
@@ -77,3 +78,8 @@ def test_get_score():
     x = np.random.rand(1, 5)
     score = get_score(x, 0)
     assert score == x[0][0]
+
+
+def test_is_bhwc_layout():
+    assert is_bhwc_layout(np.empty((1, 224, 224, 3)))
+    assert is_bhwc_layout(np.empty((1, 3, 224, 224))) == False

tests/unit/explanation/test_visualization.py

@@ -59,6 +59,10 @@ def test_resize():
     input_saliency_map = np.random.randint(0, 255, (1, 3, 3), dtype=np.uint8)
     resized_map = resize(input_saliency_map, (5, 5))
     assert resized_map.shape == (1, 5, 5)
+    # Test resizing functionality with 700+ channels to check all classes scenario
+    input_saliency_map = np.random.randint(0, 255, (700, 3, 3), dtype=np.uint8)
+    resized_map = resize(input_saliency_map, (5, 5))
+    assert resized_map.shape == (700, 5, 5)
 
 
 def test_colormap():