feature: add CV shap explainability for SageMaker Clarify

Enabled CV explainability for SHAP in SageMaker Clarify.

- Since CV explainability parameters are part of SHAP parameters, SHAPConfig includes a new parameter 'image_config' which can be set to a non-None value. Default is None.
- To handle image config parameters, this change created a new class ImageConfig which accepts the following parameters: model_type, num_segments, feature_extraction_method, segment_compactness, max_objects, iou_threshold, context
- To enable image data to be used with Clarify, added a new accepted 'dataset_type' 'application/x-image' to the list of valid dataset_types in DataConfig

Co-authored-by: keerthanvasist <[email protected]>
Co-authored-by: Pranav Krishnan <[email protected]>

Author: 3 people

src/sagemaker/clarify.py

@@ -18,9 +18,9 @@
 import logging
 import os
 import re
+
 import tempfile
 from abc import ABC, abstractmethod
-
 from sagemaker import image_uris, s3, utils
 from sagemaker.processing import ProcessingInput, ProcessingOutput, Processor
 
@@ -67,7 +67,12 @@ def __init__(
                 optional field in all cases except when the dataset contains more than one file,
                 and `save_local_shap_values` is set to true in SHAPConfig.
         """
-        if dataset_type not in ["text/csv", "application/jsonlines", "application/x-parquet"]:
+        if dataset_type not in [
+            "text/csv",
+            "application/jsonlines",
+            "application/x-parquet",
+            "application/x-image",
+        ]:
             raise ValueError(
                 f"Invalid dataset_type '{dataset_type}'."
                 f" Please check the API documentation for the supported dataset types."
@@ -212,7 +217,14 @@ def __init__(
                 )
             self.predictor_config["accept_type"] = accept_type
         if content_type is not None:
-            if content_type not in ["text/csv", "application/jsonlines"]:
+            if content_type not in [
+                "text/csv",
+                "application/jsonlines",
+                "image/jpeg",
+                "image/jpg",
+                "image/png",
+                "application/x-npy",
+            ]:
                 raise ValueError(
                     f"Invalid content_type {content_type}."
                     f" Please choose text/csv or application/jsonlines."
@@ -456,6 +468,65 @@ def get_text_config(self):
         return copy.deepcopy(self.text_config)
 
 
+class ImageConfig:
+    """Config object for handling images"""
+
+    def __init__(
+        self,
+        model_type,
+        num_segments=None,
+        feature_extraction_method=None,
+        segment_compactness=None,
+        max_objects=None,
+        iou_threshold=None,
+        context=None,
+    ):
+        """Initializes all configuration parameters needed for SHAP CV explainability
+
+        Args:
+            model_type (str): Specifies the type of CV model. Options:
+            (IMAGE_CLASSIFICATION | OBJECT_DETECTION).
+            num_segments (None or int): Clarify uses SKLearn's SLIC method for image segmentation
+            to generate features/superpixels. num_segments specifies approximate
+            number of segments to be generated. Default is None. SLIC will default to
+            100 segments.
+            feature_extraction_method (None or str): method used for extracting features from the
+            image.ex. "segmentation". Default is segmentation.
+            segment_compactness (None or float): Balances color proximity and space proximity.
+            Higher values give more weight to space proximity, making superpixel
+            shapes more square/cubic. We recommend exploring possible values on a log
+            scale, e.g., 0.01, 0.1, 1, 10, 100, before refining around a chosen value.
+            max_objects (None or int): maximum number of objects displayed. Object detection
+            algorithm may detect more than max_objects number of objects in a single
+            image. The top max_objects number of objects according to confidence score
+            will be displayed.
+            iou_threshold (None or float): minimum intersection over union for the object
+            bounding box to consider its confidence score for computing SHAP values [0.0, 1.0].
+            This parameter is used for the object detection case.
+            context (None or float): refers to the portion of the image outside of the bounding box.
+            Scale is [0.0, 1.0]. If set to 1.0, whole image is considered, if set to
+            0.0 only the image inside bounding box is considered.
+        """
+        self.image_config = {}
+
+        if model_type not in ["OBJECT_DETECTION", "IMAGE_CLASSIFICATION"]:
+            raise ValueError(
+                "Clarify SHAP only supports object detection and image classification methods. "
+                "Please set model_type to OBJECT_DETECTION or IMAGE_CLASSIFICATION."
+            )
+        self.image_config["model_type"] = model_type
+        _set(num_segments, "num_segments", self.image_config)
+        _set(feature_extraction_method, "feature_extraction_method", self.image_config)
+        _set(segment_compactness, "segment_compactness", self.image_config)
+        _set(max_objects, "max_objects", self.image_config)
+        _set(iou_threshold, "iou_threshold", self.image_config)
+        _set(context, "context", self.image_config)
+
+    def get_image_config(self):
+        """Returns the image config part of an analysis config dictionary."""
+        return copy.deepcopy(self.image_config)
+
+
 class SHAPConfig(ExplainabilityConfig):
     """Config class of SHAP."""
 
@@ -469,6 +540,7 @@ def __init__(
         seed=None,
         num_clusters=None,
         text_config=None,
+        image_config=None,
     ):
         """Initializes config for SHAP.
 
@@ -497,7 +569,10 @@ def __init__(
                 computes a baseline dataset via a clustering algorithm (K-means/K-prototypes).
                 num_clusters is a parameter for this algorithm. num_clusters will be the resulting
                 size of the baseline dataset. If not provided, Clarify job will use a default value.
-            text_config (:class:`~sagemaker.clarify.TextConfig`): Config to handle text features
+            text_config (:class:`~sagemaker.clarify.TextConfig`): Config to handle text features.
+                Default is None
+            image_config (:class:`~sagemaker.clarify.ImageConfig`): Config to handle image features.
+                Default is None
         """
         if agg_method is not None and agg_method not in ["mean_abs", "median", "mean_sq"]:
             raise ValueError(
@@ -512,17 +587,11 @@ def __init__(
             "use_logit": use_logit,
             "save_local_shap_values": save_local_shap_values,
         }
-        if baseline is not None:
-            self.shap_config["baseline"] = baseline
-        if num_samples is not None:
-            self.shap_config["num_samples"] = num_samples
-        if agg_method is not None:
-            self.shap_config["agg_method"] = agg_method
-        if seed is not None:
-            self.shap_config["seed"] = seed
-        if num_clusters is not None:
-            self.shap_config["num_clusters"] = num_clusters
+        _set(baseline, "baseline", self.shap_config)
+        _set(num_samples, "num_samples", self.shap_config)
+        _set(agg_method, "agg_method", self.shap_config)
         _set(seed, "seed", self.shap_config)
+        _set(num_clusters, "num_clusters", self.shap_config)
         if text_config:
             _set(text_config.get_text_config(), "text_config", self.shap_config)
             if not save_local_shap_values:
@@ -531,6 +600,8 @@ def __init__(
                     "Consider setting save_local_shap_values=True to inspect local text "
                     "explanations."
                 )
+        if image_config:
+            _set(image_config.get_image_config(), "image_config", self.shap_config)
 
     def get_explainability_config(self):
         """Returns config."""

tests/unit/test_clarify.py

@@ -28,6 +28,7 @@
     SageMakerClarifyProcessor,
     SHAPConfig,
     TextConfig,
+    ImageConfig,
 )
 
 JOB_NAME_PREFIX = "my-prefix"
@@ -254,17 +255,34 @@ def test_shap_config():
     seed = 123
     granularity = "sentence"
     language = "german"
+    model_type = "IMAGE_CLASSIFICATION"
+    num_segments = 2
+    feature_extraction_method = "segmentation"
+    segment_compactness = 10
+    max_objects = 4
+    iou_threshold = 0.5
+    context = 1.0
     text_config = TextConfig(
         granularity=granularity,
         language=language,
     )
+    image_config = ImageConfig(
+        model_type=model_type,
+        num_segments=num_segments,
+        feature_extraction_method=feature_extraction_method,
+        segment_compactness=segment_compactness,
+        max_objects=max_objects,
+        iou_threshold=iou_threshold,
+        context=context,
+    )
     shap_config = SHAPConfig(
         baseline=baseline,
         num_samples=num_samples,
         agg_method=agg_method,
         use_logit=use_logit,
         seed=seed,
         text_config=text_config,
+        image_config=image_config,
     )
     expected_config = {
         "shap": {
@@ -278,6 +296,15 @@ def test_shap_config():
                 "granularity": granularity,
                 "language": language,
             },
+            "image_config": {
+                "model_type": model_type,
+                "num_segments": num_segments,
+                "feature_extraction_method": feature_extraction_method,
+                "segment_compactness": segment_compactness,
+                "max_objects": max_objects,
+                "iou_threshold": iou_threshold,
+                "context": context,
+            },
         }
     }
     assert expected_config == shap_config.get_explainability_config()
@@ -359,6 +386,50 @@ def test_invalid_text_config():
     assert "Invalid language invalid. Please choose among ['chinese'," in str(error.value)
 
 
+def test_image_config():
+    model_type = "IMAGE_CLASSIFICATION"
+    num_segments = 2
+    feature_extraction_method = "segmentation"
+    segment_compactness = 10
+    max_objects = 4
+    iou_threshold = 0.5
+    context = 1.0
+    image_config = ImageConfig(
+        model_type=model_type,
+        num_segments=num_segments,
+        feature_extraction_method=feature_extraction_method,
+        segment_compactness=segment_compactness,
+        max_objects=max_objects,
+        iou_threshold=iou_threshold,
+        context=context,
+    )
+    expected_config = {
+        "model_type": model_type,
+        "num_segments": num_segments,
+        "feature_extraction_method": feature_extraction_method,
+        "segment_compactness": segment_compactness,
+        "max_objects": max_objects,
+        "iou_threshold": iou_threshold,
+        "context": context,
+    }
+
+    assert expected_config == image_config.get_image_config()
+
+
+def test_invalid_image_config():
+    model_type = "OBJECT_SEGMENTATION"
+    num_segments = 2
+    with pytest.raises(ValueError) as error:
+        ImageConfig(
+            model_type=model_type,
+            num_segments=num_segments,
+        )
+    assert (
+        "Clarify SHAP only supports object detection and image classification methods. "
+        "Please set model_type to OBJECT_DETECTION or IMAGE_CLASSIFICATION." in str(error.value)
+    )
+
+
 def test_invalid_shap_config():
     with pytest.raises(ValueError) as error:
         SHAPConfig(
@@ -665,6 +736,7 @@ def _run_test_explain(
     model_scores,
     expected_predictor_config,
     expected_text_config=None,
+    expected_image_config=None,
 ):
     with patch.object(SageMakerClarifyProcessor, "_run", return_value=None) as mock_method:
         explanation_configs = None
@@ -684,21 +756,6 @@ def _run_test_explain(
             job_name="test",
             experiment_config={"ExperimentName": "AnExperiment"},
         )
-        expected_shap_config = {
-            "baseline": [
-                [
-                    0.26124998927116394,
-                    0.2824999988079071,
-                    0.06875000149011612,
-                ]
-            ],
-            "num_samples": 100,
-            "agg_method": "mean_sq",
-            "use_logit": False,
-            "save_local_shap_values": True,
-        }
-        if expected_text_config:
-            expected_shap_config["text_config"] = expected_text_config
         expected_analysis_config = {
             "dataset_type": "text/csv",
             "headers": [
@@ -710,9 +767,6 @@ def _run_test_explain(
             ],
             "label": "Label",
             "joinsource_name_or_index": "F4",
-            "methods": {
-                "shap": expected_shap_config,
-            },
             "predictor": expected_predictor_config,
         }
         expected_explanation_configs = {}
@@ -732,6 +786,8 @@ def _run_test_explain(
             }
             if expected_text_config:
                 expected_explanation_configs["shap"]["text_config"] = expected_text_config
+            if expected_image_config:
+                expected_explanation_configs["shap"]["image_config"] = expected_image_config
         if pdp_config:
             expected_explanation_configs["pdp"] = {
                 "features": ["F1", "F2"],
@@ -963,3 +1019,70 @@ def test_shap_with_text_config(
         expected_predictor_config,
         expected_text_config=expected_text_config,
     )
+
+
+@patch("sagemaker.utils.name_from_base", return_value=JOB_NAME)
+def test_shap_with_image_config(
+    name_from_base,
+    clarify_processor,
+    clarify_processor_with_job_name_prefix,
+    data_config,
+    model_config,
+):
+    model_type = "IMAGE_CLASSIFICATION"
+    num_segments = 2
+    feature_extraction_method = "segmentation"
+    segment_compactness = 10
+    max_objects = 4
+    iou_threshold = 0.5
+    context = 1.0
+    image_config = ImageConfig(
+        model_type=model_type,
+        num_segments=num_segments,
+        feature_extraction_method=feature_extraction_method,
+        segment_compactness=segment_compactness,
+        max_objects=max_objects,
+        iou_threshold=iou_threshold,
+        context=context,
+    )
+
+    shap_config = SHAPConfig(
+        baseline=[
+            [
+                0.26124998927116394,
+                0.2824999988079071,
+                0.06875000149011612,
+            ]
+        ],
+        num_samples=100,
+        agg_method="mean_sq",
+        image_config=image_config,
+    )
+
+    expected_image_config = {
+        "model_type": model_type,
+        "num_segments": num_segments,
+        "feature_extraction_method": feature_extraction_method,
+        "segment_compactness": segment_compactness,
+        "max_objects": max_objects,
+        "iou_threshold": iou_threshold,
+        "context": context,
+    }
+    expected_predictor_config = {
+        "model_name": "xgboost-model",
+        "instance_type": "ml.c5.xlarge",
+        "initial_instance_count": 1,
+    }
+
+    _run_test_explain(
+        name_from_base,
+        clarify_processor,
+        clarify_processor_with_job_name_prefix,
+        data_config,
+        model_config,
+        shap_config,
+        None,
+        None,
+        expected_predictor_config,
+        expected_image_config=expected_image_config,
+    )