Files
2026-07-13 13:22:52 +08:00

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)