Suport 2D image attr in visualize_image_attr (#1761)

Summary:
Pull Request resolved: #1761

Objective: plan to unify the image attr logic between llm attr and the previous visualize_image_attr.

But visualize_image_attr is mainly written aiming for gradient-based attr. So this may take a while. As the 1st step:

- Revise `_normalize_attr` for readability.
- Remove `reduction` from `_normalize_attr` (how to aggregate attr (mean/sum/max) is not a part of normalization)
- Support `(H, W)` attr

Reviewed By: craymichael

Differential Revision: D90046963

fbshipit-source-id: 82dddaa591bc395ca8c511db17a19522213c2f3e

Author: aobo-y

captum/attr/_utils/visualization.py

@@ -197,34 +197,32 @@ def _cumulative_sum_threshold(
 
 
 def _normalize_attr(
-    attr: npt.NDArray,
+    attr: npt.NDArray,  # 2D (H, W)
     sign: str,
     outlier_perc: Union[int, float] = 2,
-    reduction_axis: Optional[int] = None,
 ) -> npt.NDArray:
-    attr_combined = attr
-    if reduction_axis is not None:
-        attr_combined = np.sum(attr, axis=reduction_axis)
-
-    # Choose appropriate signed values and rescale, removing given outlier percentage.
-    if VisualizeSign[sign].value == VisualizeSign.all.value:
-        threshold = _cumulative_sum_threshold(
-            np.abs(attr_combined), 100.0 - outlier_perc
-        )
-    elif VisualizeSign[sign].value == VisualizeSign.positive.value:
-        attr_combined = (attr_combined > 0) * attr_combined
-        threshold = _cumulative_sum_threshold(attr_combined, 100.0 - outlier_perc)
-    elif VisualizeSign[sign].value == VisualizeSign.negative.value:
-        attr_combined = (attr_combined < 0) * attr_combined
-        threshold = -1 * _cumulative_sum_threshold(
-            np.abs(attr_combined), 100.0 - outlier_perc
-        )
-    elif VisualizeSign[sign].value == VisualizeSign.absolute_value.value:
-        attr_combined = np.abs(attr_combined)
-        threshold = _cumulative_sum_threshold(attr_combined, 100.0 - outlier_perc)
+    sign_type = VisualizeSign[sign]
+
+    # Apply sign-specific transformation to filter/transform attribution values
+    if sign_type == VisualizeSign.all:
+        pass  # Keep all values as-is
+    elif sign_type == VisualizeSign.positive:
+        attr = np.maximum(attr, 0)
+    elif sign_type == VisualizeSign.negative:
+        attr = np.minimum(attr, 0)
+    elif sign_type == VisualizeSign.absolute_value:
+        attr = np.abs(attr)
     else:
         raise AssertionError("Visualize Sign type is not valid.")
-    return _normalize_scale(attr_combined, threshold)
+
+    # Compute threshold from absolute values, removing given outlier percentage
+    threshold = _cumulative_sum_threshold(np.abs(attr), 100.0 - outlier_perc)
+
+    # For negative sign, threshold should be negative to match the sign of values
+    if sign_type == VisualizeSign.negative:
+        threshold = -threshold
+
+    return _normalize_scale(attr, threshold)
 
 
 def _create_default_plot(
@@ -371,9 +369,8 @@ def visualize_image_attr(
     Args:
 
         attr (numpy.ndarray): Numpy array corresponding to attributions to be
-                    visualized. Shape must be in the form (H, W, C), with
-                    channels as last dimension. Shape must also match that of
-                    the original image if provided.
+                    visualized. Shape must be in the form (H, W, C) or (H, W).
+                    Shape must also match that of the original image if provided.
         original_image (numpy.ndarray, optional): Numpy array corresponding to
                     original image. Shape must be in the form (H, W, C), with
                     channels as the last dimension. Image can be provided either
@@ -510,8 +507,12 @@ def visualize_image_attr(
         "alpha_scaling": _visualize_alpha_scaling,
         "original_image": _visualize_original_image,
     }
+
+    # if the attr contains channel, aggregate them by sum
+    if len(attr.shape) == 3:
+        attr = np.sum(attr, axis=2)
     # Choose appropriate signed attributions and normalize.
-    norm_attr = _normalize_attr(attr, sign, outlier_perc, reduction_axis=2)
+    norm_attr = _normalize_attr(attr, sign, outlier_perc)
 
     # Set default colormap and bounds based on sign.
     default_cmap, vmin, vmax = _initialize_cmap_and_vmin_vmax(sign)
@@ -971,7 +972,7 @@ def visualize_timeseries_attr(
     else:
         plt_axis_list = plt_axis
 
-    norm_attr = _normalize_attr(attr, sign, outlier_perc, reduction_axis=None)
+    norm_attr = _normalize_attr(attr, sign, outlier_perc)
 
     # Set default colormap and bounds based on sign.
     default_cmap, vmin, vmax = _initialize_cmap_and_vmin_vmax(sign)