128 lines
4.7 KiB
Python
128 lines
4.7 KiB
Python
import matplotlib.pyplot as plt
|
|
import numpy as np
|
|
import pytest
|
|
|
|
import shap
|
|
|
|
|
|
def set_reproducible_mpl_rcparams() -> None:
|
|
"""Set some matplotlib rcParams to ensure consistency between versions.
|
|
|
|
In matplotlib 3.10, the default value of "image.interpolation_stage" changed from "data" to "auto"
|
|
which can lead to slighly different results.
|
|
|
|
Careful: the @pytest.mark.mpl_image_compare decorator will override rcParams,
|
|
so this change must be done *after* the fixtures are called.
|
|
"""
|
|
plt.rcParams["image.interpolation"] = "bilinear"
|
|
plt.rcParams["image.interpolation_stage"] = "data"
|
|
|
|
|
|
@pytest.fixture
|
|
def imagenet50_example() -> tuple[np.ndarray, np.ndarray]:
|
|
# Return a subset of the imagenet50 dataset, normalised for plotting
|
|
images, labels = shap.datasets.imagenet50()
|
|
images = images / 255
|
|
return images, labels
|
|
|
|
|
|
@pytest.fixture
|
|
def explanation_multi_example(imagenet50_example) -> shap.Explanation:
|
|
# Return an explanation example
|
|
images, _ = imagenet50_example
|
|
n_images = 2
|
|
n_classes = 4
|
|
|
|
images = images[:n_images]
|
|
shap_values_single = (images - images.mean()) / images.max(keepdims=True)
|
|
assert shap_values_single.shape == images.shape
|
|
|
|
# Just repeat the same SHAP values for each class
|
|
shap_values_multi = np.stack([shap_values_single for _ in range(n_classes)], axis=-1)
|
|
assert shap_values_multi.shape[-1] == n_classes
|
|
|
|
explanation = shap.Explanation(values=shap_values_multi, data=images, output_names=[x for x in range(n_classes)])
|
|
return explanation
|
|
|
|
|
|
@pytest.mark.mpl_image_compare
|
|
def test_image_single(imagenet50_example):
|
|
set_reproducible_mpl_rcparams()
|
|
images, _ = imagenet50_example
|
|
images = images[0]
|
|
shap_values = (images - images.mean()) / images.max(keepdims=True)
|
|
explanation = shap.Explanation(values=shap_values, data=images)
|
|
shap.image_plot(explanation, show=False)
|
|
return plt.gcf()
|
|
|
|
|
|
@pytest.mark.mpl_image_compare
|
|
def test_image_multi_no_labels(explanation_multi_example):
|
|
"""Multiple images, multiple classes, labels taken from explanation object"""
|
|
set_reproducible_mpl_rcparams()
|
|
shap.image_plot(explanation_multi_example, show=False)
|
|
return plt.gcf()
|
|
|
|
|
|
@pytest.mark.mpl_image_compare
|
|
def test_image_multi(explanation_multi_example):
|
|
"""Multiple images, multiple classes, a common set of labels for all rows"""
|
|
set_reproducible_mpl_rcparams()
|
|
*_, n_classes = explanation_multi_example.shape
|
|
labels = [f"Class {x + 1}" for x in range(n_classes)]
|
|
shap.image_plot(explanation_multi_example, labels=labels, show=False)
|
|
return plt.gcf()
|
|
|
|
|
|
@pytest.mark.mpl_image_compare
|
|
def test_image_multi_labels_per_row_list(explanation_multi_example):
|
|
"""Multiple images, multiple classes, a set of labels per row as list of lists"""
|
|
set_reproducible_mpl_rcparams()
|
|
n_images, *_, n_classes = explanation_multi_example.shape
|
|
labels = [[f"Class {x + 1 + y * n_classes}" for x in range(n_classes)] for y in range(n_images)]
|
|
shap.image_plot(explanation_multi_example, labels=labels, show=False)
|
|
return plt.gcf()
|
|
|
|
|
|
@pytest.mark.mpl_image_compare
|
|
def test_image_multi_labels_per_row_ndarray(explanation_multi_example):
|
|
"""Multiple images, multiple classes, a set of labels per row as np.array"""
|
|
set_reproducible_mpl_rcparams()
|
|
n_images, *_, n_classes = explanation_multi_example.shape
|
|
labels = np.array([[f"Class {x + 1 + y * n_classes}" for x in range(n_classes)] for y in range(n_images)])
|
|
shap.image_plot(explanation_multi_example, labels=labels, show=False)
|
|
return plt.gcf()
|
|
|
|
|
|
def test_random_single_image():
|
|
"""Just make sure the image_plot function doesn't crash."""
|
|
shap.image_plot(np.random.randn(3, 20, 20), np.random.randn(3, 20, 20), show=False)
|
|
|
|
|
|
def test_random_multi_image():
|
|
"""Just make sure the image_plot function doesn't crash."""
|
|
shap.image_plot([np.random.randn(3, 20, 20) for i in range(3)], np.random.randn(3, 20, 20), show=False)
|
|
|
|
|
|
def test_image_to_text_single():
|
|
"""Just make sure the image_to_text function doesn't crash."""
|
|
|
|
class MockImageExplanation:
|
|
"""Fake explanation object."""
|
|
|
|
def __init__(self, data, values, output_names):
|
|
self.data = data
|
|
self.values = values
|
|
self.output_names = output_names
|
|
|
|
test_image_height = 500
|
|
test_image_width = 500
|
|
test_word_length = 4
|
|
|
|
test_data = np.ones((test_image_height, test_image_width, 3)) * 50
|
|
test_values = np.random.rand(test_image_height, test_image_width, 3, test_word_length)
|
|
test_output_names = np.array([str(i) for i in range(test_word_length)])
|
|
|
|
shap_values_test = MockImageExplanation(test_data, test_values, test_output_names)
|
|
shap.plots.image_to_text(shap_values_test)
|