"""This file contains tests for the Image masker.""" import tempfile import numpy as np import pytest import shap from shap.utils import assert_import from shap.utils._exceptions import DimensionError try: assert_import("cv2") except ImportError: pytestmark = pytest.mark.skip("opencv not installed") def test_serialization_image_masker_inpaint_telea(): """Make sure image serialization works with inpaint telea mask.""" test_image_height = 500 test_image_width = 500 test_data = np.ones((test_image_height, test_image_width, 3)) * 50 test_shape = (test_image_height, test_image_width, 3) # initialize image masker original_image_masker = shap.maskers.Image("inpaint_telea", test_shape) with tempfile.TemporaryFile() as temp_serialization_file: # serialize independent masker original_image_masker.save(temp_serialization_file) temp_serialization_file.seek(0) # deserialize masker new_image_masker = shap.maskers.Image.load(temp_serialization_file) mask = np.ones((test_image_height, test_image_width, 3)) mask = mask.astype(int) mask[0][0] = 0 mask[4][0] = 0 # comparing masked values assert np.array_equal(original_image_masker(mask, test_data), new_image_masker(mask, test_data)) def test_serialization_image_masker_inpaint_ns(): """Make sure image serialization works with inpaint ns mask.""" test_image_height = 500 test_image_width = 500 test_data = np.ones((test_image_height, test_image_width, 3)) * 50 test_shape = (test_image_height, test_image_width, 3) # initialize image masker original_image_masker = shap.maskers.Image("inpaint_ns", test_shape) with tempfile.TemporaryFile() as temp_serialization_file: # serialize independent masker original_image_masker.save(temp_serialization_file) temp_serialization_file.seek(0) # deserialize masker new_image_masker = shap.maskers.Image.load(temp_serialization_file) mask = np.ones((test_image_height, test_image_width, 3)) mask = mask.astype(int) mask[0][0] = 0 mask[4][0] = 0 # comparing masked values assert np.array_equal(original_image_masker(mask, test_data), new_image_masker(mask, test_data)) def test_serialization_image_masker_blur(): """Make sure image serialization works with blur mask.""" test_image_height = 500 test_image_width = 500 test_data = np.ones((test_image_height, test_image_width, 3)) * 50 test_shape = (test_image_height, test_image_width, 3) # initialize image masker original_image_masker = shap.maskers.Image("blur(10,10)", test_shape) with tempfile.TemporaryFile() as temp_serialization_file: # serialize independent masker original_image_masker.save(temp_serialization_file) temp_serialization_file.seek(0) # deserialize masker new_image_masker = shap.maskers.Image.load(temp_serialization_file) mask = np.ones((test_image_height, test_image_width, 3)) mask = mask.astype(int) mask[0][0] = 0 mask[4][0] = 0 # comparing masked values assert np.array_equal(original_image_masker(mask, test_data), new_image_masker(mask, test_data)) def test_serialization_image_masker_mask(): """Make sure image serialization works.""" test_image_height = 500 test_image_width = 500 test_data = np.ones((test_image_height, test_image_width, 3)) * 50 test_shape = (test_image_height, test_image_width, 3) test_mask = np.ones((test_image_height, test_image_width, 3)) # initialize image masker original_image_masker = shap.maskers.Image(test_mask, test_shape) with tempfile.TemporaryFile() as temp_serialization_file: # serialize independent masker original_image_masker.save(temp_serialization_file) temp_serialization_file.seek(0) # deserialize masker new_image_masker = shap.maskers.Image.load(temp_serialization_file) mask = np.ones((test_image_height, test_image_width, 3)) mask = mask.astype(int) mask[0][0] = 0 mask[4][0] = 0 # comparing masked values assert np.array_equal(original_image_masker(mask, test_data), new_image_masker(mask, test_data)) def test_init_string_mask_without_shape(): """Make sure masker raises error when initializing with string mask value without shape""" with pytest.raises(TypeError): shap.maskers.Image("inpaint_telea") def test_init_ndarray_mask_without_shape(): """Make sure that shape is inferred correctly from np.array when no shape is passed""" mask_value = np.zeros((5, 5, 3)) image_masker = shap.maskers.Image(mask_value) # no shape parameter passed assert image_masker.input_shape == (5, 5, 3) assert image_masker.shape == (1, 75) # 5*5*3 = 75, when flattened def test_init_scalar_mask_with_shape(): """Make sure mask_value is expanded to a flat array of the scalar when mask_value is an int""" image_masker = shap.maskers.Image(5, shape=(5, 5, 3)) assert image_masker.input_shape == (5, 5, 3) assert image_masker.mask_value.shape == (75,) assert np.all(image_masker.mask_value == 5) def test_call_with_torch_tensor(): """x is converted from torch Tensor to numpy array before masking""" torch = pytest.importorskip("torch") image_masker = shap.maskers.Image(np.zeros((5, 5, 3))) x = torch.zeros(5, 5, 3) mask = np.ones(75, dtype=bool) result = image_masker(mask, x) assert isinstance(result[0], np.ndarray) def test_call_image_masker_shape_mismatch(): """Make sure DimensionError when image and masker shapes mismatch""" image_masker = shap.maskers.Image(np.zeros((4, 4, 3))) image = np.zeros((5, 5, 3)) mask = np.ones(48, dtype=bool) with pytest.raises(DimensionError): image_masker(mask, image) def test_call_image_masker_with_no_mask(): """Make sure the entire image is masked when mask=None is passed.""" image_masker = shap.maskers.Image(np.zeros((5, 5, 3))) image = np.ones((5, 5, 3)) result = image_masker(None, image) assert np.all(result[0] == 0) # entire image masked with mask_value (all 0) def test_inpaint_fully_masked_image(): """Make sure mean colour is used when entire image is masked during inpainting""" image_masker = shap.maskers.Image("inpaint_telea", (5, 5, 3)) image = np.zeros((5, 5, 3)) image[0, 0, :] = 100 # one pixel has 100 for all three channels, the rest are all 0 # mean = 100/ (5*5) = 4.00 result = image_masker(None, image) assert np.all(result[0] == 4.0)