diff --git a/docs/api/concepts/holistic_craft.md b/docs/api/concepts/holistic_craft.md new file mode 100644 index 00000000..6d5da5f0 --- /dev/null +++ b/docs/api/concepts/holistic_craft.md @@ -0,0 +1,351 @@ +# Holistic CRAFT + + + +[View source](https://github.com/deel-ai/xplique/blob/master/xplique/concepts/holistic_craft.py) | +📰 [CRAFT Paper](https://arxiv.org/pdf/2211.10154) | +📰 [Holistic Paper](https://arxiv.org/pdf/2306.07304) + +Holistic CRAFT (Concept Recursive Activation FacTorization) is a variant of the CRAFT method designed to extract concepts from full activation maps rather than image patches. + +This approach preserves the global spatial context and is particularly suitable for object detection models and other tasks where spatial structure across the entire image is important. + +The crop-based approach works well for classification because images of classification datasets are typically dominated by a single, well-centred object: random crops are therefore likely to contain parts of the object of interest and carry relevant signal for concept extraction. In Object Detection, the scenes generally contain multiple objects of varying sizes, often occupying only a small fraction of the image. Random crops drawn from such images are mostly background; the target objects are absent or heavily under-represented in the resulting crop dataset, making the NMF factorization blind to the very patterns it should capture. + +## Supported Object Detection Models + +Holistic CRAFT works with various object detection architectures through specialized latent extractors provided by the `xplique-adapters` package: + +**PyTorch (torchvision & ultralytics):** +- **RetinaNet** - `RetinanetExtractorBuilder` +- **Faster R-CNN** - `FasterRcnnExtractorBuilder` +- **FCOS** - `FcosExtractorBuilder` +- **SSD** - `SSDExtractorBuilder` +- **YOLO** (v11) - `YoloExtractorBuilder` +- **DETR** - `DetrExtractorBuilder` + +**TensorFlow:** +- **RetinaNet** - `RetinaNetExtractorBuilder` + +Each extractor handles the model-specific architecture to split it into the required g(.) and h(.) functions. + +## Supported Classification Models + +For standard classification models, Holistic CRAFT does not require a custom extractor per architecture. Instead, the built-in `LayeredModelExtractorBuilder` can split any layered model at a chosen intermediate layer: + +**PyTorch:** +- Any `torch.nn.Module` — `LayeredModelExtractorBuilder` (from `xplique.concepts.torch.layered_model_latent_extractor`) + +**TensorFlow:** +- Any `tf.keras.Model` — `LayeredModelExtractorBuilder` (from `xplique.concepts.tf.layered_model_latent_extractor`) + +The builder takes the model and a layer index to define the split point. Everything before that layer becomes g(.), and everything after becomes h(.). + +## Key Differences from Regular CRAFT + +| Aspect | Regular CRAFT | Holistic CRAFT | +|--------|---------------|----------------| +| **Input** | Image patches/crops | Full activation maps | +| **Use Case** | Classification tasks | Object detection, Classification | +| **Spatial Context** | Local (patch-level) | Global (full image) | +| **Concepts** | Visual patterns in patches | Spatial activation patterns | +| **Performance** | Extracts many crops per image | Processes full feature maps directly | + +## Workflow + +Holistic CRAFT follows the same core principle as CRAFT but operates on full images instead of patches: + +1. **Extract Activations**: Pass input images through the model's encoder (g) to obtain spatial activation maps from an intermediate layer +2. **Factorize Concepts**: Apply Non-negative Matrix Factorization (NMF) to these activation maps to discover recurring spatial patterns (concepts) + +!!!warning + Activations must be non-negative to use the standard NMF. Ensure a ReLU + or similar activation function is applied before the extraction layer. + Third-party NMF implementations may not have this limitation + (e.g., the Semi-NMF from the Overcomplete library). + +3. **Estimate Importance**: Use any attribution methods available in Xplique (gradient-based, perturbation-based) to rank concept importance +4. **Visualize**: Generate concept heatmaps overlaid on images to show "what" and "where" + +Like regular CRAFT, Holistic CRAFT requires splitting the model into two parts: $(g, h)$ such that $f(x) = (g \cdot h)(x)$. The model $g$ maps input to latent space (activation maps), and $h$ maps latent space to predictions. Concepts are extracted from these activation maps in latent space. + +This split is implemented through three abstractions: + +- **`LatentData`**: A container that holds the intermediate activations produced by $g$. It abstracts away framework-specific tensor formats, providing a unified interface for reading (`get_activations`) and writing (`set_activations`) activations, with the necessary shape conversions (e.g., channel-first to channel-last). + +- **`LatentExtractor`**: Wraps both $g$ (`input_to_latent_model`) and $h$ (`latent_to_logit_model`). It orchestrates the full forward pass, batching, device management, and output formatting. The `TorchLatentExtractor` and `TfLatentExtractor` subclasses provide framework-specific implementations. + +- **`LatentExtractorBuilder`**: A factory that constructs a `LatentExtractor` for a specific model architecture. It handles all the architecture-specific wiring (defining how to split the model, which layer to extract from, and how to format outputs) so that the rest of the CRAFT pipeline remains model-agnostic. + + +## Example + +### Basic Usage with Object Detection + +```python +import xplique +from xplique.concepts import HolisticCraftTorch as Craft +from xplique_adapters.concepts.torch.latent_data_retinanet import RetinanetExtractorBuilder + +# Build a latent extractor that splits the model into g(.) and h(.) +# This provides the input_to_latent (g) and latent_to_logit (h) functions +latent_extractor = RetinanetExtractorBuilder.build( + model, + device="cuda", + nb_classes=91, + extraction_location='resnet', # Choose 'resnet' or 'fpn' + extraction_layer=-1 # Extract from last ResNet feature layer +) + +# Create Holistic CRAFT instance +craft = Craft( + latent_extractor=latent_extractor, + number_of_concepts=10, + device="cuda" +) + +# Fit CRAFT on input images to discover concepts +craft.fit(input_images, class_id=class_id) + +# Display discovered concepts as heatmaps overlaid on images +craft.display_images_per_concept(images=input_images[:5]) + +# Display top 3 images for each concept ranked by activation +craft.display_top_images_per_concept(images=input_images, topk=3) + +# Estimate concept importance on the 20 first images using Gradient×Input method +# (GradientxInput is the default method) +importances_gi = craft.estimate_importance( + images=input_images[:20], + operator=xplique.Tasks.OBJECT_DETECTION, + class_id=class_id, + confidence=0.8 +) + +# Estimate concept importance on the 20 first images using Sobol method +importances_sobol = craft.estimate_importance( + images=input_images[:20], + operator=xplique.Tasks.OBJECT_DETECTION, + class_id=class_id, + confidence=0.8, + # Use Sobol method & its arguments + method="sobol", + grid_size=4, + nb_design=8, + perturbation_function="amplitude", +) + +``` + +### Using Different Attribution Methods to Compute the Concept Importances + +Holistic CRAFT supports various attribution methods for concept importance estimation: + +```python +import xplique +from xplique.concepts import PartialExplainer +from xplique.attributions import VarGrad + +# Use VarGrad for robust importance estimation +vargrad_explainer = PartialExplainer( + explainer_class=VarGrad, + operator=xplique.Tasks.OBJECT_DETECTION, + nb_samples=20, + noise=0.15 +) + +# Compute VarGrad explanation for each concept +explanation_vargrad = craft.compute_explanation_per_concept( + partial_explainer=vargrad_explainer, + images=input_images, + class_id=class_id, + confidence=0.3, +) + +# Reduce the spatial dimension of the explanation +# to compute the final concepts importances +importances_vargrad = craft.reduce_to_importance( + explanation=explanation_vargrad, +) +``` + +### Using a Different NMF Factorizer + +By default, the standard Sklearn NMF is used to factorize the concepts. +But other types of factorizers are supported, such as the ones provided +by the [Overcomplete](https://github.com/KempnerInstitute/overcomplete) project. + +```python +from overcomplete.optimization import SemiNMF +from xplique.concepts.torch.factorizer import OvercompleteFactorizer + +nb_concepts=10 + +# Create a SemiNMF factorizer which allows negative activations +factorizer = OvercompleteFactorizer( + optimizer_class=SemiNMF, + nb_concepts=nb_concepts, + device=device +) + +# Setup Craft to use this factorizer +craft = Craft( + latent_extractor=latent_extractor, + number_of_concepts=nb_concepts, + device=device, + factorizer=factorizer, +) + +craft.fit(input_images) +``` + +## Implementing Your Own Latent Extractor + +If you're working with a model architecture that isn't supported out-of-the-box, you can implement your own latent extractor by following these steps: + +### 1. Create a Custom LatentData Class + +First, create a class that stores the intermediate activations from your model: + +```python +from xplique.concepts.latent_extractor import LatentData +import torch + +class CustomLatentData(LatentData): + def __init__(self, fpn_outs: list, extraction_layer: int = 0): + super().__init__() + self.fpn_outs = fpn_outs + self.extraction_layer = extraction_layer + + def get_activations(self, as_numpy: bool = True, keep_gradients: bool = False): + """Extract activations from the specified layer.""" + activations = self.fpn_outs[self.extraction_layer] + + if not keep_gradients: + activations = activations.detach() + + # Convert from (N, C, H, W) to (N, H, W, C) for Xplique + if len(activations.shape) == 4: + activations = activations.permute(0, 2, 3, 1) + + if as_numpy: + activations = activations.cpu().numpy() + + return activations + + def set_activations(self, values: torch.Tensor) -> None: + """Set activations back into the latent data structure.""" + # Convert from (N, H, W, C) to (N, C, H, W) + if len(values.shape) == 4: + values = values.permute(0, 3, 1, 2) + self.fpn_outs[self.extraction_layer] = values + + def to(self, device: torch.device) -> 'CustomLatentData': + """Move latent data to specified device.""" + self.fpn_outs = [fpn_out.to(device) for fpn_out in self.fpn_outs] + return CustomLatentData(self.fpn_outs, self.extraction_layer) +``` + +### 2. Create a Custom ExtractorBuilder + +Next, implement a builder that splits your model into g(.) and h(.) functions: + +```python +import types +from xplique.concepts.latent_extractor import LatentExtractorBuilder +from xplique.concepts.torch.latent_extractor import TorchLatentExtractor + +class CustomExtractorBuilder(LatentExtractorBuilder): + @classmethod + def build( + cls, + model, + device: str = 'cuda', + extraction_layer: int = -1, + batch_size: int = 1 + ) -> TorchLatentExtractor: + + # Define g(.) function: input → latent activations + def g(self, x): + # Example: extract from backbone/feature pyramid + fpn_outs = self.backbone(x) + return CustomLatentData( + fpn_outs=list(fpn_outs), + extraction_layer=latent_extractor.extraction_layer + ) + + # Define h(.) function: latent activations → predictions + def h(self, latent_data: CustomLatentData): + fpn_outs = latent_data.fpn_outs + outputs = self.head(fpn_outs) + return outputs + + # Bind g and h methods to the model + model.g = types.MethodType(g, model) + model.h = types.MethodType(h, model) + + # Create output formatter (converts raw predictions to MultiBoxTensor) + output_formatter = CustomBoxFormatter() + + # Build the latent extractor + latent_extractor = TorchLatentExtractor( + model, + model.g, + model.h, + latent_data_class=CustomLatentData, + output_formatter=output_formatter, + batch_size=batch_size, + device=device + ) + + # Store extraction layer for later use + latent_extractor.extraction_layer = extraction_layer + return latent_extractor +``` + +### 3. Use Your Custom Extractor with CRAFT + +Once you have your custom extractor, you can use it just like the built-in ones: + +```python +from xplique.concepts import HolisticCraftTorch as Craft + +# Build your custom latent extractor +latent_extractor = CustomExtractorBuilder.build( + model, + device="cuda", + extraction_layer=-1, + batch_size=16 +) + +# Use it with CRAFT +craft = Craft( + latent_extractor=latent_extractor, + number_of_concepts=10, + device="cuda" +) + +# Fit and visualize concepts +craft.fit(input_images) +craft.display_images_per_concept(input_images[:5]) +``` + +### Key Points + +- **g(.) function**: Maps input images to intermediate activations at a chosen layer +- **h(.) function**: Maps latent activations back to final predictions +- **LatentData**: Handles activation extraction with proper shape conversions (PyTorch uses channel-first, Xplique expects channel-last) +- **Output formatter**: Converts model predictions to `MultiBoxTensor` format for compatibility with Xplique + + +## API Reference + +{{xplique.concepts.holistic_craft.HolisticCraft}} + +{{xplique.concepts.holistic_craft.PartialExplainer}} + +## References + +[^1]: [CRAFT: Concept Recursive Activation FacTorization for Explainability (2023).](https://arxiv.org/pdf/2211.10154.pdf) + +[^2]: [A Holistic Approach to Unifying Automatic Concept Extraction and Concept Importance Estimation (2023).](https://arxiv.org/pdf/2306.07304.pdf) diff --git a/docs/index.md b/docs/index.md index 31904c8f..3984dad6 100644 --- a/docs/index.md +++ b/docs/index.md @@ -131,12 +131,19 @@ Finally, the _Metrics_ module covers the current metrics used in explainability. ## 🚀 Quick Start -Xplique requires a version of python higher than 3.7 and several libraries including Tensorflow and Numpy. Installation can be done using Pypi: +Xplique supports Python 3.10 through 3.13. Installation includes TensorFlow, NumPy, and other +runtime dependencies: -```python +```bash pip install xplique ``` +Install the optional PyTorch dependencies for PyTorch models and concept methods: + +```bash +pip install "xplique[torch]" +``` + Now that Xplique is installed, here are some basic examples of what you can do with the available modules. ??? example "Attributions Methods" diff --git a/mkdocs.yml b/mkdocs.yml index f11f2ecc..fab486fb 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -47,6 +47,7 @@ nav: - Cav: api/concepts/cav.md - Tcav: api/concepts/tcav.md - Craft: api/concepts/craft.md + - Holistic Craft: api/concepts/holistic_craft.md - Example based: - API Description: api/example_based/api_example_based.md - Similar Examples: diff --git a/pyproject.toml b/pyproject.toml index bb81f028..4695ee1d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "xplique" -version = "1.4.2" +version = "2.0.0" description = "Explanations toolbox for TensorFlow 2" readme = "README.md" requires-python = ">=3.10,<3.14" @@ -19,7 +19,7 @@ classifiers = [ dependencies = [ "deprecated>=1.1.0", "einops>=0.1.0", - "matplotlib>=1.4.3", + "matplotlib>=3.5", "numpy>=2.0.0", "opencv-python>=3.4.11.39", "ruff>=0.15.0", @@ -31,9 +31,16 @@ dependencies = [ "tensorflow-probability[tf]>=0.20.1,<=0.25.0", ] +[project.optional-dependencies] +torch = [ + "torch>=2.5.0,<=2.10.0", + "torchvision>=0.20.0,<=0.25.0", + "overcomplete>=0.2.8", +] + [dependency-groups] dev = [ - "bump2version", + "bump-my-version", "build", "docutils", "markdown", @@ -55,6 +62,8 @@ dev = [ ] torch = [ "torch>=2.5.0,<=2.10.0", + "torchvision>=0.20.0,<=0.25.0", + "overcomplete>=0.2.8", ] [build-system] @@ -64,8 +73,11 @@ build-backend = "setuptools.build_meta" [tool.setuptools.packages.find] include = ["xplique*"] +[tool.setuptools.package-data] +"xplique.features_visualizations" = ["spectrum_decorrelated.npy"] + [tool.bumpversion] -current_version = "1.4.2" +current_version = "2.0.0" commit = true tag = false diff --git a/tests/attributions/test_hsic.py b/tests/attributions/test_hsic.py index 37e0712e..f5800896 100644 --- a/tests/attributions/test_hsic.py +++ b/tests/attributions/test_hsic.py @@ -4,6 +4,7 @@ import tensorflow_probability as tfp from xplique.attributions import HsicAttributionMethod +from xplique.attributions.base import sanitize_input_output from xplique.attributions.global_sensitivity_analysis import ( BinaryEstimator, RbfEstimator, @@ -14,6 +15,32 @@ from ..utils import generate_data, generate_model +def test_nb_design_positional_compatibility(): + model = generate_model((4, 4, 1), 2) + + method = HsicAttributionMethod(model, 2, 8) + assert method.nb_design == 8 + assert method.masks.shape[-1] == 1 + + multi_channel_method = HsicAttributionMethod(model, 2, 8, nb_channels=2) + assert multi_channel_method.masks.shape[-1] == 2 + + +def test_sanitize_input_output_forwards_keywords(): + class DecoratedExplainer: + @sanitize_input_output + def explain(self, inputs, targets, verbose=False): + return inputs, targets, verbose + + inputs, targets, verbose = DecoratedExplainer().explain( + np.ones((1, 2), dtype=np.float32), np.ones((1, 1), dtype=np.float32), verbose=True + ) + + assert isinstance(inputs, tf.Tensor) + assert isinstance(targets, tf.Tensor) + assert verbose + + def test_hsic_kernels_shape(): """ Test if the kernels are correctly computed. @@ -119,22 +146,21 @@ def test_hsic_invariant_to_design_permutation(Estimator): @pytest.mark.parametrize("Estimator", [SobolevEstimator, BinaryEstimator, RbfEstimator]) def test_hsic_detects_dependent_dimension(Estimator): - nb_design, H, W, C = 64, 2, 1, 1 # two dimensions total + nb_design, H, W, C = 64, 2, 2, 2 rng = tf.random.Generator.from_seed(7) - dep_dim = rng.uniform((nb_design,), dtype=tf.float32) # this drives Y - indep_dim = rng.uniform((nb_design,), dtype=tf.float32) - - # masks[..., 0,0,0] = dep_dim, masks[..., 1,0,0] = indep_dim - masks = tf.stack([dep_dim, indep_dim], axis=1) # (nb_design, 2) - masks = tf.reshape(masks, (nb_design, H, W, C)) # (n,2,1,1) + dependent_position = (1, 0, 1) + dep_dim = rng.uniform((nb_design,), dtype=tf.float32) + mask_dimensions = [rng.uniform((nb_design,), dtype=tf.float32) for _ in range(H * W * C)] + dependent_index = np.ravel_multi_index(dependent_position, (H, W, C)) + mask_dimensions[dependent_index] = dep_dim + masks = tf.reshape(tf.stack(mask_dimensions, axis=1), (nb_design, H, W, C)) outputs = tf.identity(dep_dim) est = Estimator() - scores = est(masks, outputs, nb_design).numpy() # shape (W,H,C) = (1,2,1) - scores_hw = np.transpose(scores, (1, 0, 2)).reshape(H) # back to [H] order - assert np.argmax(scores_hw) == 0 - assert scores_hw[0] > scores_hw[1] + 1e-3 + scores = est(masks, outputs, nb_design).numpy() + assert scores.shape == (H, W, C) + assert np.unravel_index(np.argmax(scores), scores.shape) == dependent_position def test_output_rbf_width_tfp_matches_numpy(): diff --git a/tests/attributions/test_sobol.py b/tests/attributions/test_sobol.py index 7ebffbb1..1e9aaa74 100644 --- a/tests/attributions/test_sobol.py +++ b/tests/attributions/test_sobol.py @@ -21,6 +21,17 @@ from ..utils import almost_equal, generate_data, generate_model +def test_nb_design_positional_compatibility(): + model = generate_model((4, 4, 1), 2) + + method = SobolAttributionMethod(model, 2, 8) + assert method.nb_design == 8 + assert method.masks.shape[-1] == 1 + + multi_channel_method = SobolAttributionMethod(model, 2, 8, nb_channels=2) + assert multi_channel_method.masks.shape[-1] == 2 + + def test_output_shape(): """The output size (h, w) must be the same as the input""" diff --git a/tests/commons/test_model_override.py b/tests/commons/test_model_override.py index c91e046e..1cfb3f36 100644 --- a/tests/commons/test_model_override.py +++ b/tests/commons/test_model_override.py @@ -159,3 +159,37 @@ def test_open_relu(): open_grads_3 = tape.gradient(y3d, x).numpy()[0] assert almost_equal(open_grads_3, 3.0 * x**2.0) + + +def test_override_clones_lambda_without_deserialization(): + """Ensure Lambda layers are cloned without mutating the source model.""" + lambda_layer = tf.keras.layers.Lambda( + lambda inputs, scale: inputs * scale, + output_shape=(4,), + arguments={"scale": 2.0}, + trainable=False, + dtype=tf.float64, + ) + model = tf.keras.Sequential( + [ + tf.keras.layers.Input((4,)), + tf.keras.layers.Activation(tf.nn.relu), + lambda_layer, + ] + ) + + cloned_model = override_relu_gradient(model, guided_relu_policy) + cloned_lambda = cloned_model.layers[-1] + + assert cloned_lambda is not lambda_layer + assert cloned_lambda.function is lambda_layer.function + assert cloned_lambda.arguments == lambda_layer.arguments + assert cloned_lambda.arguments is not lambda_layer.arguments + assert cloned_lambda._output_shape == lambda_layer._output_shape + assert cloned_lambda.trainable is False + assert cloned_lambda.compute_dtype == "float64" + assert model.layers[0].activation in [tf.nn.relu, tf.keras.activations.relu] + assert cloned_model.layers[0].activation not in [tf.nn.relu, tf.keras.activations.relu] + + inputs = tf.constant([[1.0, -1.0, 2.0, -2.0]]) + assert almost_equal(model(inputs), cloned_model(inputs)) diff --git a/tests/commons/test_object_detection_operator.py b/tests/commons/test_object_detection_operator.py index 37bb36d6..317b7d0b 100644 --- a/tests/commons/test_object_detection_operator.py +++ b/tests/commons/test_object_detection_operator.py @@ -191,3 +191,115 @@ def classification_op(model, inputs, targets): [normal_phis, intersection_phis, probability_phis, classification_phis], 2 ): assert not almost_equal(phi1, phi2) + + +def test_object_detection_operator_ignores_zero_padded_detections_in_a_graph(): + valid_predictions = tf.constant([[[0.0, 0.0, 2.0, 2.0, 0.5, 1.0, 0.0]]]) + padded_predictions = tf.pad(valid_predictions, [[0, 0], [0, 2], [0, 0]]) + inputs = tf.zeros((1, 1), dtype=tf.float32) + + def valid_model(_): + return valid_predictions + + def padded_model(_): + return padded_predictions + + def empty_model(_): + return tf.zeros_like(padded_predictions) + + valid_score = object_detection_operator(valid_model, inputs, valid_predictions) + padded_score = object_detection_operator(padded_model, inputs, padded_predictions) + + previous_eager_setting = tf.config.functions_run_eagerly() + try: + tf.config.run_functions_eagerly(False) + + @tf.function + def graph_operator(operator_inputs, targets): + assert tf.inside_function() + return object_detection_operator(padded_model, operator_inputs, targets) + + @tf.function + def graph_empty_operator(operator_inputs, targets): + assert tf.inside_function() + return object_detection_operator(empty_model, operator_inputs, targets) + + graph_score = graph_operator(inputs, padded_predictions) + empty_score = graph_empty_operator(inputs, padded_predictions) + finally: + tf.config.run_functions_eagerly(previous_eager_setting) + + tf.debugging.assert_near(padded_score, valid_score) + tf.debugging.assert_near(graph_score, valid_score) + tf.debugging.assert_equal(empty_score, [0.0]) + + +def test_object_detection_operator_handles_different_detection_counts_per_input(): + predictions = tf.constant( + [ + [ + [0.0, 0.0, 2.0, 2.0, 0.5, 1.0, 0.0], + [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + ], + [ + [0.0, 0.0, 1.0, 1.0, 0.7, 1.0, 0.0], + [2.0, 2.0, 4.0, 4.0, 0.8, 0.0, 1.0], + [5.0, 5.0, 6.0, 6.0, 0.4, 1.0, 0.0], + ], + [ + [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + ], + ], + dtype=tf.float32, + ) + targets = tf.constant( + [ + [ + [0.0, 0.0, 2.0, 2.0, 1.0, 1.0, 0.0], + [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + ], + [ + [0.0, 0.0, 1.0, 1.0, 1.0, 1.0, 0.0], + [2.0, 2.0, 4.0, 4.0, 1.0, 0.0, 1.0], + ], + [ + [0.0, 0.0, 1.0, 1.0, 1.0, 1.0, 0.0], + [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + ], + ], + dtype=tf.float32, + ) + inputs = tf.zeros((3, 1), dtype=tf.float32) + + def model(_): + return predictions + + def first_model(_): + return predictions[0:1, :1] + + def second_model(_): + return predictions[1:2] + + expected = tf.concat( + [ + object_detection_operator(first_model, inputs[:1], targets[0:1, :1]), + object_detection_operator(second_model, inputs[1:2], targets[1:2]), + tf.constant([0.0]), + ], + axis=0, + ) + + batched_score = object_detection_operator(model, inputs, targets) + + @tf.function + def graph_operator(operator_inputs, operator_targets): + assert tf.inside_function() + return object_detection_operator(model, operator_inputs, operator_targets) + + graph_score = graph_operator(inputs, targets) + + tf.debugging.assert_near(batched_score, expected) + tf.debugging.assert_near(graph_score, expected) diff --git a/tests/concepts/test_factorizer_tf.py b/tests/concepts/test_factorizer_tf.py new file mode 100644 index 00000000..f2a212fb --- /dev/null +++ b/tests/concepts/test_factorizer_tf.py @@ -0,0 +1,125 @@ +"""Tests for TensorFlow differentiable NMF factorizer encoding.""" + +import numpy as np +import pytest +import tensorflow as tf + +from xplique.concepts.tf.factorizer import TfSklearnNMFFactorizer + + +def _make_non_negative_data(seed, shape, minimum=0.05): + """Generate deterministic non-negative test data.""" + rng = np.random.default_rng(seed) + return rng.uniform(minimum, 1.0, size=shape).astype(np.float32) + + +def _nmf_encoding_objective(activations, coefficients, concept_bank, alpha_w, l1_ratio): + """Compute the fixed-dictionary NMF encoding objective.""" + n_features = activations.shape[1] + l1_reg = n_features * alpha_w * l1_ratio + l2_reg = n_features * alpha_w * (1.0 - l1_ratio) + + reconstruction = coefficients @ concept_bank + data_fit = 0.5 * np.square(activations - reconstruction).sum() + l1_penalty = l1_reg * np.abs(coefficients).sum() + l2_penalty = 0.5 * l2_reg * np.square(coefficients).sum() + + return data_fit + l1_penalty + l2_penalty + + +@pytest.mark.parametrize("alpha_w,l1_ratio", [(0.0, 0.0), (0.05, 0.3)]) +@pytest.mark.parametrize("tol", [1e-6, 0.0]) +def test_tf_encode_differentiable_matches_sklearn_transform(alpha_w, l1_ratio, tol): + """The differentiable TF solver should match sklearn's fixed-dictionary transform.""" + train = _make_non_negative_data(0, (80, 12)) + test = _make_non_negative_data(1, (20, 12)) + + factorizer = TfSklearnNMFFactorizer( + n_components=5, + alpha_W=alpha_w, + l1_ratio=l1_ratio, + max_iter=1000, + tol=tol, + random_state=0, + ) + factorizer.fit(train) + + expected = factorizer.encode(test) + actual = factorizer.encode_differentiable(tf.constant(test)) + actual_np = actual.numpy() + concept_bank = factorizer.get_concept_bank() + + assert isinstance(actual, tf.Tensor) + assert actual.shape == expected.shape + assert np.all(actual_np >= -1e-6) + + expected_objective = _nmf_encoding_objective( + test, expected, concept_bank, alpha_w=alpha_w, l1_ratio=l1_ratio + ) + actual_objective = _nmf_encoding_objective( + test, actual_np, concept_bank, alpha_w=alpha_w, l1_ratio=l1_ratio + ) + assert actual_objective <= expected_objective * 1.05 + 1e-5 + + expected_reconstruction = expected @ concept_bank + actual_reconstruction = actual_np @ concept_bank + reconstruction_error = np.linalg.norm(actual_reconstruction - expected_reconstruction) + reconstruction_norm = np.linalg.norm(expected_reconstruction) + assert reconstruction_error / (reconstruction_norm + 1e-8) < 0.05 + + +@pytest.mark.parametrize("tol", [1e-6, 0.0]) +def test_tf_encode_differentiable_preserves_gradients(tol): + """The differentiable solver should backpropagate through activations.""" + train = _make_non_negative_data(2, (60, 10)) + test = tf.constant(_make_non_negative_data(3, (16, 10))) + + factorizer = TfSklearnNMFFactorizer( + n_components=4, + alpha_W=1e-2, + max_iter=500, + tol=tol, + random_state=0, + ) + factorizer.fit(train) + + with tf.GradientTape() as tape: + tape.watch(test) + coefficients = factorizer.encode_differentiable(test) + loss = tf.reduce_sum(coefficients) + + gradients = tape.gradient(loss, test) + + assert gradients is not None + assert tf.reduce_all(tf.math.is_finite(gradients)) + assert tf.reduce_sum(tf.abs(gradients)) > 0 + + +def test_tf_encode_differentiable_rejects_negative_activations(): + """NMF encoding should fail when differentiable inputs are negative.""" + train = _make_non_negative_data(6, (30, 6), minimum=0.1) + test = tf.constant(-_make_non_negative_data(7, (10, 6), minimum=0.1)) + + factorizer = TfSklearnNMFFactorizer(n_components=3, max_iter=200, random_state=0) + factorizer.fit(train) + + with pytest.raises(tf.errors.InvalidArgumentError): + factorizer.encode_differentiable(test) + + +def test_tf_encode_differentiable_rejects_unsupported_beta_loss(): + """Unsupported sklearn NMF configurations should fail explicitly.""" + train = _make_non_negative_data(4, (40, 8), minimum=0.1) + test = tf.constant(_make_non_negative_data(5, (12, 8), minimum=0.1)) + + factorizer = TfSklearnNMFFactorizer( + n_components=3, + solver="mu", + beta_loss="kullback-leibler", + max_iter=300, + random_state=0, + ) + factorizer.fit(train) + + with pytest.raises(NotImplementedError): + factorizer.encode_differentiable(test) diff --git a/tests/concepts/test_factorizer_torch.py b/tests/concepts/test_factorizer_torch.py new file mode 100644 index 00000000..5703abe1 --- /dev/null +++ b/tests/concepts/test_factorizer_torch.py @@ -0,0 +1,124 @@ +"""Tests for PyTorch differentiable NMF factorizer encoding.""" + +import numpy as np +import pytest +import torch + +from xplique.concepts.torch.factorizer import TorchSklearnNMFFactorizer + + +def _make_non_negative_data(seed, shape, minimum=0.05): + """Generate deterministic non-negative test data.""" + rng = np.random.default_rng(seed) + return rng.uniform(minimum, 1.0, size=shape).astype(np.float32) + + +def _nmf_encoding_objective(activations, coefficients, concept_bank, alpha_w, l1_ratio): + """Compute the fixed-dictionary NMF encoding objective.""" + n_features = activations.shape[1] + l1_reg = n_features * alpha_w * l1_ratio + l2_reg = n_features * alpha_w * (1.0 - l1_ratio) + + reconstruction = coefficients @ concept_bank + data_fit = 0.5 * np.square(activations - reconstruction).sum() + l1_penalty = l1_reg * np.abs(coefficients).sum() + l2_penalty = 0.5 * l2_reg * np.square(coefficients).sum() + + return data_fit + l1_penalty + l2_penalty + + +@pytest.mark.parametrize("alpha_w,l1_ratio", [(0.0, 0.0), (0.05, 0.3)]) +@pytest.mark.parametrize("tol", [1e-6, 0.0]) +def test_torch_encode_differentiable_matches_sklearn_transform(alpha_w, l1_ratio, tol): + """The differentiable Torch solver should match sklearn's fixed-dictionary transform.""" + train = _make_non_negative_data(0, (80, 12)) + test = _make_non_negative_data(1, (20, 12)) + + factorizer = TorchSklearnNMFFactorizer( + n_components=5, + alpha_W=alpha_w, + l1_ratio=l1_ratio, + max_iter=1000, + tol=tol, + random_state=0, + ) + factorizer.fit(train) + + expected = factorizer.encode(test) + actual = factorizer.encode_differentiable(torch.tensor(test)) + actual_np = actual.detach().cpu().numpy() + concept_bank = factorizer.get_concept_bank() + + assert isinstance(actual, torch.Tensor) + assert actual.shape == expected.shape + assert np.all(actual_np >= -1e-6) + + expected_objective = _nmf_encoding_objective( + test, expected, concept_bank, alpha_w=alpha_w, l1_ratio=l1_ratio + ) + actual_objective = _nmf_encoding_objective( + test, actual_np, concept_bank, alpha_w=alpha_w, l1_ratio=l1_ratio + ) + assert actual_objective <= expected_objective * 1.05 + 1e-5 + + expected_reconstruction = expected @ concept_bank + actual_reconstruction = actual_np @ concept_bank + reconstruction_error = np.linalg.norm(actual_reconstruction - expected_reconstruction) + reconstruction_norm = np.linalg.norm(expected_reconstruction) + assert reconstruction_error / (reconstruction_norm + 1e-8) < 0.05 + + +@pytest.mark.parametrize("tol", [1e-6, 0.0]) +def test_torch_encode_differentiable_preserves_gradients(tol): + """The differentiable Torch solver should backpropagate through activations.""" + train = _make_non_negative_data(2, (60, 10)) + test = torch.tensor(_make_non_negative_data(3, (16, 10)), requires_grad=True) + + factorizer = TorchSklearnNMFFactorizer( + n_components=4, + alpha_W=1e-2, + max_iter=500, + tol=tol, + random_state=0, + ) + factorizer.fit(train) + + coefficients = factorizer.encode_differentiable(test) + loss = torch.sum(coefficients) + loss.backward() + + gradients = test.grad + + assert gradients is not None + assert torch.all(torch.isfinite(gradients)).item() + assert torch.sum(torch.abs(gradients)).item() > 0 + + +def test_torch_encode_differentiable_rejects_negative_activations(): + """NMF encoding should fail when differentiable inputs are negative.""" + train = _make_non_negative_data(6, (30, 6), minimum=0.1) + test = torch.tensor(-_make_non_negative_data(7, (10, 6), minimum=0.1)) + + factorizer = TorchSklearnNMFFactorizer(n_components=3, max_iter=200, random_state=0) + factorizer.fit(train) + + with pytest.raises(ValueError, match="NMF requires non-negative activations"): + factorizer.encode_differentiable(test) + + +def test_torch_encode_differentiable_rejects_unsupported_beta_loss(): + """Unsupported sklearn NMF configurations should fail explicitly.""" + train = _make_non_negative_data(4, (40, 8), minimum=0.1) + test = torch.tensor(_make_non_negative_data(5, (12, 8), minimum=0.1)) + + factorizer = TorchSklearnNMFFactorizer( + n_components=3, + solver="mu", + beta_loss="kullback-leibler", + max_iter=300, + random_state=0, + ) + factorizer.fit(train) + + with pytest.raises(NotImplementedError): + factorizer.encode_differentiable(test) diff --git a/tests/concepts/test_holistic_craft_classification_tf.py b/tests/concepts/test_holistic_craft_classification_tf.py new file mode 100644 index 00000000..c1db481b --- /dev/null +++ b/tests/concepts/test_holistic_craft_classification_tf.py @@ -0,0 +1,342 @@ +"""Tests for HolisticCraftTf on classification tasks with TensorFlow.""" +# pylint: disable=redefined-outer-name + +import numpy as np +import pytest +import tensorflow as tf +from PIL import Image + +import xplique +from xplique.attributions import Saliency +from xplique.attributions.gradient_input import GradientInput +from xplique.concepts import HolisticCraftTf as Craft +from xplique.concepts.holistic_craft import PartialExplainer +from xplique.concepts.tf.layered_model_latent_extractor import LayeredModelExtractorBuilder +from xplique.utils_functions.classification.tf.classifier_tensor import TfClassifierTensor +from xplique.utils_functions.common.tf.gradients_check import check_model_gradients + + +def test_classifier_tensor_targets_class_and_preserves_batch_shape(): + predictions = TfClassifierTensor(tf.constant([[0.1, 0.2, 0.7], [0.3, 0.6, 0.1]])) + + targets = predictions.to_attribution_target(class_id=1) + + np.testing.assert_array_equal(targets.tensor.numpy(), [[0.0, 1.0, 0.0], [0.0, 1.0, 0.0]]) + assert targets.to_batched_tensor().shape == (2, 3) + + with pytest.raises(ValueError): + predictions.to_attribution_target(class_id=3) + + +@pytest.fixture(params=["cpu", "gpu"]) +def device_param(request): + """Pytest fixture to provide device parameter (cpu or gpu).""" + if request.param == "gpu" and not tf.config.list_physical_devices('GPU'): + pytest.skip("GPU not available") + return request.param + + +@pytest.fixture(scope="function") +def image_data(device_param): + """Pytest fixture to create a fake image for testing.""" + device_name = f'/{device_param.upper()}:0' + with tf.device(device_name): + rng = np.random.default_rng(seed=42) + raw_image = Image.fromarray(rng.integers(0, 256, (462, 640, 3), dtype=np.uint8)) + + # Resize to 224x224 for ImageNet models + resized_image = raw_image.resize((224, 224), Image.Resampling.LANCZOS) + + # Convert to numpy array and normalize + img_np = np.array(resized_image, dtype=np.float32) + + # Standard ImageNet normalization + mean = np.array([0.485, 0.456, 0.406]) * 255.0 + std = np.array([0.229, 0.224, 0.225]) * 255.0 + img_np = (img_np - mean) / std + + # Add batch dimension + input_tensor = tf.expand_dims(img_np, axis=0) + + return raw_image, input_tensor + + +def test_image_size(image_data): + """Test that loaded image has expected dimensions.""" + image, _ = image_data + expected_size = (640, 462) + assert image.size == expected_size + + +@pytest.fixture(scope="function") +def model_data(image_data, device_param): + """Pytest fixture to load a local ResNet50 model and run predictions.""" + device_name = f'/{device_param.upper()}:0' + with tf.device(device_name): + _, input_tensor = image_data + # Do not download ImageNet weights during tests. + model = tf.keras.applications.ResNet50(weights=None) + predictions = model.predict(input_tensor, verbose=0) + + return model, predictions + + +def test_model_outputs(model_data): + """Test that model outputs have expected shape for ImageNet classes.""" + _, predictions = model_data + # ResNet50 outputs 1000 ImageNet classes + assert predictions.shape == (1, 1000), f"Expected shape (1, 1000), got {predictions.shape}" + + +@pytest.fixture(scope="module") +def imagenet_classes(): + """Pytest fixture providing ImageNet class labels for testing.""" + # Top-5 most common ImageNet classes for testing + # In a real scenario, you would load all 1000 classes + classes = [ + 'tench', 'goldfish', 'great_white_shark', 'tiger_shark', 'hammerhead', + 'electric_ray', 'stingray', 'cock', 'hen', 'ostrich' + ] + return classes + + +@pytest.fixture(scope="function") +def latent_extractor_data(model_data, device_param): + """Pytest fixture to create latent extractor from ResNet50 model.""" + device_name = f'/{device_param.upper()}:0' + with tf.device(device_name): + model, _ = model_data + + # Split ResNet50 at layer -3 (before GlobalAveragePooling2D) + # This preserves spatial dimensions needed for CRAFT + # ResNet structure: ... -> conv5_block3_out (7x7x2048) + # -> avg_pool (2048) -> predictions (1000) + latent_extractor = LayeredModelExtractorBuilder.build( + model=model, + split_layer=-3, # Split before avg_pool to preserve spatial dimensions + batch_size=1 + ) + return latent_extractor + + +def test_latent_extractor(image_data, latent_extractor_data): + """Test that latent extractor returns ClassifierTensor with correct shape.""" + _, input_tensor = image_data + latent_extractor = latent_extractor_data + + # Test the latent extractor output + results = latent_extractor(input_tensor) + + # Should return TfClassifierTensor with shape (batch, num_classes) + assert isinstance(results, TfClassifierTensor), "Results should be a TfClassifierTensor" + assert results.shape == (1, 1000), f"Expected shape (1, 1000), got {results.shape}" + + +def test_layered_latent_data_getitem(): + """Test that LayeredLatentData supports integer and slice indexing.""" + from xplique.concepts.tf.layered_model_latent_extractor import LayeredLatentData + + activations = tf.random.normal((4, 7, 7, 16)) + latent_data = LayeredLatentData(activations) + + sliced = latent_data[1:3] + assert isinstance(sliced, LayeredLatentData) + assert sliced.activations.shape[0] == 2 + + item = latent_data[0] + assert isinstance(item, LayeredLatentData) + + +def test_latent_extractor_gradients(image_data, latent_extractor_data): + """Test that gradients flow correctly through the latent extractor.""" + _, input_tensor = image_data + latent_extractor = latent_extractor_data + + # Use check_model_gradients to verify gradient flow + check = check_model_gradients(latent_extractor, input_tensor) + assert check, "Latent extractor gradients should be computed successfully." + + +def test_latent_extractor_saliency(image_data, latent_extractor_data): + """Test saliency attribution method on latent extractor.""" + image, input_tensor = image_data + latent_extractor = latent_extractor_data + + # Get predictions + predictions = latent_extractor(input_tensor) + + # Get top prediction + top_class = tf.argmax(predictions, axis=1).numpy()[0] + + # Create explainer + explainer = Saliency(latent_extractor, operator=xplique.Tasks.CLASSIFICATION, batch_size=1) + explanation = explainer.explain(input_tensor, targets=np.array([top_class])) + + # Check explanation shape + expected_shape = (1, 224, 224, 1) + assert explanation.shape == expected_shape, ( + f"Expected shape {expected_shape}, got {explanation.shape}" + ) + + +@pytest.fixture(scope="function") +def craft_data(image_data, latent_extractor_data, device_param): + """Pytest fixture to create and fit CRAFT instance.""" + device_name = f'/{device_param.upper()}:0' + with tf.device(device_name): + _, input_tensor = image_data + latent_extractor = latent_extractor_data + + # Create CRAFT instance + craft = Craft( + latent_extractor=latent_extractor, + number_of_concepts=10 + ) + + # Fit on the input (in real scenario, use multiple images) + craft.fit(input_tensor) + + return craft + + +def test_craft_reencode(image_data, craft_data): + """Test CRAFT encode and decode operations.""" + _, input_tensor = image_data + craft = craft_data + + # Encode the input + encoded_data = craft.encode(input_tensor) + + # Should have one tuple per image in the batch + expected_len = input_tensor.shape[0] + assert len(encoded_data) == expected_len, ( + f"Expected {expected_len} tuples, got {len(encoded_data)}" + ) + + # Get the first tuple + latent_data, coeffs_u = encoded_data[0] + + # For ResNet with GlobalAveragePooling output, coeffs_u should have spatial dimensions + assert len(coeffs_u.shape) == 4, "coeffs_u should be 4D (batch, height, width, concepts)" + assert coeffs_u.shape[0] == 1, "Batch dimension should be 1" + assert coeffs_u.shape[3] == 10, "Should have 10 concepts" + + # Decode back + result = craft.decode(latent_data, coeffs_u) + assert isinstance(result, TfClassifierTensor), "Decoded result should be a TfClassifierTensor" + assert result.shape == (1, 1000), f"Expected shape (1, 1000), got {result.shape}" + + +def test_craft_decoder_modes(image_data, craft_data): + """Test CRAFT concept decoder functionality.""" + _, input_tensor = image_data + craft = craft_data + + # Encode + encoded_data = craft.encode(input_tensor) + latent_data, coeffs_u = encoded_data[0] + + # Create decoder + decoder = craft.make_concept_decoder(latent_data) + + # Decoder should return tensor + output_tensor = decoder(coeffs_u) + assert hasattr(output_tensor, 'shape'), "Decoder should return a tensor" + assert output_tensor.shape == (1, 1000), f"Expected shape (1, 1000), got {output_tensor.shape}" + + +def test_craft_gradient_input(image_data, craft_data): + """Test CRAFT gradient-based importance estimation.""" + _, input_tensor = image_data + craft = craft_data + + # Use a specific class for testing (e.g., class 281 is 'tabby cat') + class_id = 281 + + # Test compute_explanation_per_concept + operator = xplique.Tasks.CLASSIFICATION + partial_explainer = PartialExplainer( + GradientInput, + operator=operator, + reducer=None, + ) + explanation = craft.compute_explanation_per_concept( + input_tensor, class_id=class_id, partial_explainer=partial_explainer + ) + + # Verify explanation shape + assert explanation.shape[0] == 1, "Should have one explanation per image" + assert explanation.shape[3] == 10, "Should match number of concepts" + + # Test estimate_importance + importances_gi = craft.estimate_importance( + input_tensor, operator, class_id, method='gradient_input' + ) + + # Verify importance scores + assert importances_gi.shape == (10,), f"Expected shape (10,), got {importances_gi.shape}" + assert np.all(np.isfinite(importances_gi)), "All importances should be finite" + + order = importances_gi.argsort()[::-1] + assert order.shape == (10,), "Should have ordering for all concepts" + assert len(np.unique(order)) == 10, "All concepts should have unique ordering" + + +def test_craft_encode_differentiable_gradients(image_data, craft_data): + """Test that encode(differentiable=True) preserves gradients.""" + _, input_tensor = image_data + craft = craft_data + + # Test differentiable encoding with gradient tape + with tf.GradientTape() as tape: + tape.watch(input_tensor) + + # Encode with differentiable mode + encoded_data = craft.encode(input_tensor, differentiable=True) + _, coeffs_u = encoded_data[0] + + # Verify coeffs_u is a tensor + assert isinstance(coeffs_u, tf.Tensor), ( + "coeffs_u should be a tf.Tensor in differentiable mode" + ) + assert tf.reduce_all(tf.math.is_finite(coeffs_u)), "coeffs_u should be finite" + assert tf.reduce_min(coeffs_u) >= -1e-6, "coeffs_u should remain non-negative" + + # Create a simple loss + loss = tf.reduce_sum(coeffs_u) + + # Check gradient flow + gradients = tape.gradient(loss, input_tensor) + + # Verify gradients flowed back to input + assert gradients is not None, "Gradients should flow back to input" + assert tf.reduce_sum(tf.abs(gradients)) > 0, "Gradients should be non-zero" + + +def test_craft_sobol_importance(image_data, craft_data): + """Test Sobol importance estimation for concepts.""" + _, input_tensor = image_data + craft = craft_data + + # Use a specific class for testing + class_id = 281 # 'tabby cat' + operator = xplique.Tasks.CLASSIFICATION + + # Estimate importance using Sobol method (this may take longer) + importances_sobol = craft.estimate_importance( + input_tensor, + operator, + class_id, + method='sobol', + grid_size=4, # Reduced for faster testing + nb_design=4 # Reduced for faster testing + ) + + # Verify importance scores + assert importances_sobol.shape == (10,), f"Expected shape (10,), got {importances_sobol.shape}" + assert np.all(np.isfinite(importances_sobol)), "All importances should be finite" + assert np.all(importances_sobol >= 0), "Sobol importances should be non-negative" + + order = importances_sobol.argsort()[::-1] + assert order.shape == (10,), "Should have ordering for all concepts" + assert len(np.unique(order)) == 10, "All concepts should have unique ordering" diff --git a/tests/concepts/test_holistic_craft_classification_torch.py b/tests/concepts/test_holistic_craft_classification_torch.py new file mode 100644 index 00000000..495a29e5 --- /dev/null +++ b/tests/concepts/test_holistic_craft_classification_torch.py @@ -0,0 +1,358 @@ +"""Tests for HolisticCraftTorch on classification tasks with PyTorch.""" +# pylint: disable=redefined-outer-name + +import numpy as np +import pytest +import torch +import torchvision.transforms as T +from PIL import Image +from torchvision import models + +import xplique +from xplique.attributions import Saliency +from xplique.attributions.gradient_input import GradientInput +from xplique.concepts import HolisticCraftTorch as Craft +from xplique.concepts.holistic_craft import PartialExplainer +from xplique.concepts.torch.layered_model_latent_extractor import LayeredModelExtractorBuilder +from xplique.utils_functions.classification.torch.classifier_tensor import TorchClassifierTensor +from xplique.utils_functions.common.torch.gradients_check import check_model_gradients +from xplique.wrappers import TorchWrapper + + +def test_classifier_tensor_targets_class_and_preserves_batch_shape(): + predictions = TorchClassifierTensor.from_predictions( + torch.tensor([[0.1, 0.2, 0.7], [0.3, 0.6, 0.1]]) + ) + + targets = predictions.to_attribution_target(class_id=1) + + torch.testing.assert_close(targets, torch.tensor([[0.0, 1.0, 0.0], [0.0, 1.0, 0.0]])) + assert targets.to_batched_tensor().shape == (2, 3) + + with pytest.raises(ValueError): + predictions.to_attribution_target(class_id=3) + + +@pytest.fixture(params=["cpu", "cuda"]) +def device_param(request): + """Pytest fixture to provide device parameter (cpu or cuda).""" + device_str = request.param + if device_str == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + device = torch.device(device_str) + return device + + +@pytest.fixture(scope="function") +def image_data(device_param): + """Pytest fixture to create a fake image for testing.""" + rng = np.random.default_rng(seed=42) + raw_image = Image.fromarray(rng.integers(0, 256, (462, 640, 3), dtype=np.uint8)) + + # Standard ImageNet normalization + transform = T.Compose([ + T.Resize((224, 224)), + T.ToTensor(), + T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) + ]) + + # Preprocess and batch the image + input_tensor = transform(raw_image).unsqueeze(0).to(device_param) + + return raw_image, input_tensor + + +def test_image_size(image_data): + """Test that loaded image has expected dimensions.""" + image, _ = image_data + expected_size = (640, 462) + assert image.size == expected_size + + +@pytest.fixture(scope="function") +def model_data(image_data, device_param): + """Pytest fixture to load a local ResNet50 model and run predictions.""" + _, input_tensor = image_data + # Do not download ImageNet weights during tests. + model = models.resnet50(weights=None).to(device_param) + model.eval() + + with torch.no_grad(): + predictions = model(input_tensor) + + return model, predictions + + +def test_model_outputs(model_data): + """Test that model outputs have expected shape for ImageNet classes.""" + _, predictions = model_data + # ResNet50 outputs 1000 ImageNet classes + assert predictions.shape == (1, 1000), f"Expected shape (1, 1000), got {predictions.shape}" + + +@pytest.fixture(scope="module") +def imagenet_classes(): + """Pytest fixture providing ImageNet class labels for testing.""" + # Top-5 most common ImageNet classes for testing + # In a real scenario, you would load all 1000 classes + classes = [ + 'tench', 'goldfish', 'great_white_shark', 'tiger_shark', 'hammerhead', + 'electric_ray', 'stingray', 'cock', 'hen', 'ostrich' + ] + return classes + + +@pytest.fixture(scope="function") +def latent_extractor_data(model_data, device_param): + """Pytest fixture to create latent extractor from ResNet50 model.""" + model, _ = model_data + + # Split ResNet50 at layer 4 (before the final classifier layer) + # ResNet structure: conv1, bn1, relu, maxpool, layer1, layer2, layer3, layer4, avgpool, fc + latent_extractor = LayeredModelExtractorBuilder.build( + model=model, + split_layer=8, # Split at avgpool (index 8), before fc layer + device=str(device_param), + batch_size=1 + ) + return latent_extractor + + +def test_latent_extractor(image_data, latent_extractor_data): + """Test that latent extractor returns ClassifierTensor with correct shape.""" + _, input_tensor = image_data + latent_extractor = latent_extractor_data + + # Test the latent extractor output + with torch.no_grad(): + results = latent_extractor(input_tensor) + + # Should return TorchClassifierTensor with shape (batch, num_classes) + assert isinstance(results, TorchClassifierTensor), "Results should be a TorchClassifierTensor" + assert results.shape == (1, 1000), f"Expected shape (1, 1000), got {results.shape}" + + +def test_layered_latent_data_getitem(device_param): + """Test that LayeredLatentData supports integer and slice indexing.""" + from xplique.concepts.torch.layered_model_latent_extractor import LayeredLatentData + + activations = torch.randn(4, 16, 7, 7).to(device_param) + latent_data = LayeredLatentData(activations) + + sliced = latent_data[1:3] + assert isinstance(sliced, LayeredLatentData) + assert sliced.activations.shape[0] == 2 + + item = latent_data[0] + assert isinstance(item, LayeredLatentData) + + +def test_latent_extractor_gradients(image_data, latent_extractor_data): + """Test that gradients flow correctly through the latent extractor.""" + _, input_tensor = image_data + latent_extractor = latent_extractor_data + + # Use check_model_gradients to verify gradient flow + check = check_model_gradients(latent_extractor, input_tensor) + assert check, "Latent extractor gradients should be computed successfully." + + +def test_latent_extractor_saliency(image_data, latent_extractor_data, device_param): + """Test saliency attribution method on latent extractor.""" + _, input_tensor = image_data + latent_extractor = latent_extractor_data + + # Get predictions + with torch.no_grad(): + predictions = latent_extractor(input_tensor) + + # Get top prediction + top_class = predictions.argmax(dim=1).item() + + # Convert input for TensorFlow format (Xplique expects channels last) + input_tensor_tf_dim = input_tensor.detach().cpu().numpy().transpose(0, 2, 3, 1) + + # Wrap model for Xplique + torch_wrapped_model = TorchWrapper( + latent_extractor, device=device_param, is_channel_first=True + ) + + # Create explainer + explainer = Saliency( + torch_wrapped_model, operator=xplique.Tasks.CLASSIFICATION, batch_size=1 + ) + explanation = explainer.explain(input_tensor_tf_dim, targets=np.array([top_class])) + + # Check explanation shape + expected_shape = (1, 224, 224, 1) + assert explanation.shape == expected_shape, ( + f"Expected shape {expected_shape}, got {explanation.shape}" + ) + + +@pytest.fixture(scope="function") +def craft_data(image_data, latent_extractor_data, device_param): + """Pytest fixture to create and fit CRAFT instance.""" + _, input_tensor = image_data + latent_extractor = latent_extractor_data + + # Create CRAFT instance + craft = Craft( + latent_extractor=latent_extractor, + number_of_concepts=10, + device=str(device_param) + ) + + # Fit on the input (in real scenario, use multiple images) + craft.fit(input_tensor) + + return craft + + +def test_craft_reencode(image_data, craft_data): + """Test CRAFT encode and decode operations.""" + _, input_tensor = image_data + craft = craft_data + + # Encode the input + encoded_data = craft.encode(input_tensor) + + # Should have one tuple per image in the batch + expected_tuples = input_tensor.shape[0] + assert len(encoded_data) == expected_tuples, ( + f"Expected {expected_tuples} tuples, got {len(encoded_data)}" + ) + + # Get the first tuple + latent_data, coeffs_u = encoded_data[0] + + # For ResNet with avgpool output, coeffs_u should have spatial dimensions + assert len(coeffs_u.shape) == 4, "coeffs_u should be 4D (batch, height, width, concepts)" + assert coeffs_u.shape[0] == 1, "Batch dimension should be 1" + assert coeffs_u.shape[3] == 10, "Should have 10 concepts" + + # Decode back + result = craft.decode(latent_data, coeffs_u) + assert isinstance(result, TorchClassifierTensor), "Result should be a TorchClassifierTensor" + assert result.shape == (1, 1000), f"Expected shape (1, 1000), got {result.shape}" + + +def test_craft_decoder_modes(image_data, craft_data): + """Test CRAFT concept decoder functionality.""" + _, input_tensor = image_data + craft = craft_data + + # Encode + encoded_data = craft.encode(input_tensor) + latent_data, coeffs_u = encoded_data[0] + + # Create decoder + decoder = craft.make_concept_decoder(latent_data) + + # Decoder should return tensor + output_tensor = decoder(coeffs_u) + assert hasattr(output_tensor, 'shape'), "Decoder should return a tensor" + assert output_tensor.shape == (1, 1000), f"Expected shape (1, 1000), got {output_tensor.shape}" + + +def test_craft_gradient_input(image_data, craft_data): + """Test CRAFT gradient-based importance estimation.""" + _, input_tensor = image_data + craft = craft_data + + # Use a specific class for testing (e.g., class 281 is 'tabby cat') + class_id = 281 + + # Test compute_explanation_per_concept + operator = xplique.Tasks.CLASSIFICATION + partial_explainer = PartialExplainer( + GradientInput, + operator=operator, + reducer=None, + ) + explanation = craft.compute_explanation_per_concept( + input_tensor, class_id=class_id, partial_explainer=partial_explainer + ) + + # Verify explanation shape + assert explanation.shape[0] == 1, "Should have one explanation per image" + assert explanation.shape[3] == 10, "Should match number of concepts" + + # Test estimate_importance + importances_gi = craft.estimate_importance( + input_tensor, operator, class_id, method='gradient_input' + ) + + # Verify importance scores + assert importances_gi.shape == (10,), ( + f"Expected shape (10,), got {importances_gi.shape}" + ) + assert np.all(np.isfinite(importances_gi)), "All importances should be finite" + + order = importances_gi.argsort()[::-1] + assert order.shape == (10,), "Should have ordering for all concepts" + assert len(np.unique(order)) == 10, "All concepts should have unique ordering" + + +def test_craft_encode_differentiable_gradients(image_data, craft_data): + """Test that encode(differentiable=True) preserves gradients.""" + _, input_tensor = image_data + craft = craft_data + + # Ensure input has gradients enabled + input_with_grad = input_tensor.clone().detach().requires_grad_(True) + + # Test differentiable encoding + encoded_data = craft.encode(input_with_grad, differentiable=True) + _, coeffs_u = encoded_data[0] + + # Verify coeffs_u is a tensor with gradients + expected_msg = "coeffs_u should be a torch.Tensor in differentiable mode" + assert isinstance(coeffs_u, torch.Tensor), expected_msg + assert torch.all(torch.isfinite(coeffs_u)).item(), "coeffs_u should be finite" + assert torch.min(coeffs_u).item() >= -1e-6, "coeffs_u should remain non-negative" + assert coeffs_u.requires_grad, "coeffs_u should have gradients enabled" + + # Create a simple loss and check gradient flow + loss = coeffs_u.sum() + loss.backward() + gradients = input_with_grad.grad + + # Verify gradients flowed back to input + assert gradients is not None, "Gradients should flow back to input" + assert gradients.abs().sum() > 0, "Gradients should be non-zero" + + +def test_craft_sobol_importance(image_data, craft_data): + """Test Sobol importance estimation for concepts.""" + _, input_tensor = image_data + craft = craft_data + + # Use a specific class for testing + class_id = 281 # 'tabby cat' + operator = xplique.Tasks.CLASSIFICATION + + # Estimate importance using Sobol method (this may take longer) + importances_sobol = craft.estimate_importance( + input_tensor, + operator, + class_id, + method='sobol', + grid_size=4, # Reduced for faster testing + nb_design=4 # Reduced for faster testing + ) + + # Verify importance scores + assert importances_sobol.shape == (10,), ( + f"Expected shape (10,), got {importances_sobol.shape}" + ) + assert np.all(np.isfinite(importances_sobol)), ( + "All importances should be finite" + ) + assert np.all(importances_sobol >= 0), ( + "Sobol importances should be non-negative" + ) + + order = importances_sobol.argsort()[::-1] + assert order.shape == (10,), "Should have ordering for all concepts" + assert len(np.unique(order)) == 10, "All concepts should have unique ordering" diff --git a/tests/concepts/test_holistic_craft_object_detection_tf.py b/tests/concepts/test_holistic_craft_object_detection_tf.py new file mode 100644 index 00000000..3ec0f2c0 --- /dev/null +++ b/tests/concepts/test_holistic_craft_object_detection_tf.py @@ -0,0 +1,671 @@ +"""Tests for HolisticCraftTf on object detection tasks with TensorFlow.""" +# pylint: disable=redefined-outer-name +from typing import List + +import numpy as np +import pytest +import tensorflow as tf +from PIL import Image + +import xplique +from xplique.attributions import Saliency +from xplique.attributions.gradient_input import GradientInput +from xplique.concepts import HolisticCraftTf as Craft +from xplique.concepts.holistic_craft import PartialExplainer +from xplique.concepts.latent_extractor import LatentData +from xplique.concepts.tf.latent_extractor import TfLatentExtractor +from xplique.utils_functions.common.tf.gradients_check import check_model_gradients +from xplique.utils_functions.object_detection.base.box_manager import BoxFormat, BoxType +from xplique.utils_functions.object_detection.tf.box_formatter import TfBaseBoxFormatter +from xplique.utils_functions.object_detection.tf.box_model_wrapper import TfBoxesModelWrapper +from xplique.utils_functions.object_detection.tf.multi_box_tensor import ( + TfMultiBoxTensor as MultiBoxTensor, +) + +# ============================================================================ +# Mock Classes - Shared across tests +# ============================================================================ + +class MockRetinaNetModel(tf.keras.Model): + """Mock RetinaNet model that returns predictions with boxes, confidence, and classes.""" + def __init__(self, nb_classes=20): + super().__init__() + self.nb_classes = nb_classes + + def call(self, x): + """Forward pass generating mock predictions based on input.""" + batch_size = tf.shape(x)[0] + + # Use input features to create predictions - ensures gradient flow + # Pool the input to a scalar per batch and use it to modulate predictions + pooled = tf.reduce_mean(x, axis=[1, 2, 3]) # (batch_size,) + pooled = tf.abs(pooled) + 1.0 # Ensure positive, add 1.0 for decent scale + pooled_expanded = tf.expand_dims(pooled, axis=1) # (batch_size, 1) + pooled_boxes = tf.expand_dims(pooled_expanded, axis=-1) # (batch_size, 1, 1) + + # Create deterministic predictions with at least one high-confidence car detection + # Car class index = 6 in PASCAL VOC + # Define a single batch of detections + single_batch_boxes = [ + [0.3, 0.3, 0.7, 0.7], # Box 0: high-confidence car + [0.1, 0.1, 0.3, 0.3], # Box 1: person + [0.5, 0.5, 0.9, 0.9], # Box 2: aeroplane + [0.2, 0.4, 0.5, 0.8], # Box 3: bird + [0.0, 0.0, 0.1, 0.1], # Box 4-9: low confidence + [0.0, 0.0, 0.1, 0.1], + [0.0, 0.0, 0.1, 0.1], + [0.0, 0.0, 0.1, 0.1], + [0.0, 0.0, 0.1, 0.1], + [0.0, 0.0, 0.1, 0.1], + ] + + single_batch_confidences = [ + 0.95, # High confidence for car + 0.85, # person + 0.75, # aeroplane + 0.6, # bird + 0.1, 0.1, 0.1, 0.1, 0.1, 0.1 # Low confidence for others + ] + + single_batch_classes = [ + 6, # car + 14, # person + 0, # aeroplane + 2, # bird + 0, 0, 0, 0, 0, 0 # Low confidence detections + ] + + # Create base predictions + boxes_base = tf.constant([single_batch_boxes], dtype=tf.float32) # (1, 10, 4) + confidences_base = tf.constant([single_batch_confidences], dtype=tf.float32) # (1, 10) + + # Create class probabilities (nb_classes=20) instead of class indices + # This ensures gradients flow through class predictions + nb_classes = 20 + class_probas_list = [] + for cls_idx in single_batch_classes: + # Create one-hot-like probabilities: 0.99 for target class, 0.01 for others + # Build list directly without using .numpy() + probas_row = [0.99 if i == cls_idx else 0.01 for i in range(nb_classes)] + class_probas_list.append(probas_row) + class_probas_base = tf.constant([class_probas_list], dtype=tf.float32) # (1, 10, 20) + + # Tile to match batch_size + boxes = tf.tile(boxes_base, [batch_size, 1, 1]) # (batch_size, 10, 4) + confidences = tf.tile(confidences_base, [batch_size, 1]) # (batch_size, 10) + class_probas = tf.tile(class_probas_base, [batch_size, 1, 1]) # (batch_size, 10, 20) + + # Modulate by input features to create gradient path + # Use stronger modulation (0.1 instead of 0.001) for better saliency gradients + pooled_mod = (pooled_expanded - 1.0) * 0.1 + 1.0 # Near 1.0 but with visible gradient + confidences = confidences * pooled_mod + + # Modulate class probabilities by input for gradient flow + # Expand pooled_mod for broadcasting: (batch_size, 1) -> (batch_size, 10, 1) + pooled_mod_3d = tf.expand_dims(pooled_mod, axis=-1) + class_probas = class_probas * pooled_mod_3d + + # Also modulate boxes slightly for stronger gradients + boxes = boxes * pooled_boxes * 0.05 + boxes * 0.95 # 5% modulation + + return { + 'boxes': boxes, + 'confidence': confidences, + 'class_probas': class_probas # Return probabilities instead of class indices + } + +class MockRetinaNetFormatter(TfBaseBoxFormatter): + """ + Mock formatter for RetinaNet model.call() outputs. + + Closely mimics RetinaNetProcessedBoxFormatter from xplique-adapters. + Handles RetinaNet-specific output format with boxes and classification logits, + converting them to unified Xplique format. + """ + def __init__(self, nb_classes: int = 20, + input_box_type: BoxType = BoxType(BoxFormat.XYWH, is_normalized=False), + output_box_type: BoxType = BoxType(BoxFormat.XYXY, is_normalized=True), + image_size: tuple = None) -> None: + """ + Initialize the mock RetinaNet box formatter. + + Parameters + ---------- + nb_classes + Number of object classes in the dataset. + input_box_type + Format of boxes from RetinaNet (default XYWH, unnormalized). + output_box_type + Desired output format (default XYXY, normalized). + image_size + Size of image for coordinate conversion. + **kwargs + Additional arguments passed to parent class. + """ + super().__init__(input_box_type, output_box_type) + + self.nb_classes = nb_classes + self.image_size = image_size + + def __call__(self, model_outputs): + """Make formatter callable.""" + return self.forward(model_outputs) + + def forward(self, predictions) -> List[MultiBoxTensor]: + """ + Process RetinaNet predictions handling both single and multi-batch inputs. + + Parameters + ---------- + predictions + Dictionary with 'boxes', 'confidence', 'class_probas' keys + + Returns + ------- + formatted_predictions + List of MultiBoxTensor objects, one per image in the batch. + """ + def process_single_batch(batch_idx): + boxes = predictions['boxes'][batch_idx] + scores = predictions['confidence'][batch_idx] + class_probas = predictions['class_probas'][batch_idx] + + scores_expanded = scores[:, tf.newaxis] + + pred_dict = { + 'boxes': boxes, + 'scores': scores_expanded, + 'probas': class_probas # Use class_probas directly (already has gradients) + } + return self.format_predictions(pred_dict, self.image_size) + + batch_size = predictions['boxes'].shape[0] + + results = [] + for batch_idx in range(batch_size): + formatted = process_single_batch(batch_idx) + results.append(formatted) + return results + +class MockTfLatentData(LatentData): + """Mock latent data container for TensorFlow, similar to PyTorch version.""" + def __init__(self, activations: np.ndarray, batch_size=1): + self.batch_size = batch_size + self.activations = activations + + def set_activations(self, values: np.ndarray): + """Set activations.""" + self.activations = values + + def get_activations(self, as_numpy: bool = False, keep_gradients: bool = False): + """Retrieve activations.""" + if as_numpy: + return self.activations.numpy() + return self.activations + +class MockExtractorBuilder: + """ + Builder for creating TfLatentExtractor instances for mock RetinaNet-style models. + + This class encapsulates the logic for creating a TfLatentExtractor configured + for mock object detection models that follow the RetinaNet output format. + + Similar to the PyTorch MockExtractorBuilder, this provides a clean interface + for constructing the latent extractor with all necessary components. + """ + + @classmethod + def build(cls, nb_classes, image_size, device='cpu', batch_size=1): + """ + Build a TfLatentExtractor for a mock RetinaNet-style model. + + Creates custom input_to_latent and latent_to_logit functions that split + the model's forward pass into feature extraction and prediction. + + Parameters + ---------- + nb_classes : int + Number of object classes in the dataset. + image_size : tuple + Image dimensions as (height, width) for coordinate conversion. + device : str + Device to run computations on. Default is 'cpu'. + Should be 'cpu' or 'gpu'. + batch_size : int + Batch size for processing. Default is 1. + + Returns + ------- + latent_extractor : TfLatentExtractor + Configured TfLatentExtractor instance for the mock model. + """ + device_name = f'/{device.upper()}:0' + with tf.device(device_name): + # Create mock model (not real RetinaNet) + mock_model = MockRetinaNetModel(nb_classes=nb_classes) + + def input_to_latent_fn(inputs): + """Simple mock latent extractor - just pool input to fixed size.""" + batch_size_actual = tf.shape(inputs)[0] + + # Pool to fixed spatial size (10x10) and apply ReLU for NMF + pooled = tf.image.resize(inputs, (10, 10), method='bilinear') # (N, 10, 10, C) + activations = tf.nn.relu(pooled) # Apply ReLU for NMF + + # Pad/truncate channels to 64 for consistency + current_channels = tf.shape(activations)[-1] + if current_channels != 64: + padding = [[0, 0], [0, 0], [0, 0], [0, tf.maximum(0, 64 - current_channels)]] + activations = tf.pad(activations, padding) + activations = activations[:, :, :, :64] + + batch_size_value = ( + batch_size_actual.numpy() + if hasattr(batch_size_actual, 'numpy') + else batch_size_actual + ) + return MockTfLatentData(activations=activations, batch_size=batch_size_value) + + def latent_to_logit_fn(latent_data): + """Simple mock decoder - reuse the existing MockRetinaNetModel logic.""" + activations = latent_data.get_activations(as_numpy=False, keep_gradients=True) + # Upsample back to expected size and pass through model + # The model expects (N, H, W, C), we have (N, 10, 10, 64) + return mock_model(activations) + + # Create formatter with proper image sizes + formatter = MockRetinaNetFormatter( + nb_classes, + image_size=image_size + ) + + # Create TfLatentExtractor with simplified mock functions + latent_extractor = TfLatentExtractor( + model=mock_model, + input_to_latent_model=input_to_latent_fn, + latent_to_logit_model=latent_to_logit_fn, + latent_data_class=MockTfLatentData, + output_formatter=formatter, + batch_size=batch_size + ) + + return latent_extractor + + +# ============================================================================ +# Test Fixtures +# ============================================================================ + +@pytest.fixture(params=["cpu", "gpu"]) +def device_param(request): + """Pytest fixture to parametrize tests with CPU and GPU devices.""" + if request.param == "gpu" and not tf.config.list_physical_devices('GPU'): + pytest.skip("GPU not available") + return request.param + +@pytest.fixture(scope="function") +def image_data(device_param): + """Pytest fixture to provide test image data.""" + device_name = f'/{device_param.upper()}:0' + with tf.device(device_name): + rng = np.random.default_rng(seed=42) + image = Image.fromarray(rng.integers(0, 256, (640, 640, 3), dtype=np.uint8)) + + img_np = np.array(image, dtype=np.float32) + input_tensor = tf.expand_dims(img_np, axis=0) + return image, input_tensor + +def test_image_size(image_data): + """Test that the image has the expected size.""" + image, _ = image_data + expected_size = (640, 640) + assert image.size == expected_size + +@pytest.fixture(scope="function") +def model_data(image_data, device_param): + """Pytest fixture to provide model and predictions.""" + device_name = f'/{device_param.upper()}:0' + with tf.device(device_name): + _, input_tensor = image_data + model = MockRetinaNetModel(nb_classes=20) + processed_results = model(input_tensor) + return model, processed_results + +def test_model_outputs(model_data): + """Test that model outputs have the expected structure.""" + _, processed_results = model_data + assert isinstance(processed_results, dict) + assert set(processed_results.keys()) == {'boxes', 'confidence', 'class_probas'} + +def test_gradients_predict(image_data, model_data): + """Test that gradients should not be computed in predict mode.""" + _, input_tensor = image_data + model, _ = model_data + check = check_model_gradients(model.predict, input_tensor) + expected_msg = ( + "Model gradients should not be computed successfully in normal / predict mode." + ) + assert check is False, expected_msg + +def test_gradients_model_call(image_data, model_data): + """Test that gradients should be computed when calling the model directly.""" + _, input_tensor = image_data + model, _ = model_data + check = check_model_gradients(model, input_tensor) + assert check, "Model gradients should be computed successfully when calling the model." + +def test_box_model_wrapper(image_data): + """Test that the box model wrapper works correctly in both output modes.""" + image, input_tensor = image_data + image_size = (image.height, image.width) + + model = MockRetinaNetModel(nb_classes=20) + formatter = MockRetinaNetFormatter( + nb_classes=20, + image_size=image_size + ) + wrapper = TfBoxesModelWrapper(model, formatter) + + # Test output_as_list mode (default) + assert wrapper.output_as_list is True + output_list = wrapper(input_tensor) + assert isinstance(output_list, list) + assert len(output_list) == 1 # One result per batch item + assert isinstance(output_list[0], MultiBoxTensor) + + # Test output_as_tensor mode + wrapper.output_as_tensor = True + assert wrapper.output_as_list is False + output_tensor = wrapper(input_tensor) + assert isinstance(output_tensor, tf.Tensor) + assert len(output_tensor.shape) == 3 # (batch, num_boxes, features) + + + +@pytest.fixture(scope="module") +def dataset_classes(): + """Pytest fixture to provide dataset class names and mappings.""" + classes = ['aeroplane', 'bicycle', 'bird', 'boat', 'bottle', 'bus', 'car', + 'cat', 'chair', 'cow', 'diningtable', 'dog', 'horse', 'motorbike', + 'person', 'pottedplant', 'sheep', 'sofa', 'train', 'tvmonitor'] + nb_classes = len(classes) + label_to_color = {'person': 'r', + 'bicycle': 'b', + 'car': 'g', + 'motorcycle': 'y', + 'truck': 'orange'} + return classes, nb_classes, label_to_color + +@pytest.fixture(scope="function") +def latent_extractor_data(dataset_classes, device_param, image_data): + """Create a latent extractor using the MockExtractorBuilder.""" + _classes_names, nb_classes, _label_to_color = dataset_classes + image, _ = image_data + image_size = (image.height, image.width) + + # Use builder to create the latent extractor + latent_extractor = MockExtractorBuilder.build( + nb_classes=nb_classes, + image_size=image_size, + device=device_param + ) + + return latent_extractor + +def test_latent_extractor(image_data, dataset_classes, latent_extractor_data): + """Test that the latent extractor works correctly in both tensor and list modes.""" + image, input_tensor = image_data + classes_names, _nb_classes, label_to_color = dataset_classes + latent_extractor = latent_extractor_data + + # Test in tensor mode + latent_extractor.output_as_tensor = True + results = latent_extractor(input_tensor) + assert results.shape == (1, 10, 25), "Latent data shape should be (1, 10, 25)." + + # Test in list mode + latent_extractor.output_as_list = True + results_list = latent_extractor(input_tensor) + expected_msg = "Should return a list of MultiBoxTensor objects in list mode" + assert isinstance(results_list, list), expected_msg + assert len(results_list) == 1, "Should have one result per batch item" + assert isinstance(results_list[0], MultiBoxTensor), "First element should be a MultiBoxTensor" + assert results_list[0].shape == (10, 25), "MultiBoxTensor shape should be (10, 25)." + + filtered_results_list = results_list[0].filter(confidence=0.85) + assert filtered_results_list.shape == tf.TensorShape([4, 25]) + +def test_latent_extractor_gradients(image_data, latent_extractor_data): + """Test that latent extractor computes gradients successfully.""" + _image, input_tensor = image_data + latent_extractor = latent_extractor_data + + check = check_model_gradients(latent_extractor, input_tensor) + assert check, "Latent extractor gradients should be computed successfully." + +def test_latent_extractor_saliency(image_data, dataset_classes, latent_extractor_data): + """Test that saliency attribution works with the latent extractor.""" + image, input_tensor = image_data + classes_names, _nb_classes, _label_to_color = dataset_classes + latent_extractor = latent_extractor_data + + # Set to tensor mode for operator compatibility + latent_extractor.output_as_tensor = True + operator = xplique.Tasks.OBJECT_DETECTION_BOX_PROBA + + # Get targets in list mode first, then filter + latent_extractor.output_as_list = True + targets = latent_extractor(input_tensor) + box_to_explain = targets[0].filter(confidence=0.9, class_id=classes_names.index('car')) + + # Set back to tensor mode for the explainer + latent_extractor.output_as_tensor = True + box_to_explain = box_to_explain.to_batched_tensor() + + explainer = Saliency(latent_extractor, operator=operator, batch_size=None) + explanation = explainer.explain(input_tensor, targets=box_to_explain) + + # Check explanation shape (should match input image dimensions) + expected_shape = (1, 640, 640, 1) + assert explanation.shape == expected_shape, ( + f"Expected shape {expected_shape}, got {explanation.shape}" + ) + +@pytest.fixture(scope="function") +def craft_data(image_data, latent_extractor_data, device_param): + """Pytest fixture to create and fit a CRAFT instance.""" + device_name = f'/{device_param.upper()}:0' + with tf.device(device_name): + _image, input_tensor = image_data + latent_extractor = latent_extractor_data + + # Ensure latent extractor is in list mode for CRAFT + latent_extractor.output_as_list = True + craft = Craft(latent_extractor = latent_extractor, + number_of_concepts = 10) + craft.fit(input_tensor) + return craft + +def test_craft_reencode(image_data, dataset_classes, craft_data): + """Test that CRAFT encoding and decoding work correctly.""" + image, input_tensor = image_data + classes_names, _nb_classes, label_to_color = dataset_classes + craft = craft_data + + encoded_data = craft.encode(input_tensor) + assert len(encoded_data) == 1, "Should have one encoded data tuple per batch item" + latent_data, coeffs_u = encoded_data[0] + assert coeffs_u.shape == (1, 10, 10, 10), "Latent data shape should be (1, 10, 10, 10)." + + decoded_data = craft.decode(latent_data, coeffs_u) + assert isinstance(decoded_data, MultiBoxTensor), "Should return an MultiBoxTensor directly" + assert decoded_data.shape == (10, 25), "MultiBoxTensor shape should be (10, 25)." + + filtered_decoded_data = decoded_data.filter(confidence=0.85) + assert filtered_decoded_data.shape == tf.TensorShape([4, 25]) + +def test_craft_importance_gradient_input(image_data, dataset_classes, craft_data): + """Test that CRAFT importance estimation with gradient input works correctly.""" + _image, input_tensor = image_data + classes_names, _nb_classes, _label_to_color = dataset_classes + craft = craft_data + + # Test compute_gradient_input + operator = xplique.Tasks.OBJECT_DETECTION + class_id = classes_names.index("person") + partial_explainer = PartialExplainer( + explainer_class=GradientInput, + operator=operator, + reducer=None, + ) + explanation = craft.compute_explanation_per_concept( + input_tensor, class_id=class_id, confidence=0.6, + partial_explainer=partial_explainer + ) + + # Verify explanation shape + assert explanation.shape[0] == 1, "Should have one explanation per image" + assert explanation.shape[1:3] == (10, 10), "Should match coeffs_u spatial dimensions" + assert explanation.shape[3] == 10, "Should match number of concepts" + + # Test estimate_importance_gradient_input + importances_gi = craft.estimate_importance( + input_tensor, operator, class_id, confidence=0.6, method='gradient_input' + ) + assert importances_gi.shape == (10,), "Should return importance scores for each concept" + + # Verify importances are computed (don't check specific order for mock model) + order = importances_gi.argsort()[::-1] + assert len(order) == 10, "Should have ordering for all concepts" + +def test_craft_explanation_scores(craft_data): + """Test reduce_to_importance, reduce_to_prevalence and reduce_to_reliability + with hardcoded fake explanations. + + Explanation layout (n_images=3, H=2, W=2, n_concepts=craft.number_of_concepts): + - image 0 : only concept 0 active (value 1.0) + - image 1 : only concept 1 active (value 2.0) + - image 2 : only concept 0 active (value 3.0) + Dominant concepts: [0, 1, 0] + + Expected importance (spatial="mean", abs=True, aggregation="mean"): + concept 0 = (1.0 + 0.0 + 3.0) / 3 = 4/3 + concept 1 = (0.0 + 2.0 + 0.0) / 3 = 2/3 + others = 0.0 + + Expected prevalence: + concept 0 = 2/3, concept 1 = 1/3, others = 0.0, sum = 1.0 + + Expected reliability (accuracy = [0.9, 0.5, 0.7]): + concept 0 = mean(0.9, 0.7) = 0.8 + concept 1 = 0.5 + others = 0.0 + """ + craft = craft_data + n_concepts = craft.number_of_concepts # 10 (matches fixture) + + # Build hardcoded explanation: (N=3, H=2, W=2, n_concepts) + explanation = np.zeros((3, 2, 2, n_concepts), dtype=np.float32) + explanation[0, :, :, 0] = 1.0 # image 0 -> concept 0 + explanation[1, :, :, 1] = 2.0 # image 1 -> concept 1 + explanation[2, :, :, 0] = 3.0 # image 2 -> concept 0 + + # --- importance --- + importances = craft.reduce_to_importance( + explanation, spatial_reducer="mean", abs_before_reduce=True, aggregation_reducer="mean" + ) + assert importances.shape == (n_concepts,) + assert np.isclose(importances[0], 4.0 / 3), "concept 0 importance should be 4/3" + assert np.isclose(importances[1], 2.0 / 3), "concept 1 importance should be 2/3" + assert np.all(importances[2:] == 0.0), "other concepts should have zero importance" + + # --- prevalence --- + prevalence = craft.reduce_to_prevalence(explanation) + assert prevalence.shape == (n_concepts,) + assert np.isclose(prevalence[0], 2.0 / 3), "concept 0 prevalence should be 2/3" + assert np.isclose(prevalence[1], 1.0 / 3), "concept 1 prevalence should be 1/3" + assert np.all(prevalence[2:] == 0.0), "other concepts should have zero prevalence" + assert np.isclose(prevalence.sum(), 1.0), "prevalence should sum to 1" + + # --- reliability --- + accuracy = np.array([0.9, 0.5, 0.7], dtype=np.float32) + reliability = craft.reduce_to_reliability(explanation, accuracy) + assert reliability.shape == (n_concepts,) + assert np.isclose(reliability[0], 0.8), "concept 0 reliability should be mean(0.9, 0.7) = 0.8" + assert np.isclose(reliability[1], 0.5), "concept 1 reliability should be 0.5" + assert np.all(reliability[2:] == 0.0), "other concepts should have zero reliability" + + +def test_craft_decoder_modes(image_data, dataset_classes, craft_data): + """Test that the decoder works in both tensor and list modes.""" + _image, input_tensor = image_data + _classes_names, _nb_classes, _label_to_color = dataset_classes + craft = craft_data + + encoded_data = craft.encode(input_tensor) + latent_data, coeffs_u = encoded_data[0] + decoder = craft.make_concept_decoder(latent_data) + + # Decoder should always return tensor (unified behavior with PyTorch) + output_tensor = decoder(coeffs_u) + assert hasattr(output_tensor, 'shape'), "Decoder should always return a tensor" + assert output_tensor.shape == (1, 10, 25), "Should have correct tensor shape" + + # For filtering, use decode directly to get MultiBoxTensor + nbc_tensor = craft.decode(latent_data, coeffs_u) + assert hasattr(nbc_tensor, 'filter'), "decode should return MultiBoxTensor with filter method" + +def test_multibox_tensor_filter(image_data, dataset_classes, latent_extractor_data): + """Test MultiBoxTensor filtering functionality.""" + _image, input_tensor = image_data + classes_names, _nb_classes, _label_to_color = dataset_classes + latent_extractor = latent_extractor_data + + # Get results in list mode to access MultiBoxTensor + latent_extractor.output_as_list = True + targets = latent_extractor(input_tensor) + nbc_tensor = targets[0] # targets is a list, get the first MultiBoxTensor + assert nbc_tensor.shape == (10, 25), "MultiBoxTensor shape should be (10, 25)." + + # Test filtering by class_id and confidence + filtered = nbc_tensor.filter(class_id=classes_names.index('person'), confidence=0.5) + assert hasattr(filtered, 'shape'), "Filtered result should have shape attribute" + assert len(filtered.shape) == 2, "Filtered result should be 2D (boxes, features)" + assert filtered.shape[1] == 25, "Should preserve feature dimension" + + # Test filtering with high confidence (should return fewer boxes) + filtered_high = nbc_tensor.filter( + class_id=classes_names.index('person'), confidence=0.9 + ) + expected_msg = "Higher confidence should return fewer or equal boxes" + assert filtered_high.shape[0] <= filtered.shape[0], expected_msg + +def test_craft_encode_differentiable_gradients(image_data, craft_data): + """Test that encode with differentiable mode preserves gradients.""" + _image, input_tensor = image_data + craft = craft_data + + # Test differentiable encoding with gradient tape + with tf.GradientTape() as tape: + tape.watch(input_tensor) + + # Encode with differentiable mode + encoded_data = craft.encode(input_tensor, differentiable=True) + _latent_data, coeffs_u = encoded_data[0] + + # Verify coeffs_u is a tensor + expected_msg = "coeffs_u should be a tf.Tensor in differentiable mode" + assert isinstance(coeffs_u, tf.Tensor), expected_msg + assert tf.reduce_all(tf.math.is_finite(coeffs_u)), "coeffs_u should be finite" + assert tf.reduce_min(coeffs_u) >= -1e-6, "coeffs_u should remain non-negative" + + # Create a simple loss + loss = tf.reduce_sum(coeffs_u) + + # Check gradient flow + gradients = tape.gradient(loss, input_tensor) + + # Verify gradients flowed back to input + assert gradients is not None, "Gradients should flow back to input" + assert tf.reduce_all(tf.math.is_finite(gradients)), "Gradients should be finite" + assert tf.reduce_sum(tf.abs(gradients)) > 0, "Gradients should be non-zero" diff --git a/tests/concepts/test_holistic_craft_object_detection_torch.py b/tests/concepts/test_holistic_craft_object_detection_torch.py new file mode 100644 index 00000000..85b06e94 --- /dev/null +++ b/tests/concepts/test_holistic_craft_object_detection_torch.py @@ -0,0 +1,647 @@ +"""Tests for HolisticCraftTorch on object detection tasks with PyTorch.""" +# pylint: disable=unused-variable,unused-argument,redefined-outer-name + +import numpy as np +import pytest +import torch +import torchvision.transforms as T +from PIL import Image + +import xplique +from xplique.attributions import Saliency +from xplique.attributions.gradient_input import GradientInput +from xplique.concepts import HolisticCraftTorch as Craft +from xplique.concepts.holistic_craft import PartialExplainer +from xplique.concepts.latent_extractor import LatentData +from xplique.concepts.torch.latent_extractor import TorchLatentExtractor +from xplique.utils_functions.common.torch.gradients_check import check_model_gradients +from xplique.utils_functions.object_detection.base.box_manager import BoxFormat, BoxType +from xplique.utils_functions.object_detection.torch.box_model_wrapper import TorchBoxesModelWrapper +from xplique.utils_functions.object_detection.torch.multi_box_tensor import ( + TorchMultiBoxTensor as MultiBoxTensor, +) +from xplique.wrappers import TorchWrapper + +# ============================================================================ +# Mock Classes - Shared across tests +# ============================================================================ + +class MockTorchvisionModel(torch.nn.Module): + """Mock Torchvision-style model that returns list of predictions with gradient flow.""" + def forward(self, x): + """Build method for the mock extractor builder.""" + batch_size = x.shape[0] + device = x.device + + # Create a simple linear layer to ensure gradient flow from input + # Use mean pooling to reduce spatial dimensions + pooled = torch.mean(x, dim=[2, 3]) # (batch, channels) + + # Return list of dicts (Torchvision format) + predictions = [] + for i in range(batch_size): + # Create predictions that depend on input via pooled features + # Use pooled features to create a scaling factor + pooled_mean = pooled[i].mean().abs() + 1.0 # Add 1.0 to ensure decent scale + + # boxes: modulated by input features + boxes_base = torch.tensor( + [[0.1, 0.1, 0.3, 0.3], [0.2, 0.2, 0.4, 0.4], [0.3, 0.3, 0.5, 0.5], + [0.4, 0.4, 0.6, 0.6], [0.5, 0.5, 0.7, 0.7]], + device=device, + dtype=x.dtype, + ) + boxes = boxes_base * pooled_mean * 100 # Scale to reasonable box sizes + + # scores: ensure some boxes have high confidence for filtering tests + # First 2 boxes have high confidence (>0.9), rest have lower + # Multiply by a small gradient from pooled to maintain gradient flow + gradient_factor = (pooled_mean - 1.0) * 0.01 # Small variation from input + scores = torch.tensor([0.95, 0.92, 0.5, 0.3, 0.1], device=device) + gradient_factor + scores = torch.clamp(scores, 0.0, 1.0) # Ensure valid range + + # labels: ensure first box is 'car' (index 3 in COCO) for saliency test filtering + # This ensures at least one high-confidence box has the 'car' class + labels = torch.tensor([3, 1, 7, 9, 11], device=device) + + pred = { + 'boxes': boxes, + 'scores': scores, + 'labels': labels + } + predictions.append(pred) + return predictions + + +class MockTorchvisionBoxFormatter: + """Mock formatter for Torchvision-style predictions.""" + def __init__(self, nb_classes=20): + self.nb_classes = nb_classes + self.input_box_type = BoxType(BoxFormat.XYXY, is_normalized=False) + self.output_box_type = BoxType(BoxFormat.XYXY, is_normalized=False) + + def __call__(self, model_outputs): + """Make formatter callable.""" + return self.format_outputs(model_outputs) + + def format_outputs(self, model_outputs): + """Convert Torchvision outputs (list of dicts) to list of MultiBoxTensor.""" + results = [] + + for pred in model_outputs: + boxes = pred['boxes'] # (num_boxes, 4) + scores = pred['scores'] # (num_boxes,) + labels = pred['labels'] # (num_boxes,) + + # Create one-hot encoded class probabilities on the same device as boxes + num_boxes = boxes.shape[0] + device = boxes.device + class_probs = torch.zeros(num_boxes, self.nb_classes, device=device) + class_probs.scatter_(1, labels.unsqueeze(1), 1.0) + class_probs = class_probs * scores.unsqueeze(1) + + # Concatenate: [boxes (4), confidence (1), class_probs (nb_classes)] + combined = torch.cat([ + boxes, + scores.unsqueeze(1), + class_probs + ], dim=1) + + results.append(MultiBoxTensor(combined)) + + return results + +class MockTorchvisionLatentData(LatentData): + """Mock latent data with fixed reproducible activations in (N, H, W, C) format.""" + + def __init__(self, activations: np.ndarray, batch_size=1, device='cpu'): + """ + Parameters + ---------- + activations : np.ndarray + Initial activations in (N, H, W, C) format. + batch_size : int + Batch size + device : str or torch.device + Device for tensors + """ + self.batch_size = batch_size + self.device = device + self.activations = activations + + def __len__(self): + """Return batch size.""" + return self.batch_size + + def detach(self): + """Detach all tensors from computation graph.""" + self.activations = self.activations.detach() + return self + + def get_activations(self, as_numpy=True, keep_gradients=False): + """ + Get activations in (N, H, W, C) format. + + Returns + ------- + activations : np.ndarray or torch.Tensor + Shape (batch, H, W, C) + """ + activations = self.activations + + if not keep_gradients: + activations = activations.detach() + + if as_numpy: + activations = activations.cpu().numpy() + + return activations + + def set_activations(self, values: np.ndarray): + """ + Set activations from (N, H, W, C) format. + + Parameters + ---------- + values : np.ndarray or torch.Tensor + Shape (batch, H, W, C) + """ + # Convert numpy to tensor if needed + if isinstance(values, np.ndarray): + values = torch.from_numpy(values).to(self.device) + + self.activations = values + + def to(self, device): + """Move all data to specified device.""" + self.device = device + self.activations = self.activations.to(device) + return self + + +class MockExtractorBuilder: + """ + Builder for creating LatentExtractor instances for mock Torchvision-style models. + + This class encapsulates the logic for creating a TorchLatentExtractor configured + for mock object detection models that follow the Torchvision output format. + """ + + @classmethod + def build(cls, model, device='cpu', batch_size=1): + """ + Build a LatentExtractor for a mock Torchvision-style model. + + Creates custom input_to_latent and latent_to_logit functions that split + the model's forward pass into feature extraction and prediction. + + Parameters + ---------- + model : torch.nn.Module + Mock Torchvision-style object detection model + device : str or torch.device + Device to run computations on. Default is 'cpu'. + batch_size : int + Batch size for processing. Default is 1. + + Returns + ------- + latent_extractor : TorchLatentExtractor + Configured TorchLatentExtractor instance for the mock model. + """ + def input_to_latent_fn(inputs): + """Simple mock latent extractor - just pool input to fixed size.""" + batch_size = inputs.shape[0] + device = inputs.device + + # Pool to fixed spatial size (10x10) and apply ReLU for NMF + pooled = torch.nn.functional.adaptive_avg_pool2d(inputs, (10, 10)) # (N, C, 10, 10) + activations = torch.relu(pooled.permute(0, 2, 3, 1)) # (N, 10, 10, C) with ReLU + + return MockTorchvisionLatentData(activations=activations, + batch_size=batch_size, + device=device) + + def latent_to_logit_fn(latent_data): + """Simple mock decoder - reuse the existing MockTorchvisionModel logic.""" + activations = latent_data.get_activations(as_numpy=False, keep_gradients=True) + # Convert back to channel-first for model: (N, 10, 10, C) -> (N, C, 10, 10) + inputs_chw = activations.permute(0, 3, 1, 2) + return model(inputs_chw) + + # Create mock Torchvision formatter + formatter = MockTorchvisionBoxFormatter(nb_classes=20) + + # Create TorchLatentExtractor with simplified mock functions + latent_extractor = TorchLatentExtractor( + model=model, + input_to_latent_model=input_to_latent_fn, + latent_to_logit_model=latent_to_logit_fn, + latent_data_class=MockTorchvisionLatentData, + output_formatter=formatter, + batch_size=batch_size, + device=str(device) + ) + + return latent_extractor + + +# ============================================================================ +# Test Fixtures +# ============================================================================ + +@pytest.fixture(scope="session", params=["cpu", "cuda"]) +def device_param(request): + """Fixture providing the device parameter for tests.""" + device_str = request.param + if device_str == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + device = torch.device(device_str) + return device + +@pytest.fixture(scope="function") +def image_data(device_param): + """Fixture providing image data for tests.""" + rng = np.random.default_rng(seed=42) + raw_image = Image.fromarray(rng.integers(0, 256, (462, 640, 3), dtype=np.uint8)) + + # standard PyTorch mean-std input image normalization + transform = T.Compose([ + T.Resize((800, 800)), + T.ToTensor(), + T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) + ]) + + # preprocess and batch the image + input_tensor = transform(raw_image).unsqueeze(0).to(device_param) + + return raw_image, input_tensor + +def test_image_size(image_data): + """Fixture providing a mock model for tests.""" + image, _ = image_data + expected_size = (640, 462) + assert image.size == expected_size + + +@pytest.fixture(scope="function") +def model_data(image_data, device_param): + """Create a mock Torchvision-style object detection model.""" + _, input_tensor = image_data + + model = MockTorchvisionModel().to(device_param) + model.eval() + processed_results = model(input_tensor) + return model, processed_results + +def test_model_outputs(model_data): + """Fixture providing model data including model, formatter, and device.""" + _, processed_results = model_data + # Torchvision format: list of dicts with 'boxes', 'scores', 'labels' + assert isinstance(processed_results, list) + assert len(processed_results) > 0 + assert 'boxes' in processed_results[0] + assert 'scores' in processed_results[0] + assert 'labels' in processed_results[0] + + +def test_model_outputs_are_deterministic(image_data, model_data): + """The mock must produce identical detections for the same input.""" + _, input_tensor = image_data + model, first_results = model_data + second_results = model(input_tensor) + + for first, second in zip(first_results, second_results): + assert torch.equal(first['boxes'], second['boxes']) + assert torch.equal(first['scores'], second['scores']) + assert torch.equal(first['labels'], second['labels']) + +def test_gradients_model_original(image_data, model_data): + """Fixture providing a mock latent extractor for tests.""" + image, input_tensor = image_data + model, _ = model_data + check = check_model_gradients(model, input_tensor) + assert check, "Mock model gradients should be computed successfully." + +def test_multibox_tensor_to_batched(): + """Test MultiBoxTensor to_batched_tensor method.""" + # Create a mock MultiBoxTensor with shape (num_boxes, features) + tensor_data = torch.rand(5, 25) # 5 boxes, 25 features (4+1+20) + nbc = MultiBoxTensor(tensor_data) + + # Test to_batched_tensor() + batched = nbc.to_batched_tensor() + assert isinstance(batched, torch.Tensor) + assert batched.shape == (1, 5, 25) # (1, num_boxes, features) + assert torch.allclose(batched[0], tensor_data) + +def test_box_model_wrapper(): + """Test TorchBoxesModelWrapper initialization and functionality.""" + model = MockTorchvisionModel() + formatter = MockTorchvisionBoxFormatter(nb_classes=20) + wrapper = TorchBoxesModelWrapper(model, formatter) + + # Test with batch_size=2 + input_tensor = torch.randn(2, 3, 224, 224) + + # Test output_as_list mode (default) + assert wrapper.output_as_list is True + output_list = wrapper(input_tensor) + assert isinstance(output_list, list) + assert len(output_list) == 2 + + # Test output_as_tensor mode + wrapper.output_as_tensor = True + assert wrapper.output_as_list is False + output_tensor = wrapper(input_tensor) + assert isinstance(output_tensor, torch.Tensor) + assert output_tensor.shape == (2, 5, 25) # (batch, num_boxes, features) + +@pytest.fixture(scope="function") +def dataset_classes(): + """Fixture providing COCO dataset classes and metadata.""" + # COCO classes + classes = [ + 'N/A', 'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', + 'train', 'truck', 'boat', 'traffic light', 'fire hydrant', 'N/A', + 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', + 'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe', 'N/A', 'backpack', + 'umbrella', 'N/A', 'N/A', 'handbag', 'tie', 'suitcase', 'frisbee', 'skis', + 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', + 'skateboard', 'surfboard', 'tennis racket', 'bottle', 'N/A', 'wine glass', + 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple', 'sandwich', + 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', + 'chair', 'couch', 'potted plant', 'bed', 'N/A', 'dining table', 'N/A', + 'N/A', 'toilet', 'N/A', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', + 'cell phone', 'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'N/A', + 'book', 'clock', 'vase', 'scissors', 'teddy bear', 'hair drier', + 'toothbrush' + ] + nb_classes = len(classes) + label_to_color = {'person': 'r', + 'bicycle': 'b', + 'car': 'g', + 'motorcycle': 'y', + 'truck': 'orange'} + return classes, nb_classes, label_to_color + +@pytest.fixture(scope="function") +def latent_extractor_data(dataset_classes, model_data, device_param): + """Create latent extractor using MockExtractorBuilder.""" + classes_names, nb_classes, label_to_color = dataset_classes + model, _ = model_data + + # Use the builder to create the latent extractor + latent_extractor = MockExtractorBuilder.build( + model=model, + device=str(device_param), + batch_size=1 + ) + + return latent_extractor + +def test_latent_extractor(image_data, dataset_classes, latent_extractor_data): + """Fixture providing latent extractor with model data.""" + image, input_tensor = image_data + classes_names, nb_classes, label_to_color = dataset_classes + latent_extractor = latent_extractor_data + + results = latent_extractor(input_tensor) + assert isinstance(results, list), "Results should be a list of MultiBoxTensor objects" + assert len(results) == 1, "Should have one result per batch item" + # Torchvision model: 5 boxes, each with 4+1+20 features (boxes, confidence, class_probs) + expected_shape = torch.Size([5, 25]) + error_msg = (f"MultiBoxTensor shape should be {expected_shape}, " + f"got {results[0].shape}") + assert results[0].shape == expected_shape, error_msg + + filtered_results = results[0].filter(confidence=0.5) + assert filtered_results.shape == torch.Size([3, 25]) + +def test_latent_extractor_gradients(image_data, latent_extractor_data): + """Fixture providing crop data for tests.""" + image, input_tensor = image_data + latent_extractor = latent_extractor_data + + check = check_model_gradients(latent_extractor, input_tensor) + assert check, "Latent extractor gradients should be computed successfully." + +def test_latent_extractor_saliency(image_data, dataset_classes, + latent_extractor_data, device_param): + """Test latent extractor saliency computation.""" + image, input_tensor = image_data + classes_names, nb_classes, label_to_color = dataset_classes + latent_extractor = latent_extractor_data + + targets = latent_extractor(input_tensor) + class_id_car = classes_names.index('car') + filtered_targets = targets[0].filter(confidence=0.9, class_id=class_id_car) + box_to_explain = np.expand_dims(filtered_targets.detach().cpu().numpy(), axis=0) + + latent_extractor.output_as_tensor = True + input_tensor_tf_dim = input_tensor.detach().cpu().numpy().transpose(0, 2, 3, 1) + + torch_wrapped_model = TorchWrapper(latent_extractor, device=device_param, + is_channel_first=True) + + operator = xplique.Tasks.OBJECT_DETECTION + explainer = Saliency(torch_wrapped_model, operator=operator, batch_size=1) + explanation = explainer.explain(input_tensor_tf_dim, targets=box_to_explain) + assert explanation.shape == (1, 800, 800, 1) + +@pytest.fixture(scope="function") +def craft_data(image_data, latent_extractor_data, device_param): + """Fixture providing CRAFT instance for tests.""" + image, input_tensor = image_data + latent_extractor = latent_extractor_data + + craft = Craft(latent_extractor = latent_extractor, + number_of_concepts = 10, + device = str(device_param)) + craft.fit(input_tensor) + return craft + +def test_craft_reencode(image_data, dataset_classes, craft_data): + """Fixture providing crop selection data for tests.""" + image, input_tensor = image_data + classes_names, nb_classes, label_to_color = dataset_classes + craft = craft_data + + # The encode method now returns a list of tuples [(latent_data, coeffs_u), ...] + latent_data_coeffs_u_list = craft.encode(input_tensor) + + # Should have one tuple per image in the batch + expected_len = input_tensor.shape[0] + actual_len = len(latent_data_coeffs_u_list) + error_msg = f"Expected {expected_len} tuples, got {actual_len}" + assert len(latent_data_coeffs_u_list) == expected_len, error_msg + + # Get the first tuple (since we're only testing with one image) + latent_data, coeffs_u = latent_data_coeffs_u_list[0] + + # Torchvision mock uses spatial feature maps (H=10, W=10) with 10 concepts + expected_shape = (1, 10, 10, 10) + error_msg = f"Latent data shape should be {expected_shape}, got {coeffs_u.shape}" + assert coeffs_u.shape == expected_shape, error_msg + + result = craft.decode(latent_data, coeffs_u) + assert isinstance(result, MultiBoxTensor), "Decoded result should be an MultiBoxTensor directly" + # Result should have 5 boxes, 25 features each (4 box coords + 1 confidence + 20 class probs) + expected_shape = torch.Size([5, 25]) + error_msg = f"Decoded MultiBoxTensor shape should be [5, 25], got {result.shape}" + assert result.shape == expected_shape, error_msg + + filtered_result = result.filter(confidence=0.5) + assert filtered_result.shape == torch.Size([3, 25]) + +def test_craft_decoder_modes(image_data, dataset_classes, craft_data, device_param): + """Test CRAFT encode/decode functionality.""" + image, input_tensor = image_data + classes_names, nb_classes, label_to_color = dataset_classes + craft = craft_data + + encoded_data = craft.encode(input_tensor) + latent_data, coeffs_u = encoded_data[0] + decoder = craft.make_concept_decoder(latent_data) + + # Decoder should always return tensor (unified behavior with TensorFlow) + output_tensor = decoder(coeffs_u) + assert hasattr(output_tensor, 'shape'), "Decoder should always return a tensor" + # Torchvision model: (1, num_boxes, features) = (1, 5, 25) + expected_shape = torch.Size([1, 5, 25]) + error_msg = f"Expected shape {expected_shape}, got {output_tensor.shape}" + assert output_tensor.shape == expected_shape, error_msg + + # For filtering, use decode directly to get MultiBoxTensor + nbc_tensor = craft.decode(latent_data, coeffs_u) + error_msg = "decode should return MultiBoxTensor with filter method" + assert hasattr(nbc_tensor, 'filter'), error_msg + + +def test_craft_importance_gradient_input(image_data, dataset_classes, craft_data): + """Test CRAFT importance estimation by gradient input.""" + image, input_tensor = image_data + classes_names, nb_classes, label_to_color = dataset_classes + craft = craft_data + + # Test compute_gradient_input + operator = xplique.Tasks.OBJECT_DETECTION + class_id = classes_names.index("person") + partial_explainer = PartialExplainer( + explainer_class=GradientInput, + operator=operator, + reducer=None, + ) + explanation = craft.compute_explanation_per_concept( + input_tensor, class_id=class_id, confidence=0.6, + partial_explainer=partial_explainer + ) + + # Verify explanation shape - Torchvision mock produces 4D explanations + # Shape: (batch, H, W, num_concepts) = (1, 10, 10, 10) + assert explanation.shape[0] == 1, "Should have one explanation per image" + assert explanation.shape[1] == 10, "Should match feature map height" + assert explanation.shape[2] == 10, "Should match feature map width" + assert explanation.shape[3] == 10, "Should match number of concepts" + + importances_gi = craft.estimate_importance( + input_tensor, operator, class_id, confidence=0.6, method='gradient_input' + ) + + # Verify importances are computed and have correct shape + assert importances_gi.shape[0] == 10, "Should have importance for each concept" + assert np.all(np.isfinite(importances_gi)), "All importances should be finite" + + order = importances_gi.argsort()[::-1] + + # For mock data, we just verify we got a valid ordering (not checking specific order) + # Real Torchvision models would produce different orderings based on actual features + assert len(order) == 10, "Should have ordered all 10 concepts" + assert len(np.unique(order)) == 10, "All concepts should have unique ordering" + + +def test_craft_explanation_scores(craft_data): + """Test reduce_to_importance, reduce_to_prevalence and reduce_to_reliability + with hardcoded fake explanations. + + Explanation layout (n_images=3, H=2, W=2, n_concepts=craft.number_of_concepts): + - image 0 : only concept 0 active (value 1.0) + - image 1 : only concept 1 active (value 2.0) + - image 2 : only concept 0 active (value 3.0) + Dominant concepts: [0, 1, 0] + + Expected importance (spatial="mean", abs=True, aggregation="mean"): + concept 0 = (1.0 + 0.0 + 3.0) / 3 = 4/3 + concept 1 = (0.0 + 2.0 + 0.0) / 3 = 2/3 + others = 0.0 + + Expected prevalence: + concept 0 = 2/3, concept 1 = 1/3, others = 0.0, sum = 1.0 + + Expected reliability (accuracy = [0.9, 0.5, 0.7]): + concept 0 = mean(0.9, 0.7) = 0.8 + concept 1 = 0.5 + others = 0.0 + """ + craft = craft_data + n_concepts = craft.number_of_concepts # 10 (matches fixture) + + # Build hardcoded explanation: (N=3, H=2, W=2, n_concepts) + explanation = np.zeros((3, 2, 2, n_concepts), dtype=np.float32) + explanation[0, :, :, 0] = 1.0 # image 0 → concept 0 + explanation[1, :, :, 1] = 2.0 # image 1 → concept 1 + explanation[2, :, :, 0] = 3.0 # image 2 → concept 0 + + # --- importance --- + importances = craft.reduce_to_importance( + explanation, spatial_reducer="mean", abs_before_reduce=True, aggregation_reducer="mean" + ) + assert importances.shape == (n_concepts,) + assert np.isclose(importances[0], 4.0 / 3), "concept 0 importance should be 4/3" + assert np.isclose(importances[1], 2.0 / 3), "concept 1 importance should be 2/3" + assert np.all(importances[2:] == 0.0), "other concepts should have zero importance" + + # --- prevalence --- + prevalence = craft.reduce_to_prevalence(explanation) + assert prevalence.shape == (n_concepts,) + assert np.isclose(prevalence[0], 2.0 / 3), "concept 0 prevalence should be 2/3" + assert np.isclose(prevalence[1], 1.0 / 3), "concept 1 prevalence should be 1/3" + assert np.all(prevalence[2:] == 0.0), "other concepts should have zero prevalence" + assert np.isclose(prevalence.sum(), 1.0), "prevalence should sum to 1" + + # --- reliability --- + accuracy = np.array([0.9, 0.5, 0.7], dtype=np.float32) + reliability = craft.reduce_to_reliability(explanation, accuracy) + assert reliability.shape == (n_concepts,) + assert np.isclose(reliability[0], 0.8), "concept 0 reliability should be mean(0.9, 0.7) = 0.8" + assert np.isclose(reliability[1], 0.5), "concept 1 reliability should be 0.5" + assert np.all(reliability[2:] == 0.0), "other concepts should have zero reliability" + + +def test_craft_encode_differentiable_gradients(image_data, craft_data): + """Test that encode(differentiable=True) preserves gradients through the encoding pipeline.""" + image, input_tensor = image_data + craft = craft_data + + # Ensure input has gradients enabled + input_with_grad = input_tensor.clone().detach().requires_grad_(True) + + # Test differentiable encoding + encoded_data = craft.encode(input_with_grad, differentiable=True) + latent_data, coeffs_u = encoded_data[0] + + # Verify coeffs_u is a tensor with gradients + expected_type_msg = "coeffs_u should be a torch.Tensor in differentiable mode" + assert isinstance(coeffs_u, torch.Tensor), expected_type_msg + assert torch.all(torch.isfinite(coeffs_u)).item(), "coeffs_u should be finite" + assert torch.min(coeffs_u).item() >= -1e-6, "coeffs_u should remain non-negative" + assert coeffs_u.requires_grad, "coeffs_u should have gradients enabled" + + # Create a simple loss and check gradient flow + loss = coeffs_u.sum() + loss.backward() + gradients = input_with_grad.grad + + # Verify gradients flowed back to input + assert gradients is not None, "Gradients should flow back to input" + assert gradients.abs().sum() > 0, "Gradients should be non-zero" diff --git a/tests/concepts/test_holistic_craft_regressions.py b/tests/concepts/test_holistic_craft_regressions.py new file mode 100644 index 00000000..3f77818a --- /dev/null +++ b/tests/concepts/test_holistic_craft_regressions.py @@ -0,0 +1,292 @@ +"""Regression tests for framework-agnostic HolisticCraft behavior.""" + +from contextlib import contextmanager + +import matplotlib.pyplot as plt +import numpy as np +import pytest +from sklearn.exceptions import NotFittedError + +from xplique.concepts.craft import Factorization +from xplique.concepts.holistic_craft import HolisticCraft, PartialExplainer +from xplique.concepts.latent_extractor import LatentData + + +class _LatentData(LatentData): + def __init__(self, activations): + self.activations = activations + + def get_activations(self, as_numpy=True, keep_gradients=False): + return self.activations + + def set_activations(self, values): + self.activations = values + + +class _Prediction: + def filter(self, class_id=None, confidence=None): + return self + + def to_attribution_target(self, class_id=None): + return self + + def to_batched_tensor(self): + return np.ones((1, 1), dtype=np.float32) + + def __len__(self): + return 1 + + +class _Extractor: + def __init__(self, latent_data, batch_size=1): + self.latent_data = latent_data + self.batch_size = batch_size + self.forced_batch_sizes = [] + + @contextmanager + def temporary_force_batch_size(self, batch_size): + old_batch_size = self.batch_size + self.forced_batch_sizes.append(batch_size) + self.batch_size = batch_size + try: + yield + finally: + self.batch_size = old_batch_size + + def input_to_latent_generator(self, inputs, resize=None, keep_gradients=False): + yield from self.latent_data + + def latent_to_logit(self, latent_data): + return _Prediction() + + +class _Factorizer: + is_fitted = True + requires_positive_activations = False + + def encode(self, activations): + return activations + + +class _UnfittedFactorizer(_Factorizer): + is_fitted = False + + +class _ArrayLike: + def __init__(self, values): + self.values = values + + +class _Explainer: + def __init__(self, model, batch_size): + pass + + def explain(self, coeffs_u, targets): + return _ArrayLike(np.ones_like(coeffs_u)) + + +class _BadShapeExplainer(_Explainer): + def explain(self, coeffs_u, targets): + return _ArrayLike(np.ones(coeffs_u.shape[:-1] + (1,))) + + +class _Framework: + float32 = np.float32 + + +class _Craft(HolisticCraft): + def __init__(self, latent_data, batch_size=1): + factorizer = _Factorizer() + super().__init__( + _Extractor(latent_data, batch_size), number_of_concepts=2, factorizer=factorizer + ) + self.factorization = Factorization(None, 0, None, factorizer, None, np.eye(2)) + self.framework = "tf" + self._framework_module = _Framework + + def latent_to_concept_differentiable(self, latent_data): + return latent_data.activations + + def _to_numpy(self, tensor): + return tensor if isinstance(tensor, np.ndarray) else tensor.values + + def _to_tensor(self, array, dtype=None): + return np.asarray(array, dtype=dtype) + + def make_concept_decoder(self, latent_data): + return object() + + +class _DummyModel: + def __call__(self, _): + return np.ones((1, 1), dtype=np.float32) + + +class _MinimalStructuredPrediction: + def __init__(self, empty): + self.is_empty = empty + + def filter(self, class_id=None, confidence=None): + return self + + def to_attribution_target(self, class_id=None): + return self + + def to_batched_tensor(self): + return np.ones((1, 1), dtype=np.float32) + + +def test_factorization_preserves_its_positional_field_order(): + factorization = Factorization("inputs", 3, "crops", "reducer", "crops_u", "concept_bank") + + assert factorization.inputs == "inputs" + assert factorization.class_id == 3 + assert factorization.crops == "crops" + assert factorization.reducer == "reducer" + assert factorization.crops_u == "crops_u" + assert factorization.concept_bank_w == "concept_bank" + assert factorization.coeffs_u is None + + +def test_token_explanations_use_token_axes_and_framework_conversion(): + craft = _Craft([_LatentData(np.ones((1, 2, 2), dtype=np.float32))]) + explanation = craft.compute_explanation_per_concept( + np.ones((1, 2, 2, 1)), PartialExplainer(_Explainer) + ) + + assert explanation.shape == (1, 2, 2) + token_explanations = np.array([[[1.0, 0.0], [3.0, 0.0]], [[0.0, 2.0], [0.0, 4.0]]]) + np.testing.assert_allclose( + craft.reduce_to_importance(token_explanations, spatial_reducer="mean"), [1.0, 1.5] + ) + np.testing.assert_allclose(craft.reduce_to_prevalence(token_explanations), [0.5, 0.5]) + np.testing.assert_allclose( + craft.reduce_to_reliability(token_explanations, np.array([1.0, 0.5])), [1.0, 0.5] + ) + + +def test_explanations_use_configured_batch_size_for_perturbations(): + explainer_batch_sizes = [] + + class RecordingExplainer(_Explainer): + def __init__(self, model, batch_size): + super().__init__(model, batch_size) + explainer_batch_sizes.append(batch_size) + + craft = _Craft([_LatentData(np.ones((1, 2, 2), dtype=np.float32))], batch_size=4) + explanation = craft.compute_explanation_per_concept( + np.ones((1, 2, 2, 1)), PartialExplainer(RecordingExplainer) + ) + + assert explanation.shape == (1, 2, 2) + assert explainer_batch_sizes == [4] + assert craft.latent_extractor.forced_batch_sizes == [1] + assert craft.latent_extractor.batch_size == 4 + + +def test_invalid_explanation_shape_and_empty_extractions_raise_value_errors(): + craft = _Craft([_LatentData(np.ones((1, 2, 2), dtype=np.float32))]) + + with pytest.raises(ValueError, match="Explanation shape"): + craft.compute_explanation_per_concept( + np.ones((1, 2, 2, 1)), PartialExplainer(_BadShapeExplainer) + ) + + empty_craft = _Craft([]) + with pytest.raises(ValueError, match="No activations"): + empty_craft.fit(np.ones((1, 2, 2, 1))) + with pytest.raises(ValueError, match="No activations"): + empty_craft.transform(np.ones((1, 2, 2, 1))) + with pytest.raises(ValueError, match="No latent data"): + empty_craft.compute_explanation_per_concept( + np.ones((1, 2, 2, 1)), PartialExplainer(_Explainer) + ) + + +def test_check_if_fitted_requires_factorization_and_fitted_factorizer(): + craft = _Craft([]) + craft.factorizer = _UnfittedFactorizer() + craft.factorization = Factorization(None, 0, None, craft.factorizer, None, np.eye(2)) + + with pytest.raises(NotFittedError): + craft.check_if_fitted() + + craft.factorizer = _Factorizer() + craft.factorization = None + with pytest.raises(NotFittedError): + craft.check_if_fitted() + + +def test_estimate_importance_accepts_operator_in_method_kwargs(): + craft = _Craft([_LatentData(np.ones((1, 2, 2), dtype=np.float32))]) + captured = {} + + def _capture_partial_explainer( + _, partial_explainer, class_id=None, confidence=None, verbose=False + ): + captured.update(partial_explainer.kwargs) + return np.ones((1, 2), dtype=np.float32) + + craft.compute_explanation_per_concept = _capture_partial_explainer + + craft.estimate_importance( + images=np.ones((1, 2, 2, 1), dtype=np.float32), + operator=_DummyModel(), + class_id=0, + method="gradient_input", + reducer="sum", + spatial_reducer="mean", + aggregation_reducer="mean", + ) + assert isinstance(captured["operator"], _DummyModel) + assert captured["reducer"] == "sum" + + captured.clear() + craft.estimate_importance( + images=np.ones((1, 2, 2, 1), dtype=np.float32), + operator=_DummyModel(), + class_id=0, + method="sobol", + nb_channels=7, + spatial_reducer="mean", + aggregation_reducer="mean", + ) + assert isinstance(captured["operator"], _DummyModel) + assert captured["nb_channels"] == 7 + + +def test_compute_explanation_accepts_structured_predictions_without_len(): + craft = _Craft([_LatentData(np.ones((1, 2, 2), dtype=np.float32))]) + + craft.decode = lambda latent_data, coeffs_u: _MinimalStructuredPrediction(empty=True) + explanation = craft.compute_explanation_per_concept( + np.ones((1, 2, 2, 1), dtype=np.float32), + PartialExplainer(_Explainer), + ) + assert explanation.shape == (1, 2, 2) + + craft.decode = lambda latent_data, coeffs_u: _MinimalStructuredPrediction(empty=False) + explanation = craft.compute_explanation_per_concept( + np.ones((1, 2, 2, 1), dtype=np.float32), + PartialExplainer(_Explainer), + ) + assert explanation.shape == (1, 2, 2) + + +def test_display_validates_concept_order_and_handles_a_single_column(): + craft = _Craft([]) + images = np.ones((2, 4, 4, 3), dtype=np.float32) + coeffs_u = np.ones((2, 2, 2, 2), dtype=np.float32) + + with pytest.raises(ValueError, match="between 0"): + craft.display_images_per_concept(images, coeffs_u, order=[2]) + with pytest.raises(ValueError, match="more IDs"): + craft.display_images_per_concept(images, coeffs_u, order=[0, 1, 0]) + + figure = craft.display_images_per_concept(images, coeffs_u, order=[0]) + assert len(figure.axes) == 2 + plt.close(figure) + + figure = craft.display_top_images_per_concept(images, topk=2, coeffs_u=coeffs_u, order=[0]) + assert len(figure.axes) == 2 + plt.close(figure) diff --git a/tests/concepts/test_latent_extractor_tf.py b/tests/concepts/test_latent_extractor_tf.py new file mode 100644 index 00000000..7e431d54 --- /dev/null +++ b/tests/concepts/test_latent_extractor_tf.py @@ -0,0 +1,89 @@ +"""Regression tests for TensorFlow latent extractors.""" + +import numpy as np +import pytest +import tensorflow as tf + +from xplique.concepts.tf.latent_extractor import TfLatentExtractor +from xplique.concepts.tf.layered_model_latent_extractor import LayeredModelExtractorBuilder + + +def _identity_extractor(batch_size=2): + return TfLatentExtractor( + model=lambda inputs: inputs, + input_to_latent_model=lambda inputs: inputs, + latent_to_logit_model=lambda latent_data: latent_data, + batch_size=batch_size, + ) + + +def _functional_classifier(): + inputs = tf.keras.Input(shape=(2, 2, 1)) + hidden = tf.keras.layers.Conv2D(4, 1, activation="relu", name="hidden")(inputs) + pooled = tf.keras.layers.GlobalAveragePooling2D(name="pool")(hidden) + outputs = tf.keras.layers.Dense(2, name="classifier")(pooled) + return tf.keras.Model(inputs, outputs) + + +def test_functional_builder_preserves_residual_graph(): + inputs = tf.keras.Input(shape=(2, 2, 1)) + split = tf.keras.layers.Conv2D(4, 1, activation="relu", name="split")(inputs) + left = tf.keras.layers.Conv2D(4, 1, name="left")(split) + right = tf.keras.layers.Conv2D(4, 1, name="right")(split) + merged = tf.keras.layers.Add(name="merge")([left, right]) + pooled = tf.keras.layers.GlobalAveragePooling2D(name="pool")(merged) + outputs = tf.keras.layers.Dense(2, name="classifier")(pooled) + model = tf.keras.Model(inputs, outputs) + + split_index = model.layers.index(model.get_layer("split")) + extractor = LayeredModelExtractorBuilder.build(model, split_layer=split_index) + samples = tf.constant(np.arange(8, dtype=np.float32).reshape(2, 2, 2, 1)) + + latent_data = extractor.input_to_latent(samples) + predictions = extractor.latent_to_logit(latent_data).tensor + + tf.debugging.assert_near(predictions, model(samples)) + tf.debugging.assert_near(latent_data.activations, model.get_layer("split")(samples)) + + +@pytest.mark.parametrize("split_layer", ["hidden", True, -5, 4]) +def test_functional_builder_validates_split_layer(split_layer): + with pytest.raises(ValueError, match="split_layer"): + LayeredModelExtractorBuilder.build(_functional_classifier(), split_layer=split_layer) + + +def test_functional_builder_requires_split_to_be_a_graph_cut(): + inputs = tf.keras.Input(shape=(2, 2, 1)) + split = tf.keras.layers.Conv2D(2, 1, name="split")(inputs) + bypass = tf.keras.layers.Conv2D(2, 1, name="bypass")(inputs) + outputs = tf.keras.layers.Add()([split, bypass]) + model = tf.keras.Model(inputs, outputs) + + with pytest.raises(ValueError, match="graph cut"): + LayeredModelExtractorBuilder.build( + model, split_layer=model.layers.index(model.get_layer("split")) + ) + + +def test_classifier_tensor_mode_keeps_existing_batch_dimension(): + extractor = LayeredModelExtractorBuilder.build(_functional_classifier(), split_layer=1) + extractor.output_as_tensor = True + + predictions = extractor(tf.ones((2, 2, 2, 1))) + + assert isinstance(predictions, tf.Tensor) + assert predictions.shape == (2, 2) + + +def test_input_validation_and_single_image_batching(): + extractor = _identity_extractor() + + latent_data = extractor.input_to_latent(tf.ones((2, 2, 1))) + + assert latent_data.shape == (1, 2, 2, 1) + with pytest.raises(ValueError, match="rank 3"): + extractor.input_to_latent(tf.ones((2, 2))) + with pytest.raises(ValueError, match="at least one"): + list(extractor.input_to_latent_generator(tf.zeros((0, 2, 2, 1)))) + with pytest.raises(ValueError, match="batch_size"): + _identity_extractor(batch_size=0) diff --git a/tests/concepts/test_latent_extractor_torch.py b/tests/concepts/test_latent_extractor_torch.py new file mode 100644 index 00000000..053acadb --- /dev/null +++ b/tests/concepts/test_latent_extractor_torch.py @@ -0,0 +1,103 @@ +"""Regression tests for PyTorch latent extractors.""" +# ruff: noqa: E402 + +import numpy as np +import pytest + +torch = pytest.importorskip("torch") + +from xplique.concepts.torch.holistic_craft import HolisticCraftTorch +from xplique.concepts.torch.latent_extractor import TorchLatentExtractor +from xplique.concepts.torch.layered_model_latent_extractor import LayeredModelExtractorBuilder + + +def _identity_extractor(device=None, batch_size=2): + return TorchLatentExtractor( + model=torch.nn.Identity(), + input_to_latent_model=lambda inputs: inputs * 2, + latent_to_logit_model=lambda latent_data: latent_data, + device=device, + batch_size=batch_size, + ) + + +def test_torch_builder_includes_the_requested_layer_for_positive_and_negative_indices(): + model = torch.nn.Sequential( + torch.nn.Conv2d(1, 4, 1), + torch.nn.ReLU(), + torch.nn.AdaptiveAvgPool2d((1, 1)), + torch.nn.Flatten(), + torch.nn.Linear(4, 2), + ) + samples = torch.randn(2, 1, 2, 2) + + for split_layer in (1, -2): + extractor = LayeredModelExtractorBuilder.build(model, split_layer, device="cpu") + latent_data = extractor.input_to_latent(samples) + predictions = extractor.latent_to_logit(latent_data) + split_index = split_layer % len(model) + + assert torch.allclose(latent_data.activations, model[: split_index + 1](samples)) + assert torch.allclose(predictions, model(samples)) + + +@pytest.mark.parametrize("split_layer", ["layer", True, -3, 2]) +def test_torch_builder_validates_split_layer(split_layer): + model = torch.nn.Sequential(torch.nn.Linear(3, 2), torch.nn.ReLU()) + + with pytest.raises(ValueError, match="split_layer"): + LayeredModelExtractorBuilder.build(model, split_layer, device="cpu") + + +def test_generator_exits_no_grad_before_yielding(): + extractor = _identity_extractor(device="cpu") + generator = extractor.input_to_latent_generator(torch.ones((2, 3, 2, 2), requires_grad=True)) + + latent_data = next(generator) + value_created_by_the_consumer = torch.ones((), requires_grad=True) * 2 + + assert not latent_data.requires_grad + assert value_created_by_the_consumer.requires_grad + generator.close() + + +def test_device_selection_transfer_and_input_validation(): + import tensorflow as tf + + extractor = _identity_extractor() + + expected_device = "cuda" if torch.cuda.is_available() else "cpu" + assert extractor.device.type == expected_device + craft = HolisticCraftTorch(extractor) + assert craft.device == extractor.device + np.testing.assert_array_equal(craft._to_numpy(tf.constant([1.0])), [1.0]) + assert extractor.to("cpu").device.type == "cpu" + assert next(extractor.input_to_latent_generator(torch.ones((1, 3, 2, 2)))).device.type == "cpu" + + with pytest.raises(ValueError, match="Invalid PyTorch device"): + _identity_extractor(device="not-a-device") + if not torch.cuda.is_available(): + with pytest.raises(ValueError, match="CUDA was requested"): + _identity_extractor(device="cuda") + with pytest.raises(ValueError, match="rank 3"): + extractor.input_to_latent(torch.ones((2, 2))) + with pytest.raises(ValueError, match="at least one"): + list(extractor.input_to_latent_generator(torch.empty((0, 3, 2, 2)))) + with pytest.raises(ValueError, match="batch_size"): + _identity_extractor(batch_size=0) + + +def test_classifier_tensor_mode_keeps_existing_batch_dimension(): + model = torch.nn.Sequential( + torch.nn.Conv2d(1, 4, 1), + torch.nn.ReLU(), + torch.nn.AdaptiveAvgPool2d((1, 1)), + torch.nn.Flatten(), + torch.nn.Linear(4, 2), + ) + extractor = LayeredModelExtractorBuilder.build(model, split_layer=1, device="cpu") + extractor.output_as_tensor = True + + predictions = extractor(torch.ones((2, 1, 2, 2))) + + assert predictions.shape == (2, 2) diff --git a/tests/example_based/test_cole.py b/tests/example_based/test_cole.py index 7bdbec09..4e912114 100644 --- a/tests/example_based/test_cole.py +++ b/tests/example_based/test_cole.py @@ -155,6 +155,7 @@ def euclidean_dist(x, z): cases_dataset=x_train, targets_dataset=y_train, k=k, + batch_size=7, distance=euclidean_dist, projection=projection, ) diff --git a/tests/features_visualizations/test_maco.py b/tests/features_visualizations/test_maco.py index 1159ad66..d1234bc6 100644 --- a/tests/features_visualizations/test_maco.py +++ b/tests/features_visualizations/test_maco.py @@ -75,6 +75,24 @@ def test_init_maco_buffer(): assert phase.shape[1:] == spectrum_size +def test_init_maco_buffer_uses_packaged_spectrum(monkeypatch): + """Ensure the default ImageNet initialization does not require network access.""" + + def fail_download(*args, **kwargs): + del args, kwargs + raise AssertionError("init_maco_buffer must not download its default spectrum") + + monkeypatch.setattr(tf.keras.utils, "get_file", fail_download) + + magnitude, phase = init_maco_buffer((16, 24, 3)) + + assert magnitude.shape == (3, 16, 13) + assert phase.shape == (3, 16, 13) + assert magnitude.dtype == tf.float32 + assert phase.dtype == tf.float32 + assert np.all(np.isfinite(magnitude)) + + def test_maco_image_param(): """Ensure we can reconstruct an image from magnitude and phase""" img_size_to_magnitude_size = { diff --git a/tests/plots/test_metric_plots.py b/tests/plots/test_metric_plots.py index b7f7203b..2aa29c39 100644 --- a/tests/plots/test_metric_plots.py +++ b/tests/plots/test_metric_plots.py @@ -8,7 +8,7 @@ def test_bar_plot(): # test metrics barplot - cmap = matplotlib.cm.get_cmap("Set3") + cmap = matplotlib.colormaps["Set3"] methods_colors = {"method_" + str(i): cmap(i / 9) for i in range(10)} scores = {} @@ -31,7 +31,7 @@ def test_bar_plot(): def test_curves(): # test fidelity metric curves plot - cmap = matplotlib.cm.get_cmap("Set3") + cmap = matplotlib.colormaps["Set3"] methods_colors = {"method_" + str(i): cmap(i / 9) for i in range(10)} steps = np.linspace(0, 100, num=11) diff --git a/tests/plots/test_object_detection_plots.py b/tests/plots/test_object_detection_plots.py new file mode 100644 index 00000000..ddefe8f0 --- /dev/null +++ b/tests/plots/test_object_detection_plots.py @@ -0,0 +1,180 @@ +"""Tests for object-detection plotting utilities.""" + +import warnings + +import matplotlib +import numpy as np +import pytest +import tensorflow as tf + +matplotlib.use("Agg") +import matplotlib.pyplot as plt + +from xplique.plots.image import generate_heatmap +from xplique.plots.object_detection import plot_image_detections, plot_images_detections +from xplique.utils_functions.object_detection.base.box_manager import BoxFormat, BoxType +from xplique.utils_functions.object_detection.tf.multi_box_tensor import TfMultiBoxTensor + + +class _NumpyMultiBoxTensor: + def __init__(self, predictions): + self.predictions = predictions + + def boxes(self): + return self.predictions[:, :4] + + def scores(self): + return self.predictions[:, 4] + + def probas(self): + return self.predictions[:, 5:] + + +def _make_multibox(boxes: np.ndarray, num_classes: int = 2) -> TfMultiBoxTensor: + """Build a minimal TfMultiBoxTensor: (N, 4 + 1 + num_classes).""" + n = len(boxes) + scores = np.ones((n, 1), dtype=np.float32) + probas = np.zeros((n, num_classes), dtype=np.float32) + probas[:, 0] = 1.0 + tensor = np.concatenate([boxes, scores, probas], axis=1) + return TfMultiBoxTensor(tf.constant(tensor)) + + +CLASSES = ["cat", "dog"] +COLORS = {"cat": "red", "dog": "blue"} + + +def test_pixel_boxes_correct_flag(): + """Pixel-coord boxes with is_normalized=False (default) must plot without error.""" + image = np.zeros((64, 64, 3), dtype=np.uint8) + boxes = np.array([[10.0, 10.0, 50.0, 50.0]], dtype=np.float32) + + plot_image_detections(image, _make_multibox(boxes), CLASSES, COLORS) + matplotlib.pyplot.close("all") + + +def test_normalized_boxes_correct_flag(): + """Normalized boxes [0, 1] with is_normalized=True must plot without error.""" + image = np.zeros((64, 64, 3), dtype=np.uint8) + boxes = np.array([[0.1, 0.1, 0.9, 0.9]], dtype=np.float32) + + box_type = BoxType(BoxFormat.XYXY, is_normalized=True) + plot_image_detections(image, _make_multibox(boxes), CLASSES, COLORS, box_type=box_type) + matplotlib.pyplot.close("all") + + +def test_pixel_boxes_declared_as_normalized_raises(): + """ + Pixel-coord boxes declared as is_normalized=True must raise ValueError. + + The translator skips normalization (assumes [0,1]) then denormalizes by + multiplying by image size, producing values far beyond the image bounds. + E.g. xmax=60 on a 64px image -> 60 * 64 = 3840, which exceeds 64 * 3 = 192. + """ + image = np.zeros((64, 64, 3), dtype=np.uint8) + boxes = np.array([[10.0, 10.0, 60.0, 60.0]], dtype=np.float32) + + box_type = BoxType(BoxFormat.XYXY, is_normalized=True) + # wrong: pixel coords passed as normalized + + with pytest.raises(ValueError, match="far exceed image dimensions"): + plot_image_detections(image, _make_multibox(boxes), CLASSES, COLORS, box_type=box_type) + + +def test_normalized_boxes_declared_as_pixel_warns(): + """ + Normalized boxes declared as is_normalized=False must issue UserWarning. + + The check fires on the raw input coords before translation: all values are + below 1.0, which suggests is_normalized=False may be incorrect. A warning + is used because small pixel-coordinate boxes can also be valid. + """ + image = np.zeros((64, 64, 3), dtype=np.uint8) + boxes = np.array([[0.1, 0.1, 0.9, 0.9]], dtype=np.float32) + + box_type = BoxType(BoxFormat.XYXY, is_normalized=False) + # wrong: normalized coords passed as pixel + + with pytest.warns(UserWarning, match="are all below 1.0"): + plot_image_detections(image, _make_multibox(boxes), CLASSES, COLORS, box_type=box_type) + matplotlib.pyplot.close("all") + + +def test_normalized_boxes_at_warning_limit_does_not_warn(): + """max coord = 1.0 is not below the warning threshold -> no warning.""" + image = np.zeros((64, 64, 3), dtype=np.uint8) + boxes = np.array([[0.1, 0.1, 1.0, 1.0]], dtype=np.float32) + + box_type = BoxType(BoxFormat.XYXY, is_normalized=False) + + with warnings.catch_warnings(record=True) as recorded: + warnings.simplefilter("always") + plot_image_detections(image, _make_multibox(boxes), CLASSES, COLORS, box_type=box_type) + assert not recorded + matplotlib.pyplot.close("all") + + +def test_normalized_boxes_above_warning_limit_does_not_warn(): + """max coord = 1.1 is above the warning threshold -> no warning.""" + image = np.zeros((64, 64, 3), dtype=np.uint8) + boxes = np.array([[0.1, 0.1, 1.1, 1.1]], dtype=np.float32) + + box_type = BoxType(BoxFormat.XYXY, is_normalized=False) + + with warnings.catch_warnings(record=True) as recorded: + warnings.simplefilter("always") + plot_image_detections(image, _make_multibox(boxes), CLASSES, COLORS, box_type=box_type) + assert not recorded + matplotlib.pyplot.close("all") + + +def test_generate_heatmap_resizes_rank_two_attribution(): + heatmap = generate_heatmap( + np.arange(6, dtype=np.float32).reshape(2, 3), + size=(8, 12), + clip_percentile=None, + normalize_value=False, + ) + + assert heatmap.shape == (8, 12) + + +def test_detection_heatmap_uses_the_image_coordinate_extent(): + image = np.zeros((20, 40, 3), dtype=np.float32) + detections = _NumpyMultiBoxTensor(np.array([[5, 5, 15, 10, 0.9, 1.0, 0.0]])) + + fig = plot_image_detections( + image, + detections, + classes_labels=["object", "other"], + label_to_color={"object": "red"}, + heatmap=np.zeros((2, 2), dtype=np.float32), + ) + try: + image_artist, heatmap_artist = fig.axes[0].images + assert heatmap_artist.get_extent() == image_artist.get_extent() + finally: + plt.close(fig) + + +def test_plot_images_detections_rejects_empty_lists(): + with pytest.raises(ValueError, match="at least one element"): + plot_images_detections( + [], + [], + classes_labels=["object", "other"], + label_to_color={"object": "red"}, + ) + + +def test_plot_images_detections_accepts_numpy_image_batch(): + images = np.zeros((1, 20, 40, 3), dtype=np.float32) + detections = _NumpyMultiBoxTensor(np.array([[5, 5, 15, 10, 0.9, 1.0, 0.0]])) + + fig = plot_images_detections( + images, + [detections], + classes_labels=["object", "other"], + label_to_color={"object": "red"}, + ) + plt.close(fig) diff --git a/tests/utils_functions/test_box_manager.py b/tests/utils_functions/test_box_manager.py new file mode 100644 index 00000000..16107653 --- /dev/null +++ b/tests/utils_functions/test_box_manager.py @@ -0,0 +1,238 @@ +""" +Tests for BoxManager and BoxCoordinatesTranslator across NumPy, PyTorch, and TensorFlow backends. +""" + +import numpy as np +import pytest + +from xplique.utils_functions.object_detection.base.box_manager import ( + BoxFormat, + BoxType, + NumpyBoxCoordinatesTranslator, + NumpyBoxManager, +) + +try: + import torch + + from xplique.utils_functions.object_detection.torch.box_manager import ( + TorchBoxCoordinatesTranslator, + TorchBoxManager, + ) + + HAS_TORCH = True +except ImportError: + HAS_TORCH = False + +import tensorflow as tf + +from xplique.utils_functions.object_detection.tf.box_manager import ( + TfBoxCoordinatesTranslator, + TfBoxManager, +) + + +class BaseBoxManagerTests: + """Shared tests for all backend BoxManager implementations.""" + + box_manager_cls = None + translator_cls = None + + def make_tensor(self, data): + raise NotImplementedError + + def make_image_size(self, size): + """Return a backend-appropriate representation of a square (size x size) image.""" + raise NotImplementedError + + def allclose(self, a, b): + raise NotImplementedError + + # --- Unit tests --- + + def test_normalize_boxes(self): + raw_boxes = self.make_tensor([[50, 50, 100, 100], [30, 30, 60, 60]]) + result = self.box_manager_cls.normalize_boxes(raw_boxes, self.make_image_size(200)) + expected = self.make_tensor([[0.25, 0.25, 0.5, 0.5], [0.15, 0.15, 0.3, 0.3]]) + assert self.allclose(result, expected) + + def test_box_cxcywh_to_xyxy(self): + boxes = self.make_tensor([[50, 50, 20, 20], [0.3, 0.3, 0.1, 0.1]]) + result = self.box_manager_cls.box_cxcywh_to_xyxy(boxes) + expected = self.make_tensor([[40, 40, 60, 60], [0.25, 0.25, 0.35, 0.35]]) + assert self.allclose(result, expected) + + def test_box_xyxy_to_cxcywh(self): + boxes = self.make_tensor([[40, 40, 60, 60], [0.25, 0.25, 0.35, 0.35]]) + result = self.box_manager_cls.box_xyxy_to_cxcywh(boxes) + expected = self.make_tensor([[50, 50, 20, 20], [0.3, 0.3, 0.1, 0.1]]) + assert self.allclose(result, expected) + + def test_box_xywh_to_xyxy(self): + boxes = self.make_tensor([[40, 40, 20, 20], [0.25, 0.25, 0.1, 0.1]]) + result = self.box_manager_cls.box_xywh_to_xyxy(boxes) + expected = self.make_tensor([[40, 40, 60, 60], [0.25, 0.25, 0.35, 0.35]]) + assert self.allclose(result, expected) + + def test_box_conversion_preserves_trailing_columns(self): + boxes = self.make_tensor([[40, 40, 20, 20, 0.8, 1.0]]) + result = self.box_manager_cls.box_xywh_to_xyxy(boxes) + expected = self.make_tensor([[40, 40, 60, 60, 0.8, 1.0]]) + assert self.allclose(result, expected) + + def test_box_xyxy_to_xywh(self): + boxes = self.make_tensor([[40, 40, 60, 60], [0.25, 0.25, 0.35, 0.35]]) + result = self.box_manager_cls.box_xyxy_to_xywh(boxes) + expected = self.make_tensor([[40, 40, 20, 20], [0.25, 0.25, 0.1, 0.1]]) + assert self.allclose(result, expected) + + def test_denormalize_boxes(self): + boxes = self.make_tensor([[0.25, 0.25, 0.5, 0.5], [0.15, 0.15, 0.3, 0.3]]) + result = self.box_manager_cls.denormalize_boxes(boxes, self.make_image_size(400)) + expected = self.make_tensor([[100, 100, 200, 200], [60, 60, 120, 120]]) + assert self.allclose(result, expected) + + def test_denormalize_boxes_with_tuple(self): + """TF-specific: denormalize accepts a plain Python tuple (like PIL Image.size).""" + boxes = self.make_tensor([[0.25, 0.25, 0.5, 0.5]]) + result = self.box_manager_cls.denormalize_boxes(boxes, (640, 480)) + expected = self.make_tensor([[160, 120, 320, 240]]) + assert self.allclose(result, expected) + + # --- Integration tests for BoxCoordinatesTranslator --- + + def test_translator_normalized_cxcywh_to_normalized_xyxy(self): + """DETR typical case: normalized CXCYWH -> normalized XYXY.""" + translator = self.translator_cls( + input_box_type=BoxType(BoxFormat.CXCYWH, is_normalized=True), + output_box_type=BoxType(BoxFormat.XYXY, is_normalized=True), + ) + boxes = self.make_tensor([[0.3, 0.4, 0.1, 0.2]]) + result = translator.translate(boxes) + expected = self.make_tensor([[0.25, 0.3, 0.35, 0.5]]) + assert self.allclose(result, expected) + + def test_translator_pixel_xyxy_to_normalized_xyxy(self): + """FCOS typical case: pixel XYXY -> normalized XYXY.""" + translator = self.translator_cls( + input_box_type=BoxType(BoxFormat.XYXY, is_normalized=False), + output_box_type=BoxType(BoxFormat.XYXY, is_normalized=True), + ) + boxes = self.make_tensor([[50, 50, 100, 100]]) + result = translator.translate(boxes, image_size=self.make_image_size(200)) + expected = self.make_tensor([[0.25, 0.25, 0.5, 0.5]]) + assert self.allclose(result, expected) + + def test_translator_normalized_xyxy_to_normalized_cxcywh(self): + translator = self.translator_cls( + input_box_type=BoxType(BoxFormat.XYXY, is_normalized=True), + output_box_type=BoxType(BoxFormat.CXCYWH, is_normalized=True), + ) + boxes = self.make_tensor([[0.25, 0.3, 0.35, 0.5]]) + result = translator.translate(boxes) + expected = self.make_tensor([[0.3, 0.4, 0.1, 0.2]]) + assert self.allclose(result, expected) + + def test_translator_normalized_xyxy_to_pixel_xyxy(self): + translator = self.translator_cls( + input_box_type=BoxType(BoxFormat.XYXY, is_normalized=True), + output_box_type=BoxType(BoxFormat.XYXY, is_normalized=False), + ) + boxes = self.make_tensor([[0.25, 0.25, 0.5, 0.5]]) + result = translator.translate(boxes, image_size=self.make_image_size(400)) + expected = self.make_tensor([[100, 100, 200, 200]]) + assert self.allclose(result, expected) + + +class TestNumpyBoxManager(BaseBoxManagerTests): + box_manager_cls = NumpyBoxManager + translator_cls = NumpyBoxCoordinatesTranslator + + def make_tensor(self, data): + return np.array(data, dtype=np.float32) + + def make_image_size(self, size): + return (size, size) + + def allclose(self, a, b): + return np.allclose(a, b) + + +@pytest.mark.skipif(not HAS_TORCH, reason="PyTorch not available") +class TestTorchBoxManager(BaseBoxManagerTests): + box_manager_cls = TorchBoxManager if HAS_TORCH else None + translator_cls = TorchBoxCoordinatesTranslator if HAS_TORCH else None + + def make_tensor(self, data): + return torch.tensor(data, dtype=torch.float32) + + def make_image_size(self, size): + return torch.Size([size, size]) + + def allclose(self, a, b): + return torch.allclose(a, b) + + +class TestTfBoxManager(BaseBoxManagerTests): + box_manager_cls = TfBoxManager + translator_cls = TfBoxCoordinatesTranslator + + def make_tensor(self, data): + return tf.constant(data, dtype=tf.float32) + + def make_image_size(self, size): + return tf.constant([size, size], dtype=tf.float32) + + def allclose(self, a, b): + return bool(tf.reduce_all(tf.abs(a - b) < 1e-5)) + + +def test_numpy_box_manager_handles_integer_boxes_and_metadata(): + boxes = np.array([[50, 25, 100, 50, 1, 2]], dtype=np.int32) + + normalized = NumpyBoxManager.normalize_boxes(boxes, (200, 100)) + denormalized = NumpyBoxManager.denormalize_boxes(normalized, (200, 100)) + + assert np.issubdtype(normalized.dtype, np.floating) + np.testing.assert_allclose(normalized, [[0.25, 0.25, 0.5, 0.5, 1.0, 2.0]]) + np.testing.assert_allclose(denormalized, [[50, 25, 100, 50, 1.0, 2.0]]) + + +@pytest.mark.skipif(not HAS_TORCH, reason="PyTorch not available") +def test_torch_box_manager_handles_integer_boxes_and_metadata(): + boxes = torch.tensor([[50, 25, 100, 50, 1, 2]], dtype=torch.int64) + + normalized = TorchBoxManager.normalize_boxes(boxes, (200, 100)) + denormalized = TorchBoxManager.denormalize_boxes(normalized, (200, 100)) + + assert normalized.dtype == torch.float32 + assert torch.allclose(normalized, torch.tensor([[0.25, 0.25, 0.5, 0.5, 1.0, 2.0]])) + assert torch.allclose(denormalized, torch.tensor([[50, 25, 100, 50, 1.0, 2.0]])) + + +def test_tf_box_manager_handles_integer_boxes_and_metadata_in_a_graph(): + @tf.function + def normalize_and_denormalize(boxes): + normalized = TfBoxManager.normalize_boxes(boxes, (200, 100)) + return normalized, TfBoxManager.denormalize_boxes(normalized, (200, 100)) + + boxes = tf.constant([[50, 25, 100, 50, 1, 2]], dtype=tf.int32) + normalized, denormalized = normalize_and_denormalize(boxes) + + assert normalized.dtype == tf.float32 + np.testing.assert_allclose(normalized, [[0.25, 0.25, 0.5, 0.5, 1.0, 2.0]]) + np.testing.assert_allclose(denormalized, [[50, 25, 100, 50, 1.0, 2.0]]) + + +def test_numpy_translator_uses_enum_members_for_format_checks(): + translator = NumpyBoxCoordinatesTranslator( + input_box_type=BoxType(BoxFormat.CXCYWH, is_normalized=True), + output_box_type=BoxType(BoxFormat.XYXY, is_normalized=True), + ) + + translator.input_box_type.format = BoxFormat.XYWH.value + boxes = np.array([[0.3, 0.4, 0.1, 0.2]], dtype=np.float32) + result = translator.translate(boxes) + + np.testing.assert_allclose(result, boxes) + assert not np.allclose(result, [[0.25, 0.3, 0.35, 0.5]]) diff --git a/tests/utils_functions/test_box_model_wrapper.py b/tests/utils_functions/test_box_model_wrapper.py new file mode 100644 index 00000000..14914fac --- /dev/null +++ b/tests/utils_functions/test_box_model_wrapper.py @@ -0,0 +1,135 @@ +"""Regression tests for padded object-detection tensor outputs.""" + +import pytest +import tensorflow as tf + +from xplique.concepts.tf.holistic_craft import ConceptDecoderTf +from xplique.concepts.tf.latent_extractor import TfLatentExtractor +from xplique.utils_functions.object_detection.tf.box_model_wrapper import TfBoxesModelWrapper +from xplique.utils_functions.object_detection.tf.multi_box_tensor import TfMultiBoxTensor + +try: + import torch + + from xplique.concepts.torch.holistic_craft import ConceptDecoderTorch + from xplique.concepts.torch.latent_extractor import TorchLatentExtractor + from xplique.utils_functions.object_detection.torch.box_model_wrapper import ( + TorchBoxesModelWrapper, + ) + + HAS_TORCH = True +except ImportError: + HAS_TORCH = False + + +class _TfVariableDetectionFormatter: + def __call__(self, predictions): + return [TfMultiBoxTensor(predictions[0, :0]), TfMultiBoxTensor(predictions[1, :3])] + + +class _TfDecoderParent: + @staticmethod + def _decode_coefficients(_latent_data, coefficients): + return [ + TfMultiBoxTensor(coefficients[0, :1]), + TfMultiBoxTensor(coefficients[1, :3]), + ] + + +def test_tf_box_wrapper_pads_variable_detection_counts(): + wrapper = TfBoxesModelWrapper( + tf.keras.layers.Lambda(lambda value: value), _TfVariableDetectionFormatter() + ) + wrapper.output_as_tensor = True + + @tf.function + def run(inputs): + return wrapper(inputs) + + predictions = run(tf.ones((2, 3, 7))) + + assert predictions.shape == (2, 3, 7) + tf.debugging.assert_equal(predictions[0], tf.zeros((3, 7))) + + +def test_tf_latent_extractor_pads_variable_detection_counts(): + extractor = TfLatentExtractor( + model=tf.keras.layers.Lambda(lambda value: value), + input_to_latent_model=lambda samples: samples[:, 0], + latent_to_logit_model=lambda latent_data: latent_data, + output_formatter=_TfVariableDetectionFormatter(), + ) + extractor.output_as_tensor = True + + predictions = extractor(tf.ones((2, 1, 3, 7))) + + assert predictions.shape == (2, 3, 7) + tf.debugging.assert_equal(predictions[0], tf.zeros((3, 7))) + + +def test_tf_concept_decoder_pads_batched_variable_detection_counts(): + decoder = ConceptDecoderTf(_TfDecoderParent(), latent_data=None) + coefficients = tf.Variable(tf.ones((2, 3, 7))) + + with tf.GradientTape() as tape: + predictions = decoder(coefficients) + loss = tf.reduce_sum(predictions) + gradients = tape.gradient(loss, coefficients) + + assert predictions.shape == (2, 3, 7) + tf.debugging.assert_equal(predictions[0, 1:], tf.zeros((2, 7))) + assert gradients is not None + tf.debugging.assert_positive(tf.reduce_sum(tf.abs(gradients))) + + +@pytest.mark.skipif(not HAS_TORCH, reason="PyTorch not available") +def test_torch_box_wrapper_pads_variable_detection_counts(): + def formatter(predictions): + return [predictions[0, :0], predictions[1, :3]] + + wrapper = TorchBoxesModelWrapper(torch.nn.Identity(), formatter) + wrapper.output_as_tensor = True + + predictions = wrapper(torch.ones((2, 3, 7))) + + assert predictions.shape == (2, 3, 7) + assert torch.equal(predictions[0], torch.zeros((3, 7))) + + +@pytest.mark.skipif(not HAS_TORCH, reason="PyTorch not available") +def test_torch_latent_extractor_pads_variable_detection_counts(): + def formatter(predictions): + return [predictions[0, :0], predictions[1, :3]] + + extractor = TorchLatentExtractor( + model=torch.nn.Identity(), + input_to_latent_model=lambda samples: samples[:, 0], + latent_to_logit_model=lambda latent_data: latent_data, + output_formatter=formatter, + device="cpu", + ) + extractor.output_as_tensor = True + + predictions = extractor(torch.ones((2, 1, 3, 7))) + + assert predictions.shape == (2, 3, 7) + assert torch.equal(predictions[0], torch.zeros((3, 7))) + + +@pytest.mark.skipif(not HAS_TORCH, reason="PyTorch not available") +def test_torch_concept_decoder_pads_batched_variable_detection_counts(): + class DecoderParent: + @staticmethod + def _decode_coefficients(_latent_data, coefficients): + return [coefficients[0, :1], coefficients[1, :3]] + + decoder = ConceptDecoderTorch(DecoderParent(), latent_data=None) + coefficients = torch.ones((2, 3, 7), requires_grad=True) + + predictions = decoder(coefficients) + predictions.sum().backward() + + assert predictions.shape == (2, 3, 7) + assert torch.equal(predictions[0, 1:], torch.zeros((2, 7))) + assert coefficients.grad is not None + assert torch.sum(torch.abs(coefficients.grad)) > 0 diff --git a/tests/utils_functions/test_box_utilities.py b/tests/utils_functions/test_box_utilities.py new file mode 100644 index 00000000..9d5520d8 --- /dev/null +++ b/tests/utils_functions/test_box_utilities.py @@ -0,0 +1,77 @@ +"""Focused tests for object-detection tensor and formatter utilities.""" + +import pytest +import tensorflow as tf + +from xplique.utils_functions.object_detection.base.box_formatter import BaseBoxFormatter +from xplique.utils_functions.object_detection.base.box_manager import BoxFormat, BoxType +from xplique.utils_functions.object_detection.tf.box_formatter import TfBaseBoxFormatter +from xplique.utils_functions.object_detection.tf.multi_box_tensor import TfMultiBoxTensor + +try: + import torch + + from xplique.utils_functions.object_detection.torch.box_formatter import TorchBaseBoxFormatter + + HAS_TORCH = True +except ImportError: + HAS_TORCH = False + + +class _TfFormatter(TfBaseBoxFormatter): + def forward(self, predictions): + return predictions + + +class _ForwardOnlyFormatter(BaseBoxFormatter): + def forward(self, predictions): + return predictions + + +def test_tf_multibox_tensor_length_is_detection_count(): + assert len(TfMultiBoxTensor(tf.zeros((0, 7)))) == 0 + assert len(TfMultiBoxTensor(tf.zeros((3, 7)))) == 3 + + +def test_base_formatter_only_requires_forward(): + formatter = _ForwardOnlyFormatter(BoxType(BoxFormat.XYXY, is_normalized=True)) + assert formatter({"predictions": "ok"}) == {"predictions": "ok"} + + with pytest.raises(NotImplementedError): + formatter.format_predictions({"predictions": "ok"}) + + +def test_tf_formatter_expands_rank_one_scores(): + formatter = _TfFormatter(BoxType(BoxFormat.XYXY, is_normalized=True)) + formatted = formatter.format_predictions( + { + "boxes": tf.constant([[0.1, 0.2, 0.3, 0.4]]), + "scores": tf.constant([0.9]), + "probas": tf.constant([[0.2, 0.8]]), + } + ) + + assert formatted.shape == (1, 7) + tf.debugging.assert_near(formatted.scores(), [0.9]) + + +if HAS_TORCH: + + class _TorchFormatter(TorchBaseBoxFormatter): + def forward(self, predictions): + return predictions + + +@pytest.mark.skipif(not HAS_TORCH, reason="PyTorch not available") +def test_torch_formatter_expands_rank_one_scores(): + formatter = _TorchFormatter(BoxType(BoxFormat.XYXY, is_normalized=True)) + formatted = formatter.format_predictions( + { + "boxes": torch.tensor([[0.1, 0.2, 0.3, 0.4]]), + "scores": torch.tensor([0.9]), + "probas": torch.tensor([[0.2, 0.8]]), + } + ) + + assert formatted.shape == (1, 7) + assert torch.allclose(formatted.scores(), torch.tensor([0.9])) diff --git a/tests/utils_functions/test_classifier_tensor.py b/tests/utils_functions/test_classifier_tensor.py new file mode 100644 index 00000000..34fdc235 --- /dev/null +++ b/tests/utils_functions/test_classifier_tensor.py @@ -0,0 +1,74 @@ +"""Tests for classifier tensor wrappers across TensorFlow and PyTorch.""" + +import numpy as np +import pytest +import tensorflow as tf + +from xplique.utils_functions.classification.tf.classifier_tensor import TfClassifierTensor + +try: + import torch + + from xplique.utils_functions.classification.torch.classifier_tensor import TorchClassifierTensor + + HAS_TORCH = True +except ImportError: + HAS_TORCH = False + + +def test_tf_classifier_tensor_validates_rank_and_returns_tf_tensor_batch(): + with pytest.raises(ValueError, match="rank 1 or 2"): + TfClassifierTensor(tf.ones((1, 2, 3))) + + classifier = TfClassifierTensor(np.array([0.1, 0.9], dtype=np.float32)) + batched = classifier.to_batched_tensor() + assert isinstance(batched, tf.Tensor) + assert batched.shape == (1, 2) + + +def test_tf_classifier_tensor_len_deprecation_and_properties(): + classifier = TfClassifierTensor(tf.constant([[0.1, 0.9], [0.3, 0.7]], dtype=tf.float32)) + + with pytest.warns(DeprecationWarning): + assert len(classifier) == 2 + + assert classifier.num_classes == 2 + assert classifier.batch_size == 2 + assert not classifier.is_empty + + +def test_tf_classifier_tensor_empty_reports_is_empty(): + classifier = TfClassifierTensor(tf.zeros((0, 3), dtype=tf.float32)) + assert classifier.batch_size == 0 + assert classifier.is_empty + + +@pytest.mark.skipif(not HAS_TORCH, reason="PyTorch not available") +def test_torch_classifier_tensor_validates_rank_and_properties(): + with pytest.raises(ValueError, match="rank 1 or 2"): + TorchClassifierTensor.from_predictions(torch.ones((1, 2, 3))) + + classifier = TorchClassifierTensor.from_predictions(torch.tensor([[0.1, 0.9], [0.3, 0.7]])) + with pytest.warns(DeprecationWarning): + assert len(classifier) == 2 + + assert classifier.num_classes == 2 + assert classifier.batch_size == 2 + assert not classifier.is_empty + + +@pytest.mark.skipif(not HAS_TORCH, reason="PyTorch not available") +def test_torch_classifier_tensor_empty_reports_is_empty(): + classifier = TorchClassifierTensor.from_predictions(torch.zeros((0, 3))) + assert classifier.batch_size == 0 + assert classifier.is_empty + + +@pytest.mark.skipif(not HAS_TORCH, reason="PyTorch not available") +def test_tf_and_torch_classifier_tensor_properties_match_on_same_input(): + values = np.array([[0.1, 0.2, 0.7], [0.4, 0.5, 0.1]], dtype=np.float32) + tf_classifier = TfClassifierTensor(values) + torch_classifier = TorchClassifierTensor.from_predictions(values) + + assert tf_classifier.num_classes == torch_classifier.num_classes + assert tf_classifier.batch_size == torch_classifier.batch_size diff --git a/tests/utils_functions/test_gradients_check_tf.py b/tests/utils_functions/test_gradients_check_tf.py new file mode 100644 index 00000000..22f3bb67 --- /dev/null +++ b/tests/utils_functions/test_gradients_check_tf.py @@ -0,0 +1,388 @@ +""" +Tests for TensorFlow gradient checking utilities. + +This module tests the check_model_gradients function with various output formats +to ensure it correctly detects gradient flow in different model architectures. +""" +# pylint: disable=redefined-outer-name + +import pytest +import tensorflow as tf + +from xplique.utils_functions.common.tf.gradients_check import check_model_gradients +from xplique.utils_functions.object_detection.tf.multi_box_tensor import ( + TfMultiBoxTensor as MultiBoxTensor, +) + + +class SimpleTensorModel(tf.keras.Model): + """Model that returns a single tensor.""" + + def __init__(self, has_gradient=True): + super().__init__() + self.has_gradient = has_gradient + self.dense = tf.keras.layers.Dense(5) + + def call(self, x): + """Forward pass of the model.""" + if self.has_gradient: + # Process input to create gradient path + pooled = tf.reduce_mean(x, axis=[1, 2]) # (batch, channels) + return self.dense(pooled) + # No gradient path - return constant + return tf.ones([tf.shape(x)[0], 5]) + + +class DictOutputModel(tf.keras.Model): + """Model that returns a dict of tensors.""" + + def __init__(self): + super().__init__() + self.dense = tf.keras.layers.Dense(5) + + def call(self, x): + """Forward pass returning dict.""" + pooled = tf.reduce_mean(x, axis=[1, 2]) + out = self.dense(pooled) + return {"scores": out, "features": out * 2} + + +class ListOutputModel(tf.keras.Model): + """Model that returns a list of tensors.""" + + def __init__(self): + super().__init__() + self.dense = tf.keras.layers.Dense(5) + + def call(self, x): + """Forward pass returning list.""" + pooled = tf.reduce_mean(x, axis=[1, 2]) + out = self.dense(pooled) + return [out, out * 2, out * 3] + + +class NestedDictModel(tf.keras.Model): + """Model that returns nested dict structure.""" + + def __init__(self): + super().__init__() + self.dense = tf.keras.layers.Dense(5) + + def call(self, x): + """Forward pass returning nested dict.""" + pooled = tf.reduce_mean(x, axis=[1, 2]) + out = self.dense(pooled) + return {"predictions": {"scores": out, "features": out * 2}, "meta": out * 3} + + +class DictWithListModel(tf.keras.Model): + """Model that returns dict containing lists of tensors.""" + + def __init__(self): + super().__init__() + self.dense = tf.keras.layers.Dense(5) + + def call(self, x): + """Forward pass returning dict with list values.""" + pooled = tf.reduce_mean(x, axis=[1, 2]) + out = self.dense(pooled) + return {"predictions": [out, out * 2], "meta": out * 3} + + +class MixedListModel(tf.keras.Model): + """Model that returns list with both tensors and dicts.""" + + def __init__(self): + super().__init__() + self.dense = tf.keras.layers.Dense(5) + + def call(self, x): + """Forward pass returning mixed list.""" + pooled = tf.reduce_mean(x, axis=[1, 2]) + out = self.dense(pooled) + return [out, {"scores": out * 2, "features": out * 3}, out * 4] + + +class MultiBoxModel(tf.keras.Model): + """Model that returns a MultiBoxTensor object.""" + + def __init__(self): + super().__init__() + self.dense = tf.keras.layers.Dense(85) # 4 boxes + 1 score + 80 classes + + def call(self, x): + """Forward pass returning MultiBoxTensor.""" + pooled = tf.reduce_mean(x, axis=[1, 2]) # (batch, channels) + out = self.dense(pooled) + + # Reshape to (batch * num_boxes, features) + reshaped = tf.reshape(out, [-1, 85]) + + # Create MultiBoxTensor (wraps the tensor with .tensor attribute) + return MultiBoxTensor(reshaped) + + +@pytest.fixture +def input_tensor(): # pylint: disable=redefined-outer-name + """Create a simple input tensor.""" + return tf.random.normal([2, 32, 32, 10]) + + +def test_simple_tensor_with_gradient(input_tensor): + """Test gradient checking with a simple tensor output.""" + model = SimpleTensorModel(has_gradient=True) + result = check_model_gradients(model, input_tensor) + assert result, "Gradients should flow through simple tensor model" + + +def test_simple_tensor_without_gradient(input_tensor): + """Test gradient checking when no gradient flows.""" + model = SimpleTensorModel(has_gradient=False) + result = check_model_gradients(model, input_tensor) + assert not result, "Should detect when no gradients flow" + + +def test_dict_output(input_tensor): + """Test gradient checking with dict output.""" + model = DictOutputModel() + result = check_model_gradients(model, input_tensor) + assert result, "Gradients should flow through dict output model" + + +def test_list_output(input_tensor): + """Test gradient checking with list output.""" + model = ListOutputModel() + result = check_model_gradients(model, input_tensor) + assert result, "Gradients should flow through list output model" + + +def test_nested_dict(input_tensor): + """Test gradient checking with nested dict structure.""" + model = NestedDictModel() + result = check_model_gradients(model, input_tensor) + assert result, "Gradients should flow through nested dict model" + + +def test_dict_with_list(input_tensor): + """Test gradient checking with dict containing lists.""" + model = DictWithListModel() + result = check_model_gradients(model, input_tensor) + assert result, "Gradients should flow through dict with list model" + + +def test_mixed_list(input_tensor): + """Test gradient checking with mixed list (tensors and dicts).""" + model = MixedListModel() + result = check_model_gradients(model, input_tensor) + assert result, "Gradients should flow through mixed list model" + + +def test_callable_function(input_tensor): + """Test gradient checking with a callable function instead of Model.""" + + def simple_func(x): + return tf.reduce_mean(x, axis=[1, 2]) + + result = check_model_gradients(simple_func, input_tensor) + assert result, "Gradients should flow through simple callable" + + +def test_empty_dict(input_tensor): + """Test gradient checking with model returning empty dict.""" + + class EmptyDictModel(tf.keras.Model): + """Model returning empty dict.""" + + def call(self, _): + """Forward pass returning empty dict.""" + return {} + + model = EmptyDictModel() + result = check_model_gradients(model, input_tensor) + assert not result, "Should detect no tensors in empty dict" + + +def test_empty_list(input_tensor): + """Test gradient checking with model returning empty list.""" + + class EmptyListModel(tf.keras.Model): + """Model returning empty list.""" + + def call(self, _): + """Forward pass returning empty list.""" + return [] + + model = EmptyListModel() + result = check_model_gradients(model, input_tensor) + assert not result, "Should detect no tensors in empty list" + + +def test_multibox_tensor_output(input_tensor): + """Test gradient checking with model that returns MultiBoxTensor.""" + model = MultiBoxModel() + result = check_model_gradients(model, input_tensor) + assert result, "Gradients should flow through MultiBoxTensor model" + + +def test_dict_with_non_tensor_values(input_tensor): + """Test gradient checking with dict containing non-tensor values.""" + + class MixedDictModel(tf.keras.Model): + """Model returning dict with mixed value types.""" + + def __init__(self): + super().__init__() + self.dense = tf.keras.layers.Dense(5) + + def call(self, x): + """Forward pass returning dict with mixed values.""" + pooled = tf.reduce_mean(x, axis=[1, 2]) + out = self.dense(pooled) + return {"scores": out, "metadata": "some string", "count": 42} + + model = MixedDictModel() + result = check_model_gradients(model, input_tensor) + assert result, "Should still detect gradients from tensor values" + + +def test_list_with_non_tensor_values(input_tensor): + """Test gradient checking with list containing non-tensor values.""" + + class MixedListWithNonTensorsModel(tf.keras.Model): + """Model returning list with mixed value types.""" + + def __init__(self): + super().__init__() + self.dense = tf.keras.layers.Dense(5) + + def call(self, x): + """Forward pass returning list with mixed values.""" + pooled = tf.reduce_mean(x, axis=[1, 2]) + out = self.dense(pooled) + return [out, "metadata", 42, out * 2] + + model = MixedListWithNonTensorsModel() + result = check_model_gradients(model, input_tensor) + assert result, "Gradients should flow through mixed list model" + + +def test_model_with_input_preprocessing_breaking_gradients(input_tensor): + """Test gradient checking with model that breaks gradients during preprocessing. + + This replicates DETR's issue: inputs are preprocessed into a new tensor format + where the gradient connection to the original inputs is severed. + + Pattern (like DETR's nested_tensor_from_tensor_list): + 1. Input comes in with gradient tracking enabled (for attribution) + 2. Preprocessing uses tf.stop_gradient to break gradient connection + 3. Model processes the gradient-disconnected tensor + 4. Gradient path is broken - backward pass can't reach original inputs + + Note: In TensorFlow, we explicitly use tf.stop_gradient to break the gradient + flow, simulating what happens in PyTorch when torch.zeros() creates a tensor + with requires_grad=False. + """ + + class PreprocessingModel(tf.keras.Model): + """Model with preprocessing that breaks gradients.""" + + def __init__(self): + super().__init__() + self.conv = tf.keras.layers.Conv2D(5, 3, padding="same") + + def call(self, x): + """Forward pass with gradient-breaking preprocessing.""" + # Break gradient connection at preprocessing step + x_no_grad = tf.stop_gradient(x) + out = self.conv(x_no_grad) + return tf.reduce_mean(out, axis=[1, 2]) + + model = PreprocessingModel() + result = check_model_gradients(model, input_tensor) + assert not result, ( + "Should detect broken gradients due to input preprocessing with tf.stop_gradient" + ) + + +def test_model_with_backward_error(input_tensor): + """Test gradient checking with model that raises error during backward. + + This simulates cases where gradient computation fails due to model architecture + issues, numerical problems, or other errors during backpropagation. + + The check_model_gradients function should return False and catch the error. + """ + + class BackwardErrorModel(tf.keras.Model): + """Model that breaks gradients during backward.""" + + def __init__(self): + super().__init__() + self.dense = tf.keras.layers.Dense(5) + + def call(self, x): + """Forward pass with detached output.""" + pooled = tf.reduce_mean(x, axis=[1, 2]) + out = self.dense(pooled) + + # Detach from gradient computation (TF equivalent of PyTorch's .detach()) + out_detached = tf.stop_gradient(out) + + # Return a tensor that has no gradient connection to input + return out_detached + + model = BackwardErrorModel() + # Should detect that no gradients flow through (detached tensor) + result = check_model_gradients(model, input_tensor) + assert not result, "Should detect no gradients due to backward error" + + +def test_model_with_nan_output(input_tensor): + """Test gradient checking with model that produces NaN outputs. + + This tests whether the function handles numerical instability gracefully. + """ + + class NaNModel(tf.keras.Model): + """Model that produces NaN outputs.""" + + def call(self, x): + """Forward pass producing NaN.""" + # Create NaN by invalid operation + return tf.zeros([tf.shape(x)[0], tf.shape(x)[3]]) / 0.0 + + model = NaNModel() + result = check_model_gradients(model, input_tensor) + # The important thing is it doesn't crash + assert not result, "Should detect no gradients due to NaN outputs" + + +def test_model_with_cancelling_outputs(input_tensor): + """Test that independent VJP probes detect gradients hidden by output summation.""" + + def cancelling_model(x): + return tf.stack([x[..., 0], -x[..., 0]], axis=-1) + + assert check_model_gradients(cancelling_model, input_tensor) + + +def test_model_with_zero_gradient(input_tensor): + """Test that zero gradients are rejected.""" + + assert not check_model_gradients(lambda x: x * 0.0, input_tensor) + + +def test_model_with_finite_gradient(input_tensor): + """Test that finite non-zero gradients are accepted.""" + + assert check_model_gradients(lambda x: x * 2.0, input_tensor) + + +@pytest.mark.parametrize("invalid_gradient", [float("nan"), float("inf")]) +def test_model_with_nonfinite_gradient(input_tensor, invalid_gradient): + """Test that NaN and infinite gradients are rejected.""" + + def invalid_gradient_model(x): + return x * tf.cast(invalid_gradient, x.dtype) + + assert not check_model_gradients(invalid_gradient_model, input_tensor) diff --git a/tests/utils_functions/test_gradients_check_torch.py b/tests/utils_functions/test_gradients_check_torch.py new file mode 100644 index 00000000..6f1d70e1 --- /dev/null +++ b/tests/utils_functions/test_gradients_check_torch.py @@ -0,0 +1,542 @@ +""" +Tests for PyTorch gradient checking utilities. + +This module tests the check_model_gradients function with various output formats +to ensure it correctly detects gradient flow in different model architectures. +""" +# pylint: disable=redefined-outer-name + +import pytest +import torch + +from xplique.utils_functions.common.torch.gradients_check import check_model_gradients +from xplique.utils_functions.object_detection.torch.multi_box_tensor import ( + TorchMultiBoxTensor as MultiBoxTensor, +) + + +class SimpleTensorModel(torch.nn.Module): + """Model that returns a single tensor.""" + + def __init__(self, has_gradient=True): + super().__init__() + self.has_gradient = has_gradient + self.linear = torch.nn.Linear(10, 5) + + def forward(self, x): + """Forward pass of the model.""" + if self.has_gradient: + # Process input to create gradient path + pooled = torch.mean(x, dim=[2, 3]) # (batch, channels) + return self.linear(pooled) + # No gradient path - return constant + return torch.ones(x.shape[0], 5, device=x.device) + + +class DictOutputModel(torch.nn.Module): + """Model that returns a dict of tensors.""" + + def __init__(self): + super().__init__() + self.linear = torch.nn.Linear(10, 5) + + def forward(self, x): + """Forward pass returning dict.""" + pooled = torch.mean(x, dim=[2, 3]) + out = self.linear(pooled) + return {"scores": out, "features": out * 2} + + +class ListOutputModel(torch.nn.Module): + """Model that returns a list of tensors.""" + + def __init__(self): + super().__init__() + self.linear = torch.nn.Linear(10, 5) + + def forward(self, x): + """Forward pass returning list.""" + pooled = torch.mean(x, dim=[2, 3]) + out = self.linear(pooled) + return [out, out * 2, out * 3] + + +class TorchvisionFormatModel(torch.nn.Module): + """Model that returns list of dicts (Torchvision format).""" + + def __init__(self): + super().__init__() + self.linear = torch.nn.Linear(10, 5) + + def forward(self, x): + """Forward pass returning list of dicts.""" + batch_size = x.shape[0] + pooled = torch.mean(x, dim=[2, 3]) + + results = [] + for i in range(batch_size): + out = self.linear(pooled[i : i + 1]) + results.append({"boxes": out, "scores": out * 2, "labels": out * 3}) + return results + + +class NestedDictModel(torch.nn.Module): + """Model that returns nested dict structure. + + IMPORTANT: Only nested tensors have gradients. Top-level 'meta' is constant. + This tests whether the function checks nested structures. + """ + + def __init__(self): + super().__init__() + self.linear = torch.nn.Linear(10, 5) + + def forward(self, x): + """Forward pass returning nested dict.""" + pooled = torch.mean(x, dim=[2, 3]) + out = self.linear(pooled) + return { + "predictions": {"scores": out, "features": out * 2}, + "meta": torch.ones(x.shape[0], 5, device=x.device), # Constant, no gradient + } + + +class DictWithListModel(torch.nn.Module): + """Model that returns dict containing lists of tensors. + + IMPORTANT: Only tensors in the list have gradients. 'meta' is constant. + This tests whether the function checks nested lists. + """ + + def __init__(self): + super().__init__() + self.linear = torch.nn.Linear(10, 5) + + def forward(self, x): + """Forward pass returning dict with list values.""" + pooled = torch.mean(x, dim=[2, 3]) + out = self.linear(pooled) + return { + "predictions": [out, out * 2], + "meta": torch.ones(x.shape[0], 5, device=x.device), # Constant, no gradient + } + + +class MixedListModel(torch.nn.Module): + """Model that returns list with both tensors and dicts. + + IMPORTANT: Only dict tensors have gradients. Direct tensors are constants. + This tests whether the function checks ALL elements in mixed lists. + """ + + def __init__(self): + super().__init__() + self.linear = torch.nn.Linear(10, 5) + + def forward(self, x): + """Forward pass returning mixed list.""" + pooled = torch.mean(x, dim=[2, 3]) + out = self.linear(pooled) + return [ + torch.ones(x.shape[0], 5, device=x.device), # Constant, no gradient + {"scores": out * 2, "features": out * 3}, + torch.zeros(x.shape[0], 5, device=x.device), # Constant, no gradient + ] + + +class NestedListInDictModel(torch.nn.Module): + """Model that returns list of dicts containing lists. + + Only tensors in nested lists have gradients. 'meta' is constant. + This tests whether the function checks deeply nested structures. + """ + + def __init__(self): + super().__init__() + self.linear = torch.nn.Linear(10, 5) + + def forward(self, x): + """Forward pass returning list of dicts with nested lists.""" + batch_size = x.shape[0] + pooled = torch.mean(x, dim=[2, 3]) + + results = [] + for i in range(batch_size): + out = self.linear(pooled[i : i + 1]) + results.append( + { + "predictions": [out, out * 2], + "meta": torch.ones(1, 5, device=x.device), # Constant, no gradient + } + ) + return results + + +class MultiBoxModel(torch.nn.Module): + """Model that returns MultiBoxTensor.""" + + def __init__(self): + super().__init__() + self.linear = torch.nn.Linear(10, 85) # 4 boxes + 1 score + 80 classes + + def forward(self, x): + """Forward pass returning MultiBoxTensor.""" + pooled = torch.mean(x, dim=[2, 3]) # (batch, channels) + out = self.linear(pooled) + + # Create MultiBoxTensor (9 boxes, 85 features each) + # Reshape to (batch * num_boxes, features) + reshaped = out.view(-1, 85) + + # Convert to MultiBoxTensor + return MultiBoxTensor(reshaped) + + +@pytest.fixture +def input_tensor(): + """Create a simple input tensor.""" + return torch.randn(2, 10, 32, 32) + + +def test_simple_tensor_with_gradient(input_tensor): + """Test gradient checking with a simple tensor output.""" + model = SimpleTensorModel(has_gradient=True) + result = check_model_gradients(model, input_tensor) + assert result, "Gradients should flow through simple tensor model" + + +def test_simple_tensor_without_gradient(input_tensor): + """Test gradient checking when no gradient flows.""" + model = SimpleTensorModel(has_gradient=False) + result = check_model_gradients(model, input_tensor) + assert not result, "Should detect when no gradients flow" + + +def test_dict_output(input_tensor): + """Test gradient checking with dict output.""" + model = DictOutputModel() + result = check_model_gradients(model, input_tensor) + assert result, "Gradients should flow through dict output model" + + +def test_list_output(input_tensor): + """Test gradient checking with list output.""" + model = ListOutputModel() + result = check_model_gradients(model, input_tensor) + assert result, "Gradients should flow through list output model" + + +def test_torchvision_format(input_tensor): + """Test gradient checking with Torchvision format (list of dicts).""" + model = TorchvisionFormatModel() + result = check_model_gradients(model, input_tensor) + assert result, "Gradients should flow through Torchvision format model" + + +def test_nested_dict(input_tensor): + """Test gradient checking with nested dict structure. + + Should FAIL with current implementation (only checks top-level 'meta' which has no gradient). + Should PASS with recursive implementation (finds gradients in nested 'predictions' dict). + """ + model = NestedDictModel() + result = check_model_gradients(model, input_tensor) + assert result, "Should find gradients in nested dict structure" + + +def test_dict_with_list(input_tensor): + """Test gradient checking with dict containing lists. + + Should FAIL with current implementation (only checks 'meta' which has no gradient). + Should PASS with recursive implementation (finds gradients in 'predictions' list). + """ + model = DictWithListModel() + result = check_model_gradients(model, input_tensor) + assert result, "Should find gradients in nested list structure" + + +def test_mixed_list(input_tensor): + """Test gradient checking with mixed list (tensors and dicts). + + Should FAIL with current implementation (checks direct tensors which have no gradient). + Should PASS with recursive implementation (finds gradients in nested dict). + """ + model = MixedListModel() + result = check_model_gradients(model, input_tensor) + assert result, "Should find gradients in nested dict within mixed list" + + +def test_nested_list_in_dict(input_tensor): + """Test gradient checking with list of dicts containing lists. + + Should FAIL with current implementation (only checks 'meta' which has no gradient). + Should PASS with recursive implementation (finds gradients in nested lists). + """ + model = NestedListInDictModel() + result = check_model_gradients(model, input_tensor) + assert result, "Should find gradients in deeply nested list structures" + + +def test_multibox_tensor_output(input_tensor): + """Test gradient checking with model that returns MultiBoxTensor.""" + model = MultiBoxModel() + result = check_model_gradients(model, input_tensor) + assert result, "Gradients should flow through MultiBoxTensor (it extends torch.Tensor)" + + +def test_invalid_input_type(): + """Test that invalid input type raises error.""" + model = SimpleTensorModel() + invalid_input = "not a tensor" # type: ignore + with pytest.raises(TypeError): + check_model_gradients(model, invalid_input) # type: ignore + + +def test_callable_function(input_tensor): + """Test gradient checking with a callable function instead of nn.Module.""" + + def simple_func(x): + return torch.mean(x, dim=[2, 3]) + + result = check_model_gradients(simple_func, input_tensor) + assert result, "Gradients should flow through simple callable" + + +def test_empty_dict(input_tensor): + """Test gradient checking with model returning empty dict.""" + + class EmptyDictModel(torch.nn.Module): + """Model returning empty dict.""" + + def forward(self, _): + """Forward pass returning empty dict.""" + return {} + + model = EmptyDictModel() + result = check_model_gradients(model, input_tensor) + assert not result, "Should detect no tensors in empty dict" + + +def test_empty_list(input_tensor): + """Test gradient checking with model returning empty list.""" + + class EmptyListModel(torch.nn.Module): + """Model returning empty list.""" + + def forward(self, _): + """Forward pass returning empty list.""" + return [] + + model = EmptyListModel() + result = check_model_gradients(model, input_tensor) + assert not result, "Should detect no tensors in empty list" + + +def test_dict_with_non_tensor_values(input_tensor): + """Test gradient checking with dict containing non-tensor values.""" + + class MixedDictModel(torch.nn.Module): + """Model returning dict with mixed value types.""" + + def __init__(self): + super().__init__() + self.linear = torch.nn.Linear(10, 5) + + def forward(self, x): + """Forward pass returning dict with mixed values.""" + pooled = torch.mean(x, dim=[2, 3]) + out = self.linear(pooled) + return {"scores": out, "metadata": "some string", "count": 42} + + model = MixedDictModel() + result = check_model_gradients(model, input_tensor) + assert result, "Should still detect gradients from tensor values" + + +def test_list_with_non_tensor_values(input_tensor): + """Test gradient checking with list containing non-tensor values.""" + + class MixedListWithNonTensorsModel(torch.nn.Module): + """Model returning list with mixed value types.""" + + def __init__(self): + super().__init__() + self.linear = torch.nn.Linear(10, 5) + + def forward(self, x): + """Forward pass returning list with mixed values.""" + pooled = torch.mean(x, dim=[2, 3]) + out = self.linear(pooled) + return [out, "metadata", 42, out * 2] + + model = MixedListWithNonTensorsModel() + result = check_model_gradients(model, input_tensor) + assert result, "Should still detect gradients from tensor values in mixed list" + + +def test_model_with_input_preprocessing_breaking_gradients(input_tensor): + """Test gradient checking with model that breaks gradients during preprocessing. + + This replicates DETR's issue: inputs are preprocessed into a new tensor format + where the gradient connection to the original inputs is severed. + + Pattern (like DETR's nested_tensor_from_tensor_list): + 1. Input comes in with requires_grad=True (for attribution) + 2. Preprocessing creates new tensor with torch.zeros() (requires_grad=False) + 3. Data is copied from input to new tensor + 4. Gradient path is broken - backward pass can't reach original inputs + """ + + class PreprocessingModel(torch.nn.Module): + """Model with preprocessing that breaks gradients.""" + + def __init__(self): + super().__init__() + self.conv = torch.nn.Conv2d(10, 5, 3, padding=1) + + def forward(self, x): + """Forward pass with gradient-breaking preprocessing.""" + # Simulate DETR's nested tensor preprocessing + # Create new tensor format (e.g., for padding to common size) + batch_size, channels, height, width = x.shape + + # Create padded tensor with torch.zeros (requires_grad=False!) + padded = torch.zeros( + batch_size, channels, height + 10, width + 10, device=x.device, dtype=x.dtype + ) + + # Copy input data into new tensor (breaks gradient connection) + for i, pad_img in enumerate(padded): + pad_img[:, :height, :width].copy_(x[i]) + + # Continue processing with the model + out = self.conv(padded) + return torch.mean(out, dim=[2, 3]) + + model = PreprocessingModel() + result = check_model_gradients(model, input_tensor) + assert not result, ( + "Should detect broken gradients due to input preprocessing with torch.zeros()" + ) + + +def test_model_with_backward_error(input_tensor): + """Test gradient checking with model that raises error during backward. + + This simulates cases where gradient computation fails due to model architecture + issues, numerical problems, or other errors during backpropagation. + + The check_model_gradients function should return False and catch the error. + """ + + class BackwardErrorModel(torch.nn.Module): + """Model that breaks gradients during backward.""" + + def __init__(self): + super().__init__() + self.linear = torch.nn.Linear(10, 5) + + def forward(self, x): + """Forward pass with detached output.""" + pooled = torch.mean(x, dim=[2, 3]) + out = self.linear(pooled) + + # Create a custom operation that will fail during backward + # by using a non-differentiable operation + out_detached = out.detach() + + # Return a tensor that has no gradient connection to input + return out_detached + + model = BackwardErrorModel() + # Should detect that no gradients flow through (detached tensor) + result = check_model_gradients(model, input_tensor) + assert not result, "Should detect no gradients due to backward error" + + +def test_model_with_nan_output(input_tensor): + """Test gradient checking with model that produces NaN outputs. + + This tests whether the function handles numerical instability gracefully. + """ + + class NaNModel(torch.nn.Module): + """Model that produces NaN outputs.""" + + def forward(self, x): + """Forward pass producing NaN.""" + # Create NaN by invalid operation + return torch.zeros_like(x[:, :, 0, 0]) / 0.0 + + model = NaNModel() + result = check_model_gradients(model, input_tensor) + # Result could be True or False depending on whether NaN propagates + # The important thing is it doesn't crash + assert not result, "Should detect no gradients due to NaN outputs" + + +def test_model_with_cancelling_outputs(input_tensor): + """Test that independent VJP probes detect gradients hidden by output summation.""" + + def cancelling_model(x): + return torch.stack([x[:, 0, 0, 0], -x[:, 0, 0, 0]], dim=-1) + + assert check_model_gradients(cancelling_model, input_tensor) + + +def test_model_with_zero_gradient(input_tensor): + """Test that zero gradients are rejected.""" + + assert not check_model_gradients(lambda x: x * 0.0, input_tensor) + + +def test_model_with_finite_gradient(input_tensor): + """Test that finite non-zero gradients are accepted.""" + + assert check_model_gradients(lambda x: x * 2.0, input_tensor) + + +@pytest.mark.parametrize("invalid_gradient", [float("nan"), float("inf")]) +def test_model_with_nonfinite_gradient(input_tensor, invalid_gradient): + """Test that NaN and infinite gradients are rejected.""" + + def invalid_gradient_model(x): + return x * torch.tensor(invalid_gradient, device=x.device, dtype=x.dtype) + + assert not check_model_gradients(invalid_gradient_model, input_tensor) + + +def test_model_gradient_and_training_states_are_preserved(): + """Test that gradient checking does not change parameter gradients or module modes.""" + model = torch.nn.Sequential( + torch.nn.BatchNorm1d(2), + torch.nn.Dropout(), + torch.nn.Linear(2, 1), + ) + model.train() + model[0].eval() + module_states = [(module, module.training) for module in model.modules()] + parameter_gradients = [] + for parameter in model.parameters(): + parameter.grad = torch.randn_like(parameter) + parameter_gradients.append(parameter.grad) + + assert check_model_gradients(model, torch.randn(4, 2)) + assert [(module, module.training) for module in model.modules()] == module_states + for parameter, gradient in zip(model.parameters(), parameter_gradients): + assert parameter.grad is gradient + + +def test_disabled_grad_mode_is_preserved(input_tensor): + """Test that a checker call does not enable gradients for its caller.""" + with torch.no_grad(): + assert check_model_gradients(lambda x: x, input_tensor) + assert not torch.is_grad_enabled() + + +def test_invalid_input_does_not_enable_grad_mode(): + """Test that invalid inputs do not leak an enabled global grad mode.""" + with torch.no_grad(): + with pytest.raises(TypeError): + check_model_gradients(lambda x: x, torch.ones(2, dtype=torch.int64)) + assert not torch.is_grad_enabled() diff --git a/tests/utils_functions/test_object_detection.py b/tests/utils_functions/test_object_detection.py index 854978d3..bfdaae8b 100644 --- a/tests/utils_functions/test_object_detection.py +++ b/tests/utils_functions/test_object_detection.py @@ -4,7 +4,7 @@ import numpy as np -from xplique.utils_functions.object_detection import _box_iou +from xplique.utils_functions.object_detection.object_detection_operators import _box_iou from ..utils import almost_equal diff --git a/tests/wrappers/test_pytorch_wrapper.py b/tests/wrappers/test_pytorch_wrapper.py index 1d0f388b..9a5a217a 100644 --- a/tests/wrappers/test_pytorch_wrapper.py +++ b/tests/wrappers/test_pytorch_wrapper.py @@ -1,6 +1,7 @@ import numpy as np import pytest import tensorflow as tf +import torch import torch.nn as nn from xplique.attributions import ( @@ -218,3 +219,46 @@ def test_metric_dense(): assert hasattr(metric, "batch_inference_function") score = metric(explanations) assert type(score) in [np.float32, np.float64, float] + + +def test_wrapper_gradient_does_not_mutate_parameter_gradients(): + """Test that wrapper gradients are computed without touching model parameters.""" + model = nn.Linear(2, 1, bias=False) + with torch.no_grad(): + model.weight.copy_(torch.tensor([[2.0, -3.0]])) + model.eval() + original_gradient = torch.full_like(model.weight, 7.0) + model.weight.grad = original_gradient + wrapped_model = TorchWrapper(model, device="cpu", is_channel_first=False) + inputs = tf.constant([[1.0, 2.0]], dtype=tf.float32) + + with tf.GradientTape() as tape: + tape.watch(inputs) + outputs = wrapped_model(inputs) + loss = tf.reduce_sum(outputs) + gradients = tape.gradient(loss, inputs) + + np.testing.assert_allclose(gradients.numpy(), [[2.0, -3.0]]) + assert model.weight.grad is original_gradient + + +def test_wrapper_without_gradients_supports_inference_only(): + """Test that requires_grad=False gives a clear error for gradient requests.""" + model = nn.Linear(2, 1, bias=False) + with torch.no_grad(): + model.weight.copy_(torch.tensor([[2.0, -3.0]])) + model.eval() + wrapped_model = TorchWrapper( + model, + device="cpu", + is_channel_first=False, + requires_grad=False, + ) + inputs = tf.constant([[1.0, 2.0]], dtype=tf.float32) + + np.testing.assert_allclose(wrapped_model(inputs).numpy(), [[-4.0]]) + with tf.GradientTape() as tape: + tape.watch(inputs) + outputs = wrapped_model(inputs) + with pytest.raises(RuntimeError, match="requires_grad=False"): + tape.gradient(outputs, inputs) diff --git a/tox.ini b/tox.ini index 99821640..a63661b5 100644 --- a/tox.ini +++ b/tox.ini @@ -20,21 +20,58 @@ deps = tf219: tensorflow==2.19.0 tf220: tensorflow==2.20.0 package = editable -commands = pytest --cov=xplique --ignore=tests/wrappers/test_pytorch_wrapper.py --ignore=tests/concepts/test_craft_torch.py --ignore=tests/example_based/test_torch.py --disable-warnings --no-cov {posargs} + +commands = pytest --cov=xplique \ + --ignore=tests/concepts/test_factorizer_torch.py \ + --ignore=tests/concepts/test_holistic_craft_classification_torch.py \ + --ignore=tests/concepts/test_holistic_craft_object_detection_torch.py \ + --ignore=tests/concepts/test_craft_torch.py \ + --ignore=tests/utils_functions/test_gradients_check_torch.py \ + --ignore=tests/wrappers/test_pytorch_wrapper.py \ + --ignore=tests/example_based/test_torch.py \ + --disable-warnings --no-cov {posargs} + [testenv:py3{10,11,12,13}-tf{218,219,220}-torch{25}] description = Run tests on PyTorch wrapper with different TensorFlow and PyTorch versions deps = {[testenv]deps} torch25: torch==2.5.0 -commands = pytest --cov=xplique tests/wrappers/test_pytorch_wrapper.py tests/concepts/test_craft_torch.py tests/example_based/test_torch.py --disable-warnings --no-cov {posargs} + torch25: torchvision==0.20.0 +commands = pytest --cov=xplique \ + tests/concepts/test_factorizer_torch.py \ + tests/wrappers/test_pytorch_wrapper.py \ + tests/concepts/test_holistic_craft_classification_torch.py \ + tests/concepts/test_holistic_craft_object_detection_torch.py \ + tests/concepts/test_craft_torch.py \ + tests/concepts/test_latent_extractor_torch.py \ + tests/utils_functions/test_gradients_check_torch.py \ + tests/utils_functions/test_box_manager.py \ + tests/utils_functions/test_box_utilities.py \ + tests/utils_functions/test_box_model_wrapper.py \ + tests/example_based/test_torch.py \ + --disable-warnings --no-cov {posargs} [testenv:py3{10,11,12,13}-tf{220}-torch{210}] description = Run tests on PyTorch wrapper with different TensorFlow and PyTorch versions deps = {[testenv]deps} torch210: torch==2.10.0 -commands = pytest --cov=xplique tests/wrappers/test_pytorch_wrapper.py tests/concepts/test_craft_torch.py tests/example_based/test_torch.py --disable-warnings --no-cov {posargs} + torch210: torchvision==0.25.0 +commands = pytest --cov=xplique \ + tests/concepts/test_factorizer_torch.py \ + tests/wrappers/test_pytorch_wrapper.py \ + tests/concepts/test_holistic_craft_classification_torch.py \ + tests/concepts/test_holistic_craft_object_detection_torch.py \ + tests/concepts/test_craft_torch.py \ + tests/concepts/test_latent_extractor_torch.py \ + tests/utils_functions/test_gradients_check_torch.py \ + tests/utils_functions/test_box_manager.py \ + tests/utils_functions/test_box_utilities.py \ + tests/utils_functions/test_box_model_wrapper.py \ + tests/example_based/test_torch.py \ + --disable-warnings --no-cov {posargs} + [testenv:py313-lint] description = Run linters on the codebase diff --git a/xplique/__init__.py b/xplique/__init__.py index 32064274..c2ab78fe 100644 --- a/xplique/__init__.py +++ b/xplique/__init__.py @@ -6,7 +6,7 @@ techniques """ -__version__ = "1.5.2" +__version__ = "2.0.0" from . import attributions, commons, concepts, example_based, features_visualizations, plots from .commons import Tasks diff --git a/xplique/attributions/base.py b/xplique/attributions/base.py index 117c392e..f1c6f300 100644 --- a/xplique/attributions/base.py +++ b/xplique/attributions/base.py @@ -34,11 +34,12 @@ def sanitize( inputs: Union[tf.data.Dataset, tf.Tensor, np.array], targets: Optional[Union[tf.Tensor, np.array]], *args, + **kwargs, ): # ensure we have tf.tensor inputs, targets = tensor_sanitize(inputs, targets) # then enter the explanation function - return explanation_method(self, inputs, targets, *args) + return explanation_method(self, inputs, targets, *args, **kwargs) return sanitize diff --git a/xplique/attributions/global_sensitivity_analysis/gsa_attribution_method.py b/xplique/attributions/global_sensitivity_analysis/gsa_attribution_method.py index 084b8468..33c7fce5 100644 --- a/xplique/attributions/global_sensitivity_analysis/gsa_attribution_method.py +++ b/xplique/attributions/global_sensitivity_analysis/gsa_attribution_method.py @@ -62,9 +62,11 @@ class GSABaseAttributionMethod(BlackBoxExplainer): Model used for computing explanations. grid_size Cut the image in a grid of (grid_size, grid_size) to estimate an indice per cell. + nb_channels + Number of channels in the masks generation. Default is 1. nb_design Must be a power of two. Number of design, the number of forward - will be: nb_design * (grid_size**2 + 2). Generally not above 32. + will be: nb_design * (grid_size**2 * nb_channels + 2). Generally not above 32. sampler Sampler function to call to generate masks. estimator @@ -86,6 +88,7 @@ def __init__( sampler: Callable, estimator: Callable, grid_size: int = 7, + nb_channels: int = 1, nb_design: int = 32, perturbation_function: Optional[Union[Callable, str]] = "inpainting", batch_size: int = 256, @@ -104,14 +107,15 @@ def __init__( self.sampler = sampler self.estimator = estimator - masks = self.sampler(grid_size**2, nb_design) - self.masks = tf.reshape(masks, (-1, grid_size, grid_size, 1)) + masks = self.sampler(grid_size**2 * nb_channels, nb_design) + self.masks = tf.reshape(masks, (-1, grid_size, grid_size, nb_channels)) @sanitize_input_output def explain( self, inputs: Union[tf.data.Dataset, tf.Tensor, np.ndarray], targets: Optional[Union[tf.Tensor, np.ndarray]] = None, + verbose: bool = False, ) -> tf.Tensor: """ Compute the total Sobol' indices according to the explainer parameter (perturbation @@ -128,6 +132,8 @@ def explain( One-hot encoding for classification or direction {-1, +1} for regression. Tensor or numpy array. Expected shape (N, C) or (N). + verbose + Whether to print progress during the computation. Default is False. Returns ------- @@ -141,7 +147,17 @@ def explain( perturbator = self.perturbation_function(inp) outputs = None - for batch_masks in batch_tensor(self.masks, self.batch_size): + # Calculate total number of batches for progress tracking + total_masks = len(self.masks) + if verbose: + print(f"\nComputing perturbations on {total_masks} masks...") + for batch_idx, batch_masks in enumerate(batch_tensor(self.masks, self.batch_size)): + if verbose: + print( + f"\r Processing mask {batch_idx * self.batch_size + 1}/{total_masks}...", + end="", + flush=True, + ) batch_x, batch_y = self._batch_perturbations( batch_masks, perturbator, target, input_shape ) diff --git a/xplique/attributions/global_sensitivity_analysis/hsic_attribution_method.py b/xplique/attributions/global_sensitivity_analysis/hsic_attribution_method.py index 07e40c85..1575504c 100644 --- a/xplique/attributions/global_sensitivity_analysis/hsic_attribution_method.py +++ b/xplique/attributions/global_sensitivity_analysis/hsic_attribution_method.py @@ -49,6 +49,8 @@ class HsicAttributionMethod(GSABaseAttributionMethod): Function g to explain, g take 3 parameters (f, x, y) and should return a scalar, with f the model, x the inputs and y the targets. If None, use the standard operator g(f, x, y) = f(x)[y]. + nb_channels + Number of channels in the masks generation. Default is 1. """ def __init__( @@ -62,6 +64,8 @@ def __init__( batch_size: int = 256, estimator_batch_size: int = None, operator: Optional[Union[Tasks, str, OperatorSignature]] = None, + *, + nb_channels: int = 1, ): sampler = sampler if sampler is not None else TFSobolSequence(binary=True) estimator = estimator if estimator is not None else BinaryEstimator(output_kernel="rbf") @@ -80,6 +84,7 @@ def __init__( sampler=sampler, estimator=estimator, grid_size=grid_size, + nb_channels=nb_channels, nb_design=nb_design, perturbation_function=perturbation_function, batch_size=batch_size, diff --git a/xplique/attributions/global_sensitivity_analysis/hsic_estimators.py b/xplique/attributions/global_sensitivity_analysis/hsic_estimators.py index 9c483fd4..a94fc351 100644 --- a/xplique/attributions/global_sensitivity_analysis/hsic_estimators.py +++ b/xplique/attributions/global_sensitivity_analysis/hsic_estimators.py @@ -61,9 +61,7 @@ def post_process(score: tf.Tensor, masks: tf.Tensor) -> tf.Tensor: score HSIC scores after post processing. """ - # Reshape to (H, W, 1) and then swap the first two axes. - reshaped = tf.reshape(score, tf.shape(masks)[1:]) - return tf.transpose(reshaped, perm=[1, 0, 2]) + return tf.reshape(score, tf.shape(masks)[1:]) @abstractmethod def input_kernel_func(self, X: tf.Tensor, Y: tf.Tensor) -> tf.Tensor: @@ -140,8 +138,8 @@ def estimator( HSIC estimates Raw HSIC estimates in tensorflow """ - # Rearrange to get (d, nb_design) where d = H*W*1. - X = rearrange(masks, "n h w c -> (c w h) n") + # Rearrange to get (d, nb_design) in HWC order. + X = rearrange(masks, "n h w c -> (h w c) n") # Add singleton dimensions: shape becomes (d, 1, nb_design, 1) X1 = rearrange(X, "d n -> d 1 n 1") # Swap last two axes: shape becomes (d, 1, 1, nb_design) diff --git a/xplique/attributions/global_sensitivity_analysis/sobol_attribution_method.py b/xplique/attributions/global_sensitivity_analysis/sobol_attribution_method.py index e1340c2d..0625b5eb 100644 --- a/xplique/attributions/global_sensitivity_analysis/sobol_attribution_method.py +++ b/xplique/attributions/global_sensitivity_analysis/sobol_attribution_method.py @@ -26,7 +26,7 @@ class SobolAttributionMethod(GSABaseAttributionMethod): Cut the image in a grid of (grid_size, grid_size) to estimate an indice per cell. nb_design Must be a power of two. Number of design, the number of forward - will be: nb_design * (grid_size**2 + 2). Generally not above 32. + will be: nb_design * (grid_size**2 * nb_channels + 2). Generally not above 32. sampler Sampler used to generate the (quasi-)monte carlo samples, QMC (sobol sequence recommended). For more option, see the sampler module. @@ -42,6 +42,8 @@ class SobolAttributionMethod(GSABaseAttributionMethod): Function g to explain, g take 3 parameters (f, x, y) and should return a scalar, with f the model, x the inputs and y the targets. If None, use the standard operator g(f, x, y) = f(x)[y]. + nb_channels + Number of channels in the masks generation. Default is 1. """ def __init__( @@ -54,6 +56,8 @@ def __init__( perturbation_function: Optional[Union[Callable, str]] = "inpainting", batch_size=256, operator: Optional[Union[Tasks, str, OperatorSignature]] = None, + *, + nb_channels: int = 1, ): assert (nb_design & (nb_design - 1) == 0) and nb_design != 0, ( "The number of design must be a power of two." @@ -75,6 +79,7 @@ def __init__( sampler=sampler, estimator=estimator, grid_size=grid_size, + nb_channels=nb_channels, nb_design=nb_design, perturbation_function=perturbation_function, batch_size=batch_size, diff --git a/xplique/commons/model_override.py b/xplique/commons/model_override.py index ac768e0f..bea47253 100644 --- a/xplique/commons/model_override.py +++ b/xplique/commons/model_override.py @@ -189,6 +189,24 @@ def has_relu_activation(layer: tf.keras.layers.Layer) -> bool: return layer.activation in [tf.nn.relu, tf.keras.activations.relu] +def _clone_layer(layer: tf.keras.layers.Layer) -> tf.keras.layers.Layer: + """Clone a layer without deserializing the callable of an in-memory Lambda layer.""" + if not isinstance(layer, tf.keras.layers.Lambda): + return layer.__class__.from_config(layer.get_config()) + + config = layer.get_config() + for lambda_argument in ("function", "output_shape", "mask", "arguments"): + config.pop(lambda_argument, None) + + return layer.__class__( + function=layer.function, + output_shape=layer._output_shape, # pylint: disable=W0212 + mask=layer.mask, + arguments=dict(layer.arguments), + **config, + ) + + def override_relu_gradient(model: tf.keras.Model, relu_policy: Callable) -> tf.keras.Model: """ Given a model, commute all original ReLU by a new given ReLU policy. @@ -204,7 +222,7 @@ def override_relu_gradient(model: tf.keras.Model, relu_policy: Callable) -> tf.k ------- model_commuted """ - cloned_model = clone_model(model) + cloned_model = clone_model(model, clone_function=_clone_layer) cloned_model.set_weights(model.get_weights()) for layer_id in range(len(cloned_model.layers)): # pylint: disable=C0200 diff --git a/xplique/commons/operators.py b/xplique/commons/operators.py index 9a1a1690..79166029 100644 --- a/xplique/commons/operators.py +++ b/xplique/commons/operators.py @@ -164,64 +164,73 @@ def batch_loop(args): # function to loop on for `tf.map_fn` obj, obj_ref = args - if obj is None or obj.shape[0] == 0: - return tf.constant(0.0, dtype=inputs.dtype) - - # compute predicted boxes for a given image - # (nb_box_pred, 4), (nb_box_pred, 1), (nb_box_pred, nb_classes) - current_boxes, proba_detection, classification = _format_objects(obj) - size = tf.shape(current_boxes)[0] - - if tf.shape(obj_ref).shape[0] == 1: - obj_ref = tf.expand_dims(obj_ref, axis=0) - - # DRise consider the reference objectness to be 1 - # (nb_box_ref, 4), _, (nb_box_ref, nb_classes) - boxes_refs, _, class_refs = _format_objects(obj_ref) - - # (nb_box_ref, nb_box_pred, 4) - boxes_refs = tf.repeat(tf.expand_dims(boxes_refs, axis=1), repeats=size, axis=1) - - # (nb_box_ref, nb_box_pred) - intersection_score = intersection_score_fn(boxes_refs, current_boxes) - - # (nb_box_pred,) - detection_probability = tf.squeeze(proba_detection, axis=1) - - # set detection probability to 1 if it should be included - detection_probability = tf.cond( - tf.cast(include_detection_probability, tf.bool), - true_fn=lambda: detection_probability, - false_fn=lambda: tf.ones_like(detection_probability), + # Reshape instead of branching so TensorFlow retains a statically known rank. + obj_ref = tf.reshape(obj_ref, (-1, tf.shape(obj_ref)[-1])) + obj_ref = tf.ensure_shape(obj_ref, [None, None]) + + # Tensor-mode detection outputs are zero-padded to equalize batch shapes. + obj = tf.boolean_mask(obj, tf.reduce_any(tf.not_equal(obj, 0), axis=-1)) + obj_ref = tf.boolean_mask(obj_ref, tf.reduce_any(tf.not_equal(obj_ref, 0), axis=-1)) + + def compute_score(): + # compute predicted boxes for a given image + # (nb_box_pred, 4), (nb_box_pred, 1), (nb_box_pred, nb_classes) + current_boxes, proba_detection, classification = _format_objects(obj) + size = tf.shape(current_boxes)[0] + + # DRise consider the reference objectness to be 1 + # (nb_box_ref, 4), _, (nb_box_ref, nb_classes) + boxes_refs, _, class_refs = _format_objects(obj_ref) + + # (nb_box_ref, nb_box_pred, 4) + boxes_refs = tf.repeat(tf.expand_dims(boxes_refs, axis=1), repeats=size, axis=1) + + # (nb_box_ref, nb_box_pred) + intersection_score = intersection_score_fn(boxes_refs, current_boxes) + + # (nb_box_pred,) + detection_probability = tf.squeeze(proba_detection, axis=1) + + # set detection probability to 1 if it should be included + detection_probability = tf.cond( + tf.cast(include_detection_probability, tf.bool), + true_fn=lambda: detection_probability, + false_fn=lambda: tf.ones_like(detection_probability), + ) + + # (nb_box_ref, nb_box_pred, nb_classes) + class_refs = tf.repeat(tf.expand_dims(class_refs, axis=1), repeats=size, axis=1) + + # (nb_box_ref, nb_box_pred) + classification_score = tf.reduce_sum(class_refs * classification, axis=-1) / ( + tf.norm(classification, axis=-1) * tf.norm(class_refs, axis=-1) + _EPSILON + ) + + # set classification score to 1 if it should be included + classification_score = tf.cond( + tf.cast(include_classification_score, tf.bool), + true_fn=lambda: classification_score, + false_fn=lambda: tf.ones_like(classification_score), + ) + + # Compute score as defined in DRise for all possible pair of boxes + # (nb_box_ref, nb_box_pred) + boxes_pairwise_scores = ( + intersection_score * detection_probability * classification_score + ) + + # select for a reference box the most similar predicted box score + # (nb_box_ref,) + ref_boxes_scores = tf.reduce_max(boxes_pairwise_scores, axis=1) + + # get an attribution for several boxes in the same time + return tf.reduce_mean(ref_boxes_scores) + + return tf.cond( + tf.logical_and(tf.shape(obj)[0] > 0, tf.shape(obj_ref)[0] > 0), + compute_score, + lambda: tf.constant(0.0, dtype=inputs.dtype), ) - # (nb_box_ref, nb_box_pred, nb_classes) - class_refs = tf.repeat(tf.expand_dims(class_refs, axis=1), repeats=size, axis=1) - - # (nb_box_ref, nb_box_pred) - classification_score = tf.reduce_sum(class_refs * classification, axis=-1) / ( - tf.norm(classification, axis=-1) * tf.norm(class_refs, axis=-1) + _EPSILON - ) - - # set classification score to 1 if it should be included - classification_score = tf.cond( - tf.cast(include_classification_score, tf.bool), - true_fn=lambda: classification_score, - false_fn=lambda: tf.ones_like(classification_score), - ) - - # Compute score as defined in DRise for all possible pair of boxes - # (nb_box_ref, nb_box_pred) - boxes_pairwise_scores = intersection_score * detection_probability * classification_score - - # select for a reference box the most similar predicted box score - # (nb_box_ref,) - ref_boxes_scores = tf.reduce_max(boxes_pairwise_scores, axis=1) - - # get an attribution for several boxes in the same time - # () - image_score = tf.reduce_mean(ref_boxes_scores) - return image_score - objects = model(inputs) return tf.map_fn(batch_loop, (objects, targets), fn_output_signature=tf.float32) diff --git a/xplique/commons/prediction_types.py b/xplique/commons/prediction_types.py new file mode 100644 index 00000000..17d013bd --- /dev/null +++ b/xplique/commons/prediction_types.py @@ -0,0 +1,87 @@ +""" +Protocols and types for model prediction outputs. + +This module defines common interfaces for different types of model predictions +(object detection, classification, etc.) to enable polymorphic handling without +runtime type checking. +""" +# pylint: disable=unnecessary-ellipsis + +from typing import Protocol, runtime_checkable + + +@runtime_checkable +class StructuredPrediction(Protocol): + """ + Protocol for unified handling of structured model predictions. + + This protocol defines a common interface that both object detection outputs + (MultiBoxTensor) and classifier outputs (ClassifierTensor) implement, allowing + code to work with either type without isinstance checks. + + The protocol includes: + - to_batched_tensor(): Ensures output has batch dimension for attribution methods + - filter(): Filters predictions based on confidence threshold (OD only, no-op for classifiers) + - to_attribution_target(): Builds the attribution target (OD: returns filtered boxes; + classifiers: builds a one-hot vector for the requested class_id) + """ + + def to_batched_tensor(self): + """ + Convert prediction to batched tensor format. + + Ensures the output has a batch dimension, which is required by + attribution methods. For single predictions, adds a batch dimension. + For already-batched predictions, returns as-is. + + Returns + ------- + tensor + Tensor with batch dimension: (batch_size, ...) + """ + ... + + def filter(self, class_id=None, confidence=None): + """ + Filter predictions by class ID and/or confidence threshold. + + For object detection, this filters bounding boxes by class and score. + For classifiers, this is a no-op returning self (there are no boxes to filter). + + Parameters + ---------- + class_id + Optional class ID to filter by (for object detection only) + confidence + Optional minimum confidence threshold (for object detection only) + + Returns + ------- + StructuredPrediction + Filtered predictions (or self for classifiers) + """ + ... + + def to_attribution_target(self, class_id=None): + """ + Build the attribution target for the requested class. + + For object detection, the filtered boxes (from filter()) are already the + correct target — this returns self unchanged. + For classifiers, a one-hot vector is constructed for ``class_id`` since + the actual logit/probability values are not used by the attribution method + (only the class index and output shape matter). + + Parameters + ---------- + class_id + Class to target. For classifiers, builds one_hot(class_id, num_classes). + For object detection, ignored (boxes carry the class information). + If None, returns self unchanged for both types. + + Returns + ------- + StructuredPrediction + Attribution target ready for to_batched_tensor(). + """ + ... diff --git a/xplique/concepts/__init__.py b/xplique/concepts/__init__.py index e74316d7..a699a061 100644 --- a/xplique/concepts/__init__.py +++ b/xplique/concepts/__init__.py @@ -5,14 +5,19 @@ from .cav import Cav from .craft import DisplayImportancesOrder from .craft_tf import CraftManagerTf, CraftTf +from .holistic_craft import HolisticCraft, PartialExplainer +from .latent_extractor import EncodedData from .tcav import Tcav +from .tf.holistic_craft import HolisticCraftTf try: from .craft_torch import CraftManagerTorch, CraftTorch + from .torch.holistic_craft import HolisticCraftTorch __all__ = [ "CraftManagerTorch", "CraftTorch", + "HolisticCraftTorch", ] except ImportError: __all__ = [] @@ -23,5 +28,9 @@ "DisplayImportancesOrder", "CraftManagerTf", "CraftTf", + "EncodedData", + "HolisticCraft", + "HolisticCraftTf", + "PartialExplainer", "Tcav", ] diff --git a/xplique/concepts/craft.py b/xplique/concepts/craft.py index fdb77c95..42f10af6 100644 --- a/xplique/concepts/craft.py +++ b/xplique/concepts/craft.py @@ -27,12 +27,13 @@ class Factorization: """Dataclass handling data produced during the Factorization step.""" - inputs: np.ndarray + inputs: Optional[np.ndarray] class_id: int - crops: np.ndarray + crops: Optional[np.ndarray] reducer: NMF - crops_u: np.ndarray + crops_u: Optional[np.ndarray] concept_bank_w: np.ndarray + coeffs_u: Optional[np.ndarray] = None class Sensitivity: @@ -267,7 +268,12 @@ def fit( concept_bank_w = reducer.components_.astype(np.float32) self.factorization = Factorization( - inputs, class_id, crops, reducer, crops_u, concept_bank_w + class_id=class_id, + reducer=reducer, + concept_bank_w=concept_bank_w, + inputs=inputs, + crops=crops, + crops_u=crops_u, ) return crops, crops_u, concept_bank_w diff --git a/xplique/concepts/factorizer.py b/xplique/concepts/factorizer.py new file mode 100644 index 00000000..f7ccfa40 --- /dev/null +++ b/xplique/concepts/factorizer.py @@ -0,0 +1,278 @@ +""" +Factorizer protocol and base implementations for concept extraction + +Note: All factorizers expect activations in shape (spatial_flattened, channels), +where spatial_flattened = N*H*W for batched 2D feature maps. +""" + +from abc import ABC, abstractmethod +from typing import Any, Tuple, Union + +import numpy as np +from sklearn.decomposition import NMF + + +class ConceptFactorizer(ABC): + """ + Abstract base class for concept factorization methods used in CRAFT. + """ + + @abstractmethod + def fit(self, activations: np.ndarray) -> Tuple[np.ndarray, np.ndarray]: + """ + Fit the factorizer on activations. + + Parameters + ---------- + activations : np.ndarray + Activations to factorize, shape (n_samples, n_features) + + Returns + ------- + concept_bank_w : np.ndarray + Concept bank (dictionary), shape (n_concepts, n_features) + coeffs_u : np.ndarray + Coefficients for the input activations, shape (n_samples, n_concepts) + """ + ... + + @abstractmethod + def encode(self, activations: np.ndarray) -> np.ndarray: + """ + Encode activations to coefficients (non-differentiable). + + Parameters + ---------- + activations : np.ndarray + Activations to encode, shape (n_samples, n_features) + + Returns + ------- + np.ndarray + Coefficients, shape (n_samples, n_concepts) + """ + ... + + @abstractmethod + def encode_differentiable(self, activations) -> Union[np.ndarray, Any]: + """ + Encode activations to coefficients (differentiable, framework-specific). + + Parameters + ---------- + activations : Tensor + Activations to encode (torch.Tensor or tf.Tensor) + + Returns + ------- + Tensor + Coefficients (same framework as input) + """ + ... + + @abstractmethod + def decode(self, coefficients) -> Union[np.ndarray, Any]: + """ + Decode coefficients back to activations (naturally differentiable). + + Parameters + ---------- + coefficients : array or Tensor + Coefficients to decode, shape (n_samples, n_concepts) + + Returns + ------- + array or Tensor + Reconstructed activations, shape (n_samples, n_features) + """ + ... + + @abstractmethod + def get_concept_bank(self) -> np.ndarray: + """ + Get the concept bank (W matrix). + + Returns + ------- + np.ndarray + Concept bank, shape (n_concepts, n_features) + """ + ... + + @property + @abstractmethod + def is_fitted(self) -> bool: + """ + Check if the factorizer has been fitted. + + Returns + ------- + bool + True if fitted, False otherwise + """ + ... + + @property + @abstractmethod + def requires_positive_activations(self) -> bool: + """ + Whether this factorizer requires positive activations. + + Returns + ------- + bool + True if activations must be non-negative + """ + ... + + +class SklearnNMFFactorizer(ConceptFactorizer): + """ + Base sklearn NMF factorizer (framework-agnostic). + """ + + def __init__(self, n_components: int = 20, **nmf_kwargs): + """ + Parameters + ---------- + n_components : int + Number of concepts to extract + **nmf_kwargs + Additional arguments passed to sklearn.decomposition.NMF + """ + self.n_components = n_components + self.nmf_kwargs = nmf_kwargs + self._decomposer = None + self._concept_bank_w = None + + def fit(self, activations: np.ndarray) -> Tuple[np.ndarray, np.ndarray]: + """ + Fit NMF on activations. + + Parameters + ---------- + activations : np.ndarray + Activations to factorize, shape (n_samples, n_features) + + Returns + ------- + concept_bank_w : np.ndarray + Concept bank (dictionary), shape (n_concepts, n_features) + coeffs_u : np.ndarray + Coefficients for the input activations, shape (n_samples, n_concepts) + """ + self._decomposer = NMF(n_components=self.n_components, **self.nmf_kwargs) + coeffs_u = self._decomposer.fit_transform(activations) + self._concept_bank_w = self._decomposer.components_ + return self._concept_bank_w, coeffs_u + + def encode(self, activations: np.ndarray) -> np.ndarray: + """ + Encode activations to coefficients using sklearn NMF transform. + + Parameters + ---------- + activations : np.ndarray + Activations to encode, shape (n_samples, n_features) + + Returns + ------- + np.ndarray + Coefficients, shape (n_samples, n_concepts) + """ + if self._decomposer is None: + raise ValueError("Factorizer must be fitted before encoding") + + # Cast activations to concept bank dtype for consistency + if activations.dtype != self._concept_bank_w.dtype: + activations = activations.astype(self._concept_bank_w.dtype) + + return self._decomposer.transform(activations) + + @abstractmethod + def encode_differentiable(self, activations): + """ + Differentiable encoding (not implemented in base class). + + This method should be overridden by framework-specific subclasses. + + Parameters + ---------- + activations : Tensor + Activations to encode + + Returns + ------- + Tensor + Coefficients + + Raises + ------ + NotImplementedError + Always raised, must be implemented by subclasses + """ + raise NotImplementedError( + "encode_differentiable must be implemented by framework-specific subclasses" + ) + + def decode(self, coefficients) -> Union[np.ndarray, Any]: + """ + Decode coefficients to activations via matrix multiplication. + + This operation is naturally differentiable when using tensors. + + Parameters + ---------- + coefficients : array or Tensor + Coefficients to decode, shape (n_samples, n_concepts) + + Returns + ------- + array or Tensor + Reconstructed activations, shape (n_samples, n_features) + """ + if self._concept_bank_w is None: + raise ValueError("Factorizer must be fitted before decoding") + + if isinstance(coefficients, np.ndarray): + return coefficients @ self._concept_bank_w + + concept_bank_tensor = type(coefficients)(self._concept_bank_w) + return coefficients @ concept_bank_tensor + + def get_concept_bank(self) -> np.ndarray: + """ + Get the concept bank (W matrix from NMF). + + Returns + ------- + np.ndarray + Concept bank, shape (n_concepts, n_features) + """ + if self._concept_bank_w is None: + raise ValueError("Factorizer must be fitted before getting concept bank") + return self._concept_bank_w + + @property + def is_fitted(self) -> bool: + """ + Check if the factorizer has been fitted. + + Returns + ------- + bool + True if fitted, False otherwise + """ + return self._decomposer is not None and self._concept_bank_w is not None + + @property + def requires_positive_activations(self) -> bool: + """ + NMF requires non-negative activations. + + Returns + ------- + bool + True + """ + return True diff --git a/xplique/concepts/holistic_craft.py b/xplique/concepts/holistic_craft.py new file mode 100644 index 00000000..06090c0c --- /dev/null +++ b/xplique/concepts/holistic_craft.py @@ -0,0 +1,1201 @@ +""" +Framework-agnostic CRAFT implementation for holistic model explanations. +""" + +import warnings +from abc import ABC, abstractmethod +from typing import Any, Callable, List, Optional, Tuple, Union + +import cv2 +import numpy as np +from matplotlib import pyplot as plt +from matplotlib.figure import Figure +from sklearn.exceptions import NotFittedError + +from xplique.attributions.global_sensitivity_analysis.sobol_attribution_method import ( + SobolAttributionMethod, +) +from xplique.attributions.gradient_input import GradientInput +from xplique.commons.prediction_types import StructuredPrediction +from xplique.plots.image import _clip_percentile + +from .craft import Factorization, Sensitivity +from .factorizer import SklearnNMFFactorizer +from .latent_extractor import EncodedData, LatentData, LatentExtractor + + +def show_ax(img, ax, **kwargs): + """ + Display an image on a matplotlib axis with normalization. + + Converts channel-first images to channel-last format, normalizes pixel + values to [0, 1] range, and displays without axis labels. + + Parameters + ---------- + img + Image array to display, either in channel-first (C, H, W) or + channel-last (H, W, C) format + ax + Matplotlib axis object on which to display the image + kwargs + Additional keyword arguments passed to ax.imshow() + """ + img = np.array(img, dtype=np.float32) + if img.shape[0] == 3: + img = img.transpose(1, 2, 0) + + img -= img.min() + if img.max() > 0: + img /= img.max() + ax.imshow(img, **kwargs) + ax.axis("off") + + +class PartialExplainer: + """ + Wrapper for explainer classes to enable deferred instantiation. + + This class stores an explainer class and its configuration kwargs, allowing + the explainer to be instantiated later when the model and batch_size become + available during concept importance estimation. + + Parameters + ---------- + explainer_class + The explainer class to instantiate (e.g., GradientInput, SobolAttributionMethod). + Must be callable and accept 'model' and 'batch_size' as keyword arguments. + kwargs + Configuration arguments for the explainer (e.g., operator, reducer, grid_size). + Should NOT include 'model' or 'batch_size' as these will be provided during + instantiation. + + Raises + ------ + ValueError + If 'model' or 'batch_size' are provided in kwargs, since these are reserved for + later instantiation. + """ + + def __init__(self, explainer_class, **kwargs): + # Validate that model and batch_size are not provided + if "model" in kwargs or "batch_size" in kwargs: + raise ValueError( + "PartialExplainer should not receive 'model' or 'batch_size' arguments. " + "These will be provided automatically during importance estimation." + ) + + self.explainer_class = explainer_class + self.kwargs = kwargs + + def __call__(self, model, batch_size): + """ + Instantiate the explainer with the provided model and batch_size. + + Parameters + ---------- + model + The model to explain + batch_size + Batch size for processing + + Returns + ------- + explainer + Instantiated explainer object + """ + return self.explainer_class(model=model, batch_size=batch_size, **self.kwargs) + + +class HolisticCraft(ABC): + """ + Framework-agnostic CRAFT implementation for holistic model explanations. + + This base class provides concept-based explanations for various model types + (object detection, classification, etc.) by extracting and analyzing intermediate + activations using Non-negative Matrix Factorization (NMF). It supports both + TensorFlow and PyTorch through framework-specific subclasses. + + Ref. Fel et al., CRAFT Concept Recursive Activation FacTorization (2023). + https://arxiv.org/abs/2211.10154 + Ref. Fel et al., A Holistic Approach to Unifying Automatic Concept Extraction + and Concept Importance Estimation (2023). + https://arxiv.org/pdf/2306.07304 + + The workflow involves: + 1. Extracting latent activations from a computer vision model + 2. Factorizing activations into interpretable concepts using NMF + 3. Computing concept importance using gradient-based attribution methods + 4. Visualizing concepts as spatial heatmaps overlaid on input images + + Parameters + ---------- + latent_extractor + Extractor that splits the model into encoder (input to activations) and + decoder (activations to predictions) for concept extraction + number_of_concepts + Number of concepts to extract via NMF decomposition + device + Device specification for tensor operations (framework-specific) + factorizer + Optional factorizer instance implementing the ConceptFactorizer protocol. + If None, creates a default SklearnNMFFactorizer with alpha_W=1e-2 and + max_iter=200 + + Attributes + ---------- + latent_extractor + The latent extractor instance + number_of_concepts + Number of concepts extracted + batch_size + Batch size inherited from latent_extractor + factorization + Factorization object containing NMF results, populated after fit() + factorizer + Factorizer instance used for concept extraction + device + Device for tensor operations + cmaps + List of colormaps for visualization + """ + + def __init__( + self, + latent_extractor: LatentExtractor, + number_of_concepts: int = 20, + device: str = None, + factorizer: Optional[Any] = None, + ): + self.latent_extractor = latent_extractor + self.number_of_concepts = number_of_concepts + self.batch_size = latent_extractor.batch_size + self.factorization = None + self.device = device + + # Use provided factorizer or create default NMF factorizer + if factorizer is None: + self.factorizer = SklearnNMFFactorizer( + n_components=number_of_concepts, alpha_W=1e-2, max_iter=200 + ) + else: + self.factorizer = factorizer + + # Setup visualization colormaps + self.cmaps = [Sensitivity._get_alpha_cmap(cmap) for cmap in plt.get_cmap("tab10").colors] + + def check_if_fitted(self): + """Checks if the factorization model has been fitted to input data. + + Raises + ------ + NotFittedError + If the factorization model has not been fitted to input data. + """ + if not self.factorizer.is_fitted or self.factorization is None: + raise NotFittedError("The factorization model has not been fitted to input data yet.") + + def fit(self, inputs, class_id: int = 0): + """ + Fit NMF to extract concepts from latent activations. + + Processes input images through the latent extractor to obtain intermediate + activations, then applies Non-negative Matrix Factorization to discover + interpretable concepts. The concepts are spatial patterns in activation + space that recur across different images and locations. + + Activations are converted to NumPy immediately to minimize device memory usage, + which is especially important for large datasets and GPU processing. + + Parameters + ---------- + inputs + Input images to extract concepts from, as framework tensors or arrays + class_id + Target class ID for object detection (used in factorization metadata) + + """ + # Pass the data through the 1st part of the model, convert each batch to + # numpy immediately to free device memory before the next batch + activations_list = [ + latent_data.get_activations(as_numpy=True) + for latent_data in self.latent_extractor.input_to_latent_generator(inputs) + ] + if not activations_list: + raise ValueError("No activations extracted from inputs.") + activations = np.concatenate(activations_list, axis=0) + + needs_reshape = len(activations.shape) > 2 # (N,H,W,C) or (N,Tokens,C) + if needs_reshape: + activations_original_shape = activations.shape[:-1] + spatial_positions = int(np.prod(activations_original_shape[1:])) + if spatial_positions <= 1: + warnings.warn( + f"Latent activations have only {spatial_positions} spatial position(s) " + f"(shape {activations.shape}). CRAFT needs multiple spatial positions to " + f"extract meaningful concepts. Check your split_layer, it may be set " + f"too close to the end of the network (e.g. after a global pooling layer).", + UserWarning, + stacklevel=3, + ) + # Activations are already in numpy format, reshape for factorization + activations = np.reshape(activations, (-1, activations.shape[-1])) + + # Check if factorizer requires positive activations + if self.factorizer.requires_positive_activations and np.any(activations < 0): + raise ValueError( + "Factorizer requires non-negative activations but received negative values." + ) + + # Apply factorizer to the activations + concept_bank_w, coeffs_u = self.factorizer.fit(activations) + concept_bank_w = concept_bank_w.astype(np.float32) + + # Reshape coefficients back to spatial dimensions + if needs_reshape: + coeffs_u = coeffs_u.reshape(*activations_original_shape, -1) + + self.factorization = Factorization( + inputs=None, + class_id=class_id, + crops=None, + reducer=self.factorizer, + concept_bank_w=concept_bank_w, + crops_u=None, + coeffs_u=coeffs_u, + ) + + def transform(self, inputs=None, resize=None) -> np.ndarray: + """Transform inputs to concept coefficients. + + This method encodes the inputs and returns only the concept coefficients + as a concatenated numpy array, discarding the latent data. + + If inputs is None, returns the stored coefficients from fit() if available + (useful for ConvexNMF which can only encode training data). + + Parameters + ---------- + inputs + Input images to transform. If None, returns stored coefficients from fit(). + resize + Target size for resizing images + + Returns + ------- + coeffs_u + Concept coefficients for the inputs (or stored coefficients if inputs=None) + + Raises + ------ + ValueError + If inputs is None but no stored coefficients are available + """ + # If no inputs provided, return stored coefficients from fit() + if inputs is None: + self.check_if_fitted() + if self.factorization.coeffs_u is None: + raise ValueError("No stored coefficients available, and no inputs given.") + return self.factorization.coeffs_u + + # encode, but only return coeffs_u as a single tensor + encoded_data = self.encode(inputs, resize) + if not encoded_data: + raise ValueError("No activations extracted from inputs.") + # extract coeffs_u using named attribute access for clarity + coeffs_u = np.concatenate([enc.coeffs_u for enc in encoded_data], axis=0) + return coeffs_u + + def latent_to_concept(self, latent_data: LatentData) -> np.ndarray: + """ + Transform latent data to concept coefficients. + + Projects latent activations onto the learned concept space using the + fitted NMF model. This non-differentiable transform is faster than + latent_to_concept_differentiable() but cannot be used for gradient-based + methods. + + Parameters + ---------- + latent_data + Single image's latent representation containing activations + + Returns + ------- + coeffs_u + Concept coefficients, shape (H, W, n_concepts) + + Raises + ------ + ValueError + If latent_data is not a single LatentData instance + NotFittedError + If fit() has not been called yet + """ + if not isinstance(latent_data, LatentData): + raise ValueError( + f"latent_to_concept() only accepts a single LatentData as input, " + f"got {type(latent_data)}" + ) + self.check_if_fitted() + + activations = latent_data.get_activations(as_numpy=True) + needs_reshape = len(activations.shape) > 2 # (N,H,W,C) or (N,Tokens,C) + if needs_reshape: + activations_original_shape = activations.shape[:-1] + activations = np.reshape(activations, (-1, activations.shape[-1])) + + # Encode activations to coefficients using the factorizer + coeffs_u = self.factorizer.encode(activations) + if needs_reshape: + coeffs_u = np.reshape(coeffs_u, (*activations_original_shape, -1)) + return coeffs_u + + @abstractmethod + def latent_to_concept_differentiable(self, latent_data: LatentData) -> Any: + """ + Transform latent data to concept coefficients with gradient preservation. + + Uses a differentiable non-negative optimization procedure to project + activations onto concepts while maintaining the computational graph for + gradient-based attribution methods. Must be implemented by + framework-specific subclasses. + + Parameters + ---------- + latent_data + Single image's latent representation containing activations + + Returns + ------- + coeffs_u + Concept coefficients as framework tensor with gradients + """ + raise NotImplementedError + + @abstractmethod + def _to_numpy(self, tensor: Any) -> np.ndarray: + """Convert a framework-specific tensor to a numpy array. + + Parameters + ---------- + tensor + Framework tensor (torch.Tensor or tf.Tensor) or numpy array + + Returns + ------- + array + Numpy array + """ + raise NotImplementedError + + @abstractmethod + def _to_tensor(self, array: np.ndarray, dtype: Any = None) -> Any: + """Convert a numpy array to a framework-specific tensor. + + Parameters + ---------- + array + Numpy array to convert + dtype + Target framework dtype (e.g., torch.float32 or tf.float32) + + Returns + ------- + tensor + Framework-specific tensor + """ + raise NotImplementedError + + def encode( + self, + inputs: Union[np.ndarray, Any], + resize: Optional[Tuple[int, int]] = None, + differentiable: bool = False, + ) -> List[EncodedData]: + """Encode inputs to latent data and concept coefficients. + + Parameters + ---------- + inputs + Input images to encode + resize + Target size for resizing images + differentiable + If True, preserves gradients for backpropagation using differentiable + non-negative optimization. If False (default), uses standard NMF + transform which is faster but does not preserve gradients. + + Returns + ------- + encoded_data + List of EncodedData named tuples, each containing: + - latent_data: LatentData object with intermediate activations + - coeffs_u: Concept coefficients (numpy array or tensor with gradients) + + When differentiable=False, coeffs_u are numpy arrays. + When differentiable=True, coeffs_u are framework tensors + (torch.Tensor or tf.Tensor) with gradients preserved. + """ + encoded_data = [] + for latent_data in self.latent_extractor.input_to_latent_generator( + inputs, resize, keep_gradients=differentiable + ): + if differentiable: + coeffs_u = self.latent_to_concept_differentiable(latent_data) + else: + coeffs_u = self.latent_to_concept(latent_data) + encoded_data.append(EncodedData(latent_data, coeffs_u)) + return encoded_data + + def decode( + self, latent_data: LatentData, coeffs_u: Union[np.ndarray, Any] + ) -> StructuredPrediction: + """Decode concept coefficients back to predictions. + + This method accepts a single LatentData and returns a prediction tensor + that implements the StructuredPrediction protocol (either MultiBoxTensor for + object detection or ClassifierTensor for classification). + + The latent_extractor.latent_to_logit() method returns predictions that + are already formatted by the output_formatter: + - A list with 1 element (PyTorch formatters) + - A single tensor directly (TensorFlow formatters with batch_size=1) + + The formatter guarantees the output implements StructuredPrediction protocol. + This method only handles unwrapping single-element lists. + + Parameters + ---------- + latent_data + Single image's latent representation (not batched) + coeffs_u + Concept coefficients for reconstruction + + Returns + ------- + predictions + Predictions implementing StructuredPrediction protocol (has filter() + and to_batched_tensor() methods). Concrete types are MultiBoxTensor + for object detection or ClassifierTensor for classification. + + Raises + ------ + ValueError + If latent_data is not a single LatentData instance, or if + latent_to_logit returns a list with != 1 elements + """ + if not isinstance(latent_data, LatentData): + raise ValueError("decode() only accepts a single LatentData as input") + + result = self._decode_coefficients(latent_data, coeffs_u) + + # Public decoding returns one structured prediction. Internal attribution + # decoding uses _decode_coefficients() directly to handle perturbation batches. + if isinstance(result, list): + if len(result) != 1: + raise ValueError( + f"Expected single-element list for single LatentData, " + f"got {len(result)} elements" + ) + result = result[0] + + return result + + def _decode_coefficients( + self, latent_data: LatentData, coeffs_u: Union[np.ndarray, Any] + ) -> Any: + """Reconstruct latent activations and return raw formatted predictions.""" + self.check_if_fitted() + + # Convert coeffs_u to framework tensor if needed + if isinstance(coeffs_u, np.ndarray): + coeffs_u = self._to_tensor(coeffs_u, dtype=self._framework_module.float32) + + # Reconstruct activations from concepts + concept_bank_tensor = self._to_tensor( + self.factorization.concept_bank_w, dtype=self._framework_module.float32 + ) + activations = coeffs_u @ concept_bank_tensor + + # Set activations and decode through model + latent_data.set_activations(activations) + return self.latent_extractor.latent_to_logit(latent_data) + + def compute_explanation_per_concept( + self, + images: np.ndarray, + partial_explainer: PartialExplainer, + class_id: Optional[int] = None, + confidence: Optional[float] = None, + verbose: bool = False, + ) -> np.ndarray: + """ + Compute explanations per concept using the provided explainer. + + Wraps the concept decoder in a framework specific wrapper for compatibility + with Xplique attribution methods. + + For each image, creates a concept decoder and uses the specified attribution + method to compute how much each concept contributes to the filtered detections. + + Parameters + ---------- + images + Input images as numpy arrays + partial_explainer + PartialExplainer instance that creates an attribution explainer when called + with model and batch_size arguments + class_id + Target class ID for filtering detections + confidence + Confidence threshold for filtering detections + verbose + If True, prints progress information during processing + + Returns + ------- + explanations + Concatenated explanations for all images, shape (N, H, W, n_concepts) + + Raises + ------ + TypeError + If partial_explainer is not a PartialExplainer instance + """ + if not isinstance(partial_explainer, PartialExplainer): + raise TypeError( + f"partial_explainer must be a PartialExplainer instance," + f" got {type(partial_explainer).__name__}.\n" + f"Wrap your explainer class using PartialExplainer, e.g., " + f"PartialExplainer(GradientInput, operator=my_operator)" + ) + + explanation_list = [] + + # Targets and decoder metadata are prepared per image. Each explainer can + # still batch coefficient perturbations using the configured batch size. + with self.latent_extractor.temporary_force_batch_size(1): + # Encode images to get latent data and concept coefficients + # The list is composed of 1 EncodedData per image because + # object detection models can return various number of + # detection boxes per image + encoded_data_list = self.encode(images) + if not encoded_data_list: + raise ValueError("No latent data extracted from inputs.") + + total_images = len(encoded_data_list) + for i, enc in enumerate(encoded_data_list): + if verbose: + print(f"\rProcessing image {i + 1}/{total_images}...", end="", flush=True) + # Pass 1 (no gradients): plain forward pass to build attribution targets. + decoded_result = self.decode(enc.latent_data, enc.coeffs_u) + filtered_result = decoded_result.filter(class_id=class_id, confidence=confidence) + expected_explanation_shape = tuple(enc.coeffs_u.shape) + is_empty = ( + bool(filtered_result.is_empty) + if hasattr(filtered_result, "is_empty") + else len(filtered_result) == 0 + ) + if is_empty: # No detection + explanation = np.zeros(expected_explanation_shape) + if verbose: + print( + f"\nNo detection for image {i}, returning zero explanation " + f"of shape {explanation.shape}" + ) + else: + targets = self._to_numpy( + filtered_result.to_attribution_target(class_id).to_batched_tensor() + ) + decoder = self.make_concept_decoder(enc.latent_data) + explainer_instance = partial_explainer( + model=decoder, batch_size=self.batch_size + ) + + # Pass 2 (differentiable): explainer calls ConceptDecoder internally to compute + # gradients. Explain the importance of each concept w.r.t the targets. + explanation = explainer_instance.explain(enc.coeffs_u, targets) + explanation = self._to_numpy(explanation) + if explanation.shape != expected_explanation_shape: + raise ValueError( + f"Explanation shape {explanation.shape} does not match expected shape " + f"{expected_explanation_shape} for image {i}. Check that the explainer " + f"and concept decoder are correctly implemented." + ) + explanation_list.append(explanation) + if verbose: + # Print newline after all images are processed + print() + return np.concatenate(explanation_list, axis=0) + + @abstractmethod + def make_concept_decoder(self, latent_data: LatentData) -> Any: + """Creates a concept decoder for gradient-based attribution. + + The decoder is bound to a specific image's latent representation and + accepts concept coefficients as input. Suitable for computing gradients + with respect to concepts using attribution methods like GradientInput. + + Parameters + ---------- + latent_data + Image-specific latent representation + + Returns + ------- + decoder + ConceptDecoder instance with signature: (coeffs_u) -> predictions + """ + raise NotImplementedError + + def _prepare_display_concept_inputs( + self, + images: Union[np.ndarray, List[Any]], + coeffs_u: Optional[np.ndarray], + order: Optional[List[int]], + ) -> Tuple[np.ndarray, np.ndarray, List[int]]: + """Normalise images and coefficients for display methods. + + Computes concept coefficients when not provided, reshapes token-based + coefficients to spatial form, converts images to HWC numpy arrays, and + resolves the ordered list of concept IDs. + + Parameters + ---------- + images + Input images as a batch tensor or list of tensors/arrays + coeffs_u + Pre-computed concept coefficients, or None to compute via transform() + order + Optional list of concept IDs. If None, uses sequential order. + + Returns + ------- + images_np + Images as HWC numpy arrays, shape (N, H, W, C) + coeffs_u + Concept coefficients, shape (N, H, W, n_concepts) + concepts_id + Ordered list of concept IDs to display + """ + # encode images + if coeffs_u is None: + coeffs_u = self.transform(images) + # coeffs_u shape is (N, H, W, C) or (N, Tokens, C) + + if len(coeffs_u.shape) == 3: + # Reshape (N, Tokens, C) to (N, H, W, C) + num_images, num_tokens, num_concepts = coeffs_u.shape + height = width = int(np.sqrt(num_tokens)) # Only valid if Tokens is a perfect square + if height * width != num_tokens: + raise ValueError( + f"Cannot reshape coeffs_u of shape {coeffs_u.shape} to (N, H, W, C) " + f"because Tokens is not a perfect square." + ) + coeffs_u = coeffs_u.reshape(num_images, height, width, num_concepts) + elif len(coeffs_u.shape) != 4: + raise ValueError( + "coeffs_u must have shape (N, H, W, n_concepts) or (N, tokens, n_concepts)" + ) + + if coeffs_u.shape[-1] != self.number_of_concepts: + raise ValueError( + f"coeffs_u contains {coeffs_u.shape[-1]} concepts, expected " + f"{self.number_of_concepts}" + ) + + # convert images to HWC numpy format for display + if self.framework == "torch": + # channel first (C, H, W) -> channel last (H, W, C) for each image + images_np = np.stack([self._to_numpy(img.squeeze().permute(1, 2, 0)) for img in images]) + else: + images_np = np.stack([self._to_numpy(img) for img in images]) + + if order is None: + concepts_id = list(range(self.number_of_concepts)) + else: + try: + concepts_id = list(order) + except TypeError as error: + raise ValueError("order must be an iterable of concept IDs") from error + if not concepts_id: + raise ValueError("order must contain at least one concept ID") + if len(concepts_id) > self.number_of_concepts: + raise ValueError("order cannot contain more IDs than number_of_concepts") + if any( + not isinstance(concept_id, (int, np.integer)) + or isinstance(concept_id, bool) + or not 0 <= concept_id < self.number_of_concepts + for concept_id in concepts_id + ): + raise ValueError( + f"order concept IDs must be integers between 0 and " + f"{self.number_of_concepts - 1}" + ) + if len(set(concepts_id)) != len(concepts_id): + raise ValueError("order cannot contain duplicate concept IDs") + + return images_np, coeffs_u, concepts_id + + def display_concept_heatmap( + self, + image: np.ndarray, + concept_heatmap: np.ndarray, + concept_idx: int, + ax: Any, + filter_percentile: int = 80, + clip_percentile: int = 5, + ) -> None: + """Overlay a single concept heatmap on a single image. + + Displays the image on the given axis, then overlays the concept activation + heatmap after filtering low activations, resizing to image resolution, and + clipping outlier values. + + Parameters + ---------- + image + Single image as HWC numpy array, shape (H, W, C) + concept_heatmap + Raw concept activation map, shape (H', W') + concept_idx + Index of the concept, used to select the colormap + ax + Matplotlib axis on which to draw + filter_percentile + Percentile used to filter the concept heatmap + (only show concept if excess N-th percentile). Defaults to 80. + clip_percentile + Percentile value to use if clipping is needed when drawing the concept, + e.g a value of 1 will perform a clipping between percentile 1 and 99. + This parameter allows to avoid outliers in case of too extreme values. + Default to 5. + """ + dsize = (image.shape[1], image.shape[0]) # cv2 expects (width, height) + + # Display the image + show_ax(image, ax=ax) + + # only show concept if excess N-th percentile + sigma = np.percentile(concept_heatmap.flatten(), filter_percentile) + heatmap = concept_heatmap * (concept_heatmap > sigma) + + # resize the heatmap before clipping + heatmap = cv2.resize(heatmap[:, :, None], dsize=dsize, interpolation=cv2.INTER_CUBIC) + heatmap = _clip_percentile(heatmap, clip_percentile) + + # Display the heatmap overlay + cmap_idx = concept_idx % len(self.cmaps) + show_ax(heatmap, cmap=self.cmaps[::-1][cmap_idx], alpha=0.5, ax=ax) + + def display_images_per_concept( + self, + images: np.ndarray, + coeffs_u: Optional[np.ndarray] = None, + filter_percentile: int = 80, + clip_percentile: int = 5, + order: Optional[List[int]] = None, + ) -> Figure: + """ + Display concept heatmaps overlaid on images. + + Creates a grid visualization with one row per image and one column per + concept. Each cell shows the input image with a heatmap overlay indicating + where that concept is activated. + + Parameters + ---------- + images + Input images to visualize (array of shape (N, H, W, C) for tensorflow or (N, C, H, W) + for pytorch) + coeffs_u + Optional pre-computed coefficients, shape (N, H, W, C) or (N, Tokens, C). + If None, coefficients will be computed via transform(images). + filter_percentile + Percentile used to filter the concept heatmap + (only show concept if excess N-th percentile). Defaults to 80. + clip_percentile + Percentile value to use if clipping is needed when drawing the concept, + e.g a value of 1 will perform a clipping between percentile 1 and 99. + This parameter allows to avoid outliers in case of too extreme values. + Default to 5. + order + Optional list of concept IDs to specify display order. If None, + concepts are shown in sequential order + + Returns + ------- + fig + matplotlib figure with len(images) rows and number_of_concepts columns + """ + images_np, coeffs_u, concepts_id = self._prepare_display_concept_inputs( + images, coeffs_u, order + ) + + nb_cols = len(concepts_id) + nb_rows = len(images_np) + + fig, axs = plt.subplots(nb_rows, nb_cols, figsize=(2 * nb_cols, 2 * nb_rows)) + axs = np.asarray(axs).reshape(nb_rows, nb_cols) + + for i, c_i in enumerate(concepts_id): + axs[0, i].set_title(f"concept #{c_i}", fontsize=10) + + # Display a heatmap per concept, per image + for i, c_i in enumerate(concepts_id): + for image_id, image in enumerate(images_np): + self.display_concept_heatmap( + image=image, + concept_heatmap=coeffs_u[image_id, :, :, c_i], + concept_idx=c_i, + ax=axs[image_id, i], + filter_percentile=filter_percentile, + clip_percentile=clip_percentile, + ) + return fig + + def get_topk_images_per_concept( + self, + coeffs_u: np.ndarray, + topk: int = 3, + ) -> np.ndarray: + """Return the indices of the top images for each concept, ranked by mean activation. + + Parameters + ---------- + coeffs_u + Concept coefficients, shape (N, H, W, n_concepts) + topk + Number of top images to return per concept (default: 3) + + Returns + ------- + top_image_ids + Array of shape (n_concepts, topk) containing the indices of the top images + for each concept, ranked by descending mean activation + """ + # Compute mean activation per image per concept: (N, n_concepts) + mean_activations = np.mean(coeffs_u, axis=(1, 2)) + + # For each concept, find the top-k image indices by descending mean activation + top_image_ids = np.argsort(mean_activations, axis=0)[::-1, :][:topk, :].T + # top_image_ids shape: (n_concepts, topk) + return top_image_ids + + def display_top_images_per_concept( + self, + images: Union[np.ndarray, List[Any]], + topk: int = 3, + filter_percentile: int = 80, + clip_percentile: int = 5, + order: Optional[List[int]] = None, + coeffs_u: Optional[np.ndarray] = None, + ) -> Figure: + """Display top N images per concept ranked by average activation. + + Parameters + ---------- + images + Input images (as framework tensors or numpy arrays) + topk + Number of top images to display per concept (default: 3) + filter_percentile + Percentile threshold for filtering heatmaps (default: 80) + clip_percentile + Percentile for clipping heatmap values (default: 5) + order + Optional list of concept IDs to specify display order + coeffs_u + Optional pre-computed concept coefficients. If None, will call + self.transform(images) to compute them. Use this to pass the + coefficients stored in factorization.coeffs_u after fit(). + + Returns + ------- + fig + matplotlib figure with topk rows and number_of_concepts columns + """ + images_np, coeffs_u, concepts_id = self._prepare_display_concept_inputs( + images, coeffs_u, order + ) + + nb_rows = topk + nb_cols = len(concepts_id) + fig, axs = plt.subplots(nb_rows, nb_cols, figsize=(2 * nb_cols, 2 * nb_rows)) + axs = np.asarray(axs).reshape(nb_rows, nb_cols) + + for i, c_i in enumerate(concepts_id): + axs[0, i].set_title(f"concept #{c_i}", fontsize=10) + + # Get top image indices for all concepts at once: (n_concepts, topk) + topk_images_ids = self.get_topk_images_per_concept(coeffs_u, topk) + + for i, c_i in enumerate(concepts_id): + for j, image_id in enumerate(topk_images_ids[c_i]): + self.display_concept_heatmap( + image=images_np[image_id], + concept_heatmap=coeffs_u[image_id, :, :, c_i], + concept_idx=c_i, + ax=axs[j, i], + filter_percentile=filter_percentile, + clip_percentile=clip_percentile, + ) + + return fig + + def estimate_importance( + self, + images: Union[np.ndarray, List[Any]], + operator: Callable, + class_id: int, + method: str = "gradient_input", + confidence: float = 0.9, + spatial_reducer: Optional[str] = "mean", + abs_before_reduce: bool = True, + aggregation_reducer: Optional[str] = "mean", + verbose: bool = False, + **method_kwargs: Any, + ) -> np.ndarray: + """ + Estimate concept importance using the specified attribution method. + + Parameters + ---------- + images + Input images to analyze + operator + Function to extract target values from predictions + class_id + Target class ID for filtering detections + confidence + Confidence threshold for filtering detections (default: 0.9) + method + Attribution method: "gradient_input" or "sobol" (default: "gradient_input") + spatial_reducer + Reducer to use on the spatial dimension of the raw explanations. + Explanation has shape (num_images, height, width, num_concepts) and will be reduced + to (num_images, num_concepts) before final aggregation. + Either "min", "mean", "max", "sum", or `None` to ignore. Default is "mean". + abs_before_reduce + Whether to take the absolute value of the explanations before spatial reduction + (default: True) + aggregation_reducer + Reducer to use on the image dimension after spatial reduction. + Either "min", "mean", "max", "sum", or `None` to ignore. Default is "mean". + Transform spatial explanations from shape (num_images, num_concepts) to the final + importances (num_concepts,). + verbose + If True, prints progress information during processing + **method_kwargs + Additional keyword arguments for the attribution method, such as 'grid_size' + and 'nb_design' for the Sobol method + + Returns + ------- + importances + Importance scores for each concept, shape (n_concepts,) + """ + if method == "gradient_input": + method_kwargs.setdefault("operator", operator) + method_kwargs.setdefault("reducer", None) + explainer = PartialExplainer(GradientInput, **method_kwargs) + elif method == "sobol": + # set default values for Sobol-specific parameters if not provided + method_kwargs.setdefault("grid_size", 8) + method_kwargs.setdefault("nb_design", 32) + method_kwargs.setdefault("perturbation_function", "amplitude") + method_kwargs.setdefault("nb_channels", self.number_of_concepts) + method_kwargs.setdefault("operator", operator) + explainer = PartialExplainer(SobolAttributionMethod, **method_kwargs) + else: + raise ValueError(f"Unknown attribution method: {method}") + + # explainer is a PartialExplainer that creates an explainer with + # all the necessary arguments except the model and bach_size which + # will be provided in compute_explanation_per_concept when the + # concept decoder is created + explanation = self.compute_explanation_per_concept( + images, explainer, class_id, confidence, verbose + ) + return self.reduce_to_importance( + explanation, + spatial_reducer=spatial_reducer, + abs_before_reduce=abs_before_reduce, + aggregation_reducer=aggregation_reducer, + ) + + def reduce_to_importance( + self, + explanation: np.ndarray, + spatial_reducer: Optional[str] = "max", + abs_before_reduce: bool = True, + aggregation_reducer: Optional[str] = "mean", + ) -> np.ndarray: + """ + Reduce pre-computed concept explanations to global importance scores. + + Parameters + ---------- + explanation + Per-concept explanations, shape (N, H, W, n_concepts), as returned by + :meth:`compute_explanation_per_concept`. + spatial_reducer + Reducer applied over the spatial dimensions (H, W) to collapse each image + to a per-concept score vector. Either "min", "mean", "max", "sum", "median" + or `None` to skip. Default is "max". + abs_before_reduce + Whether to take the absolute value of explanations before spatial reduction. + Default is True. + aggregation_reducer + Reducer applied over the image dimension after spatial reduction to produce + a single score per concept. Either "min", "mean", "max", "sum", "median" + or `None` to skip (returns per-image scores). Default is "mean". + + Returns + ------- + importances + Importance scores for each concept, shape (n_concepts,). + """ + reducers = { + "min": np.min, + "max": np.max, + "sum": np.sum, + "mean": np.mean, + "median": np.median, + } + if abs_before_reduce: + explanation = np.abs(explanation) + spatial_axes = tuple(range(1, explanation.ndim - 1)) + if spatial_reducer is not None and spatial_axes: + explanation = reducers[spatial_reducer](explanation, axis=spatial_axes) + if aggregation_reducer is not None: + importances = reducers[aggregation_reducer](explanation, axis=0) + else: + importances = explanation + return importances + + def reduce_to_prevalence(self, explanation: np.ndarray) -> np.ndarray: + """ + Compute concept prevalence from pre-computed explanations. + + A concept is prevalent when it is frequently the most important one across + images, i.e. it dominates the most samples (argmax of per-image importance). + + Ref. Fel et al., A Holistic Approach to Unifying Automatic Concept Extraction + and Concept Importance Estimation (2023). + https://arxiv.org/pdf/2306.07304 + + Parameters + ---------- + explanation + Per-concept explanations, shape (N, H, W, n_concepts), as returned by + :meth:`compute_explanation_per_concept`. + + Returns + ------- + prevalence + Fraction of images for which each concept is dominant, shape (n_concepts,). + Values sum to 1. + """ + spatial_axes = tuple(range(1, explanation.ndim - 1)) + per_image = ( + np.mean(explanation, axis=spatial_axes) if spatial_axes else explanation + ) # (N, n_concepts) + dominant = np.argmax(per_image, axis=-1) # (N,) + prevalence = np.zeros(self.number_of_concepts) + for c in range(self.number_of_concepts): + prevalence[c] = np.sum(dominant == c) / len(dominant) + return prevalence + + def reduce_to_reliability(self, explanation: np.ndarray, accuracy: np.ndarray) -> np.ndarray: + """ + Compute concept reliability from pre-computed explanations and per-image accuracy. + + A concept is reliable when images for which it is the most important concept + also tend to be correctly predicted. Reliability is the mean accuracy of the + group of images sharing the same dominant concept. + + Ref. Fel et al., A Holistic Approach to Unifying Automatic Concept Extraction + and Concept Importance Estimation (2023). + https://arxiv.org/pdf/2306.07304 + + Parameters + ---------- + explanation + Per-concept explanations, shape (N, H, W, n_concepts), as returned by + :meth:`compute_explanation_per_concept`. + accuracy + Per-image accuracy scores, shape (N,). For classification: 0.0 or 1.0. + For object detection: per-image IoU, AP, or any scalar correctness metric + computed externally by the user. + + Returns + ------- + reliability + Mean accuracy per dominant-concept group, shape (n_concepts,). + Concepts with no dominant image get a reliability of 0.0. + """ + accuracy = np.asarray(accuracy) + spatial_axes = tuple(range(1, explanation.ndim - 1)) + per_image = ( + np.mean(explanation, axis=spatial_axes) if spatial_axes else explanation + ) # (N, n_concepts) + dominant = np.argmax(per_image, axis=-1) # (N,) + reliability = np.zeros(self.number_of_concepts) + for c in range(self.number_of_concepts): + mask = dominant == c + reliability[c] = np.mean(accuracy[mask]) if mask.any() else 0.0 + return reliability + + +class ConceptDecoder: + """ + Concept decoder module. + + Converts concept coefficients back to object detection predictions by + reconstructing activations and passing them through the decoder network. + + Attributes + ---------- + parent_craft + HolisticCraft instance used to decode predictions + latent_data + Image-specific latent representation to use for decoding + """ + + parent_craft: HolisticCraft + latent_data: LatentData + + def set_latent_data(self, latent_data: LatentData) -> None: + """ + Update the latent data for this decoder. + + Parameters + ---------- + latent_data + New latent representation to use + """ + self.latent_data = latent_data + + def _decode(self, coeffs_u): + """ + Decode concept coefficients to predictions + + Parameters + ---------- + coeffs_u + Batched concept coefficients. + + Returns + ------- + logits + Predictions as a dense batched tensor. Object detections are zero-padded + to the largest number of boxes in the batch. + """ + predictions = self.parent_craft._decode_coefficients(self.latent_data, coeffs_u) + return self._predictions_to_tensor(predictions) + + def _predictions_to_tensor(self, predictions): + """Convert framework-specific structured predictions to a dense tensor.""" + raise NotImplementedError diff --git a/xplique/concepts/latent_extractor.py b/xplique/concepts/latent_extractor.py new file mode 100644 index 00000000..1b736871 --- /dev/null +++ b/xplique/concepts/latent_extractor.py @@ -0,0 +1,265 @@ +""" +Base classes for latent extraction in deep learning models. +""" + +import contextlib +from abc import ABC, abstractmethod +from typing import Any, Callable, Generator, NamedTuple, Optional, Union + +import numpy as np + +from xplique.utils_functions.object_detection.base.box_formatter import ( + BaseBoxFormatter, +) + + +class LatentData(ABC): + """Abstract base class for storing latent representations and activations. + + This class provides an interface for managing intermediate activations + extracted from a model's latent space. Subclasses must implement methods + to get and set activations in a framework-specific manner. + """ + + @abstractmethod + def get_activations(self, as_numpy: bool = True, keep_gradients: bool = False): + """Retrieve the activations. + + Parameters + ---------- + as_numpy + If True, return activations as a NumPy array. If False, return them as a + framework-native tensor (e.g., TensorFlow or PyTorch). Default is True. + keep_gradients + If True, preserve gradients for backpropagation. Default is False. + + Returns + ------- + activations + The activations following TensorFlow ordering (batch, height, width, channels). + """ + raise NotImplementedError("get_activations method must be implemented by subclasses") + + @abstractmethod + def set_activations(self, values: np.ndarray): + """Set the activations. + + Parameters + ---------- + values + A NumPy array containing the activations, following TensorFlow ordering + (batch, height, width, channels). + """ + raise NotImplementedError("set_activations method must be implemented by subclasses") + + +class EncodedData(NamedTuple): + """Encoded representation containing latent data and concept coefficients. + + This named tuple is returned by the encode() method and contains both the + intermediate latent representation and the concept coefficients from NMF + factorization. + + Attributes + ---------- + latent_data + Intermediate latent representation from the model, containing activations + and metadata needed for decoding back to the original space. + coeffs_u + Concept coefficients (U matrix from NMF factorization). + - When differentiable=False: NumPy array + - When differentiable=True: Framework-specific tensor (torch.Tensor or + tf.Tensor) with gradients preserved for backpropagation. + """ + + latent_data: "LatentData" + coeffs_u: Union[np.ndarray, Any] # Framework-specific tensor type for differentiable case + + +class LatentExtractor(ABC): + """Extracts and manages latent representations from models. + + This class provides functionality to split a model into two parts: + - input_to_latent: extracts intermediate activations + - latent_to_logit: processes activations to final predictions + + This splitting enables concept-based explanations by allowing manipulation + of the latent space between the two model parts. + + Parameters + ---------- + model + The full model to be split + input_to_latent_model + Model segment that processes inputs to latent activations + latent_to_logit_model + Model segment that processes latent activations to final outputs + latent_data_class + Class for storing latent data (default: LatentData) + output_formatter + Optional formatter for model outputs + batch_size + Batch size for processing (default: 8) + """ + + def __init__( + self, + model: Callable, + input_to_latent_model: Callable, + latent_to_logit_model: Callable, + latent_data_class=LatentData, + output_formatter: Optional[BaseBoxFormatter] = None, + batch_size: int = 8, + ): + self.model = model + self.input_to_latent_model = input_to_latent_model + self.latent_to_logit_model = latent_to_logit_model + self.latent_data_class = latent_data_class + self.output_formatter = output_formatter + self.batch_size = batch_size + + def __call__(self, *args, **kwargs): + """Make the extractor callable, forwarding to forward method. + + Returns + ------- + predictions + Model predictions + """ + return self.forward(*args, **kwargs) + + @abstractmethod + def input_to_latent(self, inputs) -> LatentData: + """Extract latent representations from inputs. + + Parameters + ---------- + inputs + Input data to process + + Returns + ------- + latent_data + Latent representation containing activations + + Raises + ------ + NotImplementedError + This method must be implemented by subclasses + """ + raise NotImplementedError("This method should be implemented in subclasses.") + + @abstractmethod + def input_to_latent_generator( + self, inputs, resize: Optional[tuple] = None, keep_gradients: bool = False + ) -> Generator[LatentData, None, None]: + """ + Generator that yields latent data batch by batch. + + Parameters + ---------- + inputs + Input images as PyTorch tensors. + resize + Optional target size for resizing inputs. Default is None. + keep_gradients + If True, preserve gradients during processing. Default is False. + + Yields + ------ + latent_data + LatentData object for each batch, with automatic memory management. + """ + raise NotImplementedError("This method should be implemented in subclasses.") + + @abstractmethod + def latent_to_logit(self, latent_data: LatentData): + """Convert latent representations to final predictions. + + Parameters + ---------- + latent_data + Latent representation containing activations + + Returns + ------- + predictions + Model output predictions + + Raises + ------ + NotImplementedError + This method must be implemented by subclasses + """ + raise NotImplementedError("This method should be implemented in subclasses.") + + def forward(self, samples): + """Forward pass through the full model via latent space. + + Parameters + ---------- + samples + Input samples to process + + Returns + ------- + predictions + Model predictions + """ + latent_data = self.input_to_latent(samples) + return self.latent_to_logit(latent_data) + + @contextlib.contextmanager + def temporary_force_batch_size(self, batch_size: int): + """Context manager to temporarily set the batch size. + + This is used during encoding to process one sample at a time, when we need + to compute explanations per sample. + + Parameters + ---------- + batch_size + Temporary batch size to use within the context + + Yields + ------ + None + Context for processing with temporary batch size + """ + old_batch_size = self.batch_size + self.batch_size = batch_size + try: + yield + finally: + self.batch_size = old_batch_size + + +class LatentExtractorBuilder(ABC): + """Abstract base class ensuring all builders return LatentExtractor instances. + + This class provides a common interface for building LatentExtractor objects + in a framework-specific manner. Subclasses implement the build method to + construct appropriate extractors for different model architectures. + """ + + @classmethod + @abstractmethod + def build(cls, **kwargs) -> "LatentExtractor": + """Build and return a LatentExtractor. + + Parameters + ---------- + **kwargs + Framework-specific arguments for building the extractor + + Returns + ------- + extractor + Configured LatentExtractor instance + + Raises + ------ + NotImplementedError + This method must be implemented by subclasses + """ + raise NotImplementedError("build method must be implemented by subclasses") diff --git a/xplique/concepts/tf/__init__.py b/xplique/concepts/tf/__init__.py new file mode 100644 index 00000000..9490798c --- /dev/null +++ b/xplique/concepts/tf/__init__.py @@ -0,0 +1,11 @@ +""" +TensorFlow-specific implementations for concepts module. +""" + +from .holistic_craft import HolisticCraftTf +from .latent_extractor import TfLatentExtractor + +__all__ = [ + "TfLatentExtractor", + "HolisticCraftTf", +] diff --git a/xplique/concepts/tf/factorizer.py b/xplique/concepts/tf/factorizer.py new file mode 100644 index 00000000..4fdbd747 --- /dev/null +++ b/xplique/concepts/tf/factorizer.py @@ -0,0 +1,139 @@ +""" +TensorFlow-specific factorizer implementations +""" + +import numpy as np +import tensorflow as tf + +from ..factorizer import SklearnNMFFactorizer + + +class TfSklearnNMFFactorizer(SklearnNMFFactorizer): + """ + TensorFlow-compatible sklearn NMF factorizer with differentiable encoding. + """ + + def encode_differentiable(self, activations: tf.Tensor) -> tf.Tensor: + """ + Encode activations with a differentiable non-negative solver. + + Solves the fixed-dictionary NMF subproblem with restartable FISTA so + gradients still flow through ``activations``. + + Parameters + ---------- + activations : tf.Tensor + Activations to encode, shape (n_samples, n_features) + + Returns + ------- + tf.Tensor + Coefficients, shape (n_samples, n_concepts) + """ + if self._concept_bank_w is None: + raise ValueError("Factorizer must be fitted before encoding") + + if not isinstance(activations, tf.Tensor): + activations = tf.convert_to_tensor(activations) + + beta_loss = self.nmf_kwargs.get("beta_loss", "frobenius") + if beta_loss != "frobenius": + raise NotImplementedError( + "TensorFlow differentiable NMF encoding only supports beta_loss='frobenius'" + ) + + tf.debugging.assert_non_negative( + activations, message="NMF requires non-negative activations" + ) + + concept_bank_tensor = tf.constant(self._concept_bank_w, dtype=activations.dtype) + dtype = activations.dtype + eps = tf.cast(1e-8, dtype) + one = tf.cast(1.0, dtype) + four = tf.cast(4.0, dtype) + + alpha_w = tf.cast(self.nmf_kwargs.get("alpha_W", 0.0), dtype) + l1_ratio = tf.cast(self.nmf_kwargs.get("l1_ratio", 0.0), dtype) + n_features = tf.cast(tf.shape(activations)[1], dtype) + l1_reg = n_features * alpha_w * l1_ratio + l2_reg = n_features * alpha_w * (1.0 - l1_ratio) + + max_iter = int(self.nmf_kwargs.get("max_iter", 200)) + tol = float(self.nmf_kwargs.get("tol", 1e-4)) + + gram = concept_bank_tensor @ tf.transpose(concept_bank_tensor) + cross = activations @ tf.transpose(concept_bank_tensor) + diagonal = tf.linalg.diag_part(gram) + coefficients = tf.nn.relu(cross / (diagonal[tf.newaxis, :] + l2_reg + eps)) + + lipschitz = tf.linalg.norm(gram, ord="fro", axis=(0, 1)) + l2_reg + eps + step = one / lipschitz + + def fista_step(coeffs, extrapolated_coeffs, momentum): + gradient = extrapolated_coeffs @ gram - cross + l2_reg * extrapolated_coeffs + updated = tf.nn.relu(extrapolated_coeffs - step * gradient - step * l1_reg) + coeff_delta = tf.linalg.norm(updated - coeffs) / (tf.linalg.norm(updated) + eps) + + next_momentum = (one + tf.sqrt(one + four * tf.square(momentum))) / 2.0 + accelerated = updated + ((momentum - one) / next_momentum) * (updated - coeffs) + restart = tf.reduce_sum((extrapolated_coeffs - updated) * (updated - coeffs)) > 0 + + next_extrapolated = tf.where(restart, updated, accelerated) + next_momentum = tf.where(restart, one, next_momentum) + + return updated, next_extrapolated, next_momentum, coeff_delta + + extrapolated_coeffs = coefficients + momentum = one + + if tol > 0.0: + max_iter_tensor = tf.constant(max_iter) + tol_tensor = tf.cast(tol, dtype) + delta = tf.constant(float("inf"), dtype=dtype) + + def cond(iteration, coeffs, _, __, coeff_delta): + return tf.logical_and(iteration < max_iter_tensor, coeff_delta > tol_tensor) + + def body(iteration, coeffs, extrapolated, current_momentum, _): + updated, next_extrapolated, next_momentum, coeff_delta = fista_step( + coeffs, extrapolated, current_momentum + ) + return iteration + 1, updated, next_extrapolated, next_momentum, coeff_delta + + _, coefficients, _, _, _ = tf.while_loop( + cond, + body, + (tf.constant(0), coefficients, extrapolated_coeffs, momentum, delta), + parallel_iterations=1, + ) + else: + for _ in range(max_iter): + coefficients, extrapolated_coeffs, momentum, _ = fista_step( + coefficients, extrapolated_coeffs, momentum + ) + + return coefficients + + def decode(self, coefficients: tf.Tensor) -> tf.Tensor: + """ + Decode coefficients to activations via matrix multiplication. + + Parameters + ---------- + coefficients : tf.Tensor + Coefficients to decode, shape (n_samples, n_concepts) + + Returns + ------- + tf.Tensor + Reconstructed activations, shape (n_samples, n_features) + """ + if self._concept_bank_w is None: + raise ValueError("Factorizer must be fitted before decoding") + + if isinstance(coefficients, np.ndarray): + return coefficients @ self._concept_bank_w + + concept_bank_tensor = tf.constant(self._concept_bank_w, dtype=coefficients.dtype) + + return coefficients @ concept_bank_tensor diff --git a/xplique/concepts/tf/holistic_craft.py b/xplique/concepts/tf/holistic_craft.py new file mode 100644 index 00000000..b3b41521 --- /dev/null +++ b/xplique/concepts/tf/holistic_craft.py @@ -0,0 +1,220 @@ +"""TensorFlow-specific wrapper for HolisticCraft.""" + +from typing import Any, Optional, Union + +import numpy as np +import tensorflow as tf + +from xplique.concepts.factorizer import ConceptFactorizer +from xplique.utils_functions.object_detection.tf.box_model_wrapper import ( + _pad_and_stack_box_predictions, +) + +from ..holistic_craft import ConceptDecoder, HolisticCraft +from ..latent_extractor import LatentData +from .factorizer import TfSklearnNMFFactorizer +from .latent_extractor import TfLatentExtractor as LatentExtractor + + +class HolisticCraftTf(HolisticCraft): + """ + TensorFlow-specific implementation of CRAFT for holistic model explanations. + + This class is a thin wrapper around the framework-agnostic base class. + All core functionality is inherited from HolisticCraft. + + Parameters + ---------- + latent_extractor + TensorFlow latent extractor for the model. + number_of_concepts + Number of concepts to extract, by default 20. + factorizer + Optional factorizer instance. If None, creates a TfSklearnNMFFactorizer + with alpha_W=1e-2 and max_iter=200 + """ + + def __init__( + self, + latent_extractor: LatentExtractor, + number_of_concepts: int = 20, + factorizer: Optional[ConceptFactorizer] = None, + ) -> None: + """ + Initialize the TensorFlow CRAFT wrapper. + + Parameters + ---------- + latent_extractor + TensorFlow latent extractor for the model + number_of_concepts + Number of concepts to extract (default: 20) + factorizer + Optional factorizer instance. If None, creates a TfSklearnNMFFactorizer + with alpha_W=1e-2 and max_iter=200 + """ + + # Create TensorFlow-specific factorizer if none provided + if factorizer is None: + factorizer = TfSklearnNMFFactorizer( + n_components=number_of_concepts, alpha_W=1e-2, max_iter=200 + ) + + super().__init__(latent_extractor, number_of_concepts, device=None, factorizer=factorizer) + self.framework = "tf" + self._framework_module = tf + + def latent_to_concept_differentiable(self, latent_data: LatentData) -> tf.Tensor: + """ + Transform latent data to concept coefficients with gradient preservation. + + TensorFlow-specific implementation using a differentiable non-negative + optimization solver. Maintains the gradient tape for computing + attributions with respect to concepts. + + Parameters + ---------- + latent_data + Single image's latent representation containing activations + + Returns + ------- + coeffs_u + Concept coefficients as TensorFlow tensor with gradients preserved + + Raises + ------ + ValueError + If latent_data is not a single LatentData instance + NotFittedError + If fit() has not been called yet + """ + if not isinstance(latent_data, LatentData): + raise ValueError( + f"latent_to_concept_differentiable() only accepts a single " + f"LatentData as input, got {type(latent_data)}" + ) + self.check_if_fitted() + + # Get activations as tensors with gradients preserved + activations = latent_data.get_activations(as_numpy=False, keep_gradients=True) + + # Ensure we have TensorFlow tensors + if not isinstance(activations, tf.Tensor): + activations = tf.convert_to_tensor(activations) + + activations_original_shape = activations.shape[:-1] + activations_flat = tf.reshape(activations, (-1, activations.shape[-1])) + + # Use factorizer's differentiable encoding + coeffs_u = self.factorizer.encode_differentiable(activations_flat) + + # Reshape back to original dimensions + coeffs_u = tf.reshape(coeffs_u, tf.concat([activations_original_shape, [-1]], axis=0)) + return coeffs_u + + def _to_numpy(self, tensor: Union[tf.Tensor, np.ndarray]) -> np.ndarray: + """ + Convert TensorFlow tensor to numpy array. + + Parameters + ---------- + tensor + TensorFlow tensor or numpy array + + Returns + ------- + array + Numpy array + """ + if isinstance(tensor, np.ndarray): + return tensor + return tensor.numpy() + + def _to_tensor(self, array: np.ndarray, dtype: Optional[tf.DType] = None) -> tf.Tensor: + """ + Convert numpy array to TensorFlow tensor. + + Parameters + ---------- + array + Numpy array to convert + dtype + Target TensorFlow dtype (e.g., tf.float32) + + Returns + ------- + tensor + TensorFlow tensor + """ + kwargs = {} + if dtype is not None: + kwargs["dtype"] = dtype + return tf.convert_to_tensor(array, **kwargs) + + def make_concept_decoder(self, latent_data: LatentData) -> tf.keras.layers.Layer: + """ + Creates a TensorFlow concept decoder for gradient-based attribution. + + The decoder is a Keras Layer that accepts concept coefficients and returns + detection predictions. It maintains a reference to the latent_data to + reconstruct activations during the call. + + Parameters + ---------- + latent_data + Single image's latent representation + + Returns + ------- + decoder + ConceptDecoderTf instance (Keras Layer) with call method + """ + + return ConceptDecoderTf(self, latent_data) + + +class ConceptDecoderTf(tf.keras.layers.Layer, ConceptDecoder): + """ + TensorFlow concept decoder layer. + + Converts concept coefficients back to object detection predictions by + reconstructing activations and passing them through the decoder network. + + Parameters + ---------- + latent_data + Image-specific latent representation to use for decoding + **kwargs + Additional keyword arguments for Keras Layer + """ + + def __init__(self, parent_craft: HolisticCraft, latent_data: LatentData, **kwargs: Any) -> None: + super().__init__(**kwargs) + self.parent_craft = parent_craft + self.latent_data = latent_data + + def call(self, coeffs_u: tf.Tensor) -> tf.Tensor: + """ + Decode concept coefficients to predictions. + + Parameters + ---------- + coeffs_u + Batched concept coefficients. + + Returns + ------- + logits + Predictions as a dense batched tensor. Object detections are zero-padded + to the largest number of boxes in the batch. + """ + + return self._decode(coeffs_u) + + def _predictions_to_tensor(self, predictions) -> tf.Tensor: + if isinstance(predictions, (list, tuple)): + return _pad_and_stack_box_predictions(predictions) + if hasattr(predictions, "to_batched_tensor"): + return predictions.to_batched_tensor() + return predictions diff --git a/xplique/concepts/tf/latent_extractor.py b/xplique/concepts/tf/latent_extractor.py new file mode 100644 index 00000000..f93ea9ef --- /dev/null +++ b/xplique/concepts/tf/latent_extractor.py @@ -0,0 +1,204 @@ +"""TensorFlow-specific latent extractor for object detection models.""" + +from typing import Callable, Generator, List, Optional, Tuple, Union + +import tensorflow as tf + +from xplique.utils_functions.object_detection.base.box_formatter import ( + BaseBoxFormatter, +) +from xplique.utils_functions.object_detection.tf.box_model_wrapper import ( + _pad_and_stack_box_predictions, +) +from xplique.utils_functions.object_detection.tf.multi_box_tensor import TfMultiBoxTensor +from xplique.utils_functions.output_as_list_mixin import OutputAsListMixin + +from ..latent_extractor import LatentData, LatentExtractor + + +class TfLatentExtractor(OutputAsListMixin, LatentExtractor): + """ + TensorFlow-specific latent extractor for object detection models. + + This class provides TensorFlow-specific implementations for extracting intermediate + activations from object detection models and decoding them back to predictions. + It handles batching, resizing, and output formatting for TensorFlow models. + + Parameters + ---------- + model + Complete TensorFlow model (for reference, not directly used) + input_to_latent_model + TensorFlow model/function that maps inputs to latent activations + latent_to_logit_model + TensorFlow model/function that maps latent activations to predictions + latent_data_class + Class to use for storing latent data (default: LatentData) + output_formatter + Formatter to convert raw model outputs to standardized box format + batch_size + Number of samples to process at once + + Attributes + ---------- + output_as_list + Whether to return outputs as list of MultiBoxTensor (True) or + as stacked tensor (False) + """ + + def __init__( + self, + model: Callable, + input_to_latent_model: Callable, + latent_to_logit_model: Callable, + latent_data_class=LatentData, + output_formatter: Optional[BaseBoxFormatter] = None, + batch_size: int = 8, + ) -> None: + if not isinstance(batch_size, int) or isinstance(batch_size, bool) or batch_size <= 0: + raise ValueError("batch_size must be a positive integer") + super().__init__( + model, + input_to_latent_model, + latent_to_logit_model, + latent_data_class, + output_formatter, + batch_size, + ) + self.output_as_list = True + + @staticmethod + def _prepare_inputs(inputs: tf.Tensor) -> tf.Tensor: + """Validate image inputs and add a batch dimension to one image.""" + inputs = tf.convert_to_tensor(inputs) + if inputs.shape.rank not in (3, 4): + raise ValueError( + "inputs must have rank 3 (a single image) or rank 4 (a batch of images)" + ) + if inputs.shape.num_elements() == 0: + raise ValueError("inputs must contain at least one value") + if inputs.shape.rank == 3: + inputs = tf.expand_dims(inputs, axis=0) + return inputs + + def forward(self, samples: tf.Tensor) -> Union[List["TfMultiBoxTensor"], tf.Tensor]: + """ + Process samples through the complete model pipeline. + + Encodes inputs to latent representations, decodes to predictions, and + optionally formats outputs. Return format depends on output_as_list flag. + + Parameters + ---------- + samples + Input images as TensorFlow tensors + + Returns + ------- + outputs + If output_as_list=True: List of MultiBoxTensor (one per image) + If output_as_list=False: Zero-padded tensor of shape (N, max_num_boxes, features) + """ + latent_data = self.input_to_latent(samples) + outputs = self.latent_to_logit_model(latent_data) + if self.output_formatter: + outputs = self.output_formatter(outputs) + if not self.output_as_list: + if isinstance(outputs, (list, tuple)): + outputs = _pad_and_stack_box_predictions(outputs) + elif hasattr(outputs, "to_batched_tensor"): + outputs = outputs.to_batched_tensor() + else: + outputs = tf.expand_dims(outputs, axis=0) + return outputs + + def input_to_latent(self, inputs: tf.Tensor) -> LatentData: + """ + Extract latent representations from input images. + + Encodes single or batched input images into their intermediate latent + representations. Automatically handles 3D inputs by adding batch dimension. + + Parameters + ---------- + inputs + Input image(s) as TensorFlow tensor. Shape: (H, W, C) or (N, H, W, C) + + Returns + ------- + latent_data + Extracted latent activations wrapped in LatentData container + """ + inputs = self._prepare_inputs(inputs) + latent_data = self.input_to_latent_model(inputs) + return latent_data + + def input_to_latent_generator( + self, + inputs: tf.Tensor, + resize: Optional[Tuple[int, int]] = None, + keep_gradients: bool = False, + ) -> Generator[LatentData, None, None]: + # pylint: disable=unused-argument + """ + Generator that yields latent representations for batched inputs. + + Internal generator method that splits inputs into batches, optionally resizes, + encodes to latent space, and yields results incrementally. Efficiently handles + large datasets by processing one batch at a time. + + Parameters + ---------- + inputs + Input images as TensorFlow tensor. Shape: (N, H, W, C) + resize + Target size (height, width) for resizing images before encoding. + If None, uses original image sizes. + keep_gradients + Whether to keep gradients during processing (for gradient-based methods) + + Yields + ------ + latent_data + LatentData object containing encoded activations for current batch + """ + inputs = self._prepare_inputs(inputs) + batch_count = inputs.shape[0] + if batch_count is None: + raise ValueError("inputs must have a known batch dimension") + + for i in range(0, batch_count, self.batch_size): + i_end = min(i + self.batch_size, batch_count) + batch = inputs[i:i_end] + + if resize: + batch = tf.image.resize(batch, size=resize) + + latent_data = self.input_to_latent_model(batch) + del batch + yield latent_data + + def latent_to_logit(self, latent_data: LatentData) -> Union[List[TfMultiBoxTensor], tf.Tensor]: + """ + Decode latent representations into object detection predictions. + + Transforms latent activations back through the decoder portion of the model + to produce bounding box predictions, class scores, and labels. Optionally + applies output formatting to standardize the prediction format. + + Parameters + ---------- + latent_data + Latent activations to decode, wrapped in LatentData container + + Returns + ------- + output + Object detection predictions. If output_formatter is set, returns + list of MultiBoxTensor objects with standardized box format. + Otherwise, returns raw model outputs. + """ + output = self.latent_to_logit_model(latent_data) + if self.output_formatter: + output = self.output_formatter(output) + return output diff --git a/xplique/concepts/tf/layered_model_latent_extractor.py b/xplique/concepts/tf/layered_model_latent_extractor.py new file mode 100644 index 00000000..5e178f37 --- /dev/null +++ b/xplique/concepts/tf/layered_model_latent_extractor.py @@ -0,0 +1,265 @@ +""" +TensorFlow latent data and extractor builder for layered models. +""" + +from typing import Union + +import numpy as np +import tensorflow as tf + +from xplique.utils_functions.classification.tf import TfClassifierTensor + +from ..latent_extractor import LatentData, LatentExtractorBuilder +from .latent_extractor import TfLatentExtractor + + +class LayeredLatentData(LatentData): + """ + Stores latent representations (activations) from a layered TensorFlow model. + + This class encapsulates intermediate activations from any layered model + (ResNet, VGG, DenseNet, etc.) used for classification tasks. It stores + activations from a single intermediate layer of interest. + + Attributes + ---------- + activations + Tensor of intermediate activations from the model. Expected shape depends on + the extraction layer (e.g., (batch, height, width, channels) for conv layers, + or (batch, features) for fully connected layers). + """ + + def __init__(self, activations: tf.Tensor): + """ + Initialize layered model latent data with activations. + + Parameters + ---------- + activations + Intermediate activations tensor from the model. + """ + self.activations = activations + + def __len__(self) -> int: + """ + Return the batch size from the activations. + + Returns + ------- + batch_size + Number of samples in the batch. + """ + return self.activations.shape[0] + + def __getitem__(self, indices: Union[int, slice]) -> "LayeredLatentData": + """ + Get a subset of the latent data by indexing. + + Parameters + ---------- + indices + Indices or slice to extract from the batch. + + Returns + ------- + latent_data + New LayeredLatentData instance with selected samples. + """ + return LayeredLatentData(self.activations[indices]) + + def get_activations( + self, as_numpy: bool = True, keep_gradients: bool = False + ) -> Union[np.ndarray, tf.Tensor]: + """ + Extract activations as a numpy array or tensor. + + Parameters + ---------- + as_numpy + If True, convert tensors to numpy arrays. Default is True. + keep_gradients + If True, preserve gradient information. Default is False. + + Returns + ------- + activations + Activations as numpy array or TensorFlow tensor. + """ + activations = self.activations + + if as_numpy: + activations = ( + activations.numpy() if hasattr(activations, "numpy") else np.array(activations) + ) + + return activations + + def set_activations(self, values: Union[tf.Tensor, np.ndarray]) -> None: + """ + Update activations with new values. + + Parameters + ---------- + values + New activation tensor values as tf.Tensor or numpy array. + """ + if isinstance(values, tf.Tensor): + self.activations = values + else: + # Convert from numpy + self.activations = tf.constant(values) + + +class LayeredModelExtractorBuilder(LatentExtractorBuilder): + """ + Builder for creating LatentExtractor instances for generic layered TensorFlow models. + + This class provides methods to construct a TfLatentExtractor for any layered + model (ResNet, VGG, DenseNet, etc.) by specifying a split layer. It automatically + splits the model's forward pass into feature extraction (g) and classification (h). + """ + + # pylint: disable=arguments-differ + @classmethod + def build( + cls, model: tf.keras.Model, split_layer: int, batch_size: int = 1, **kwargs + ) -> "TfLatentExtractor": + """ + Build a LatentExtractor for a generic layered classifier model. + + This method creates custom g and h functions that split the model's forward pass + at a specified layer: g extracts features up to and including the split layer, + and h processes them through the remaining layers to produce predictions. + + Parameters + ---------- + model + TensorFlow/Keras model instance with sequential layers. + split_layer + Integer index of the layer to split at. Supports negative indexing + (e.g., -1 for the last layer, -2 for the second-to-last). The split + targets the layer at this index, and h processes the remaining layers. + batch_size + Batch size for processing. Default is 1. + **kwargs + Additional keyword arguments (ignored, for compatibility). + + Returns + ------- + latent_extractor + Configured TfLatentExtractor instance for the model. + + Raises + ------ + ValueError + If the model cannot be split at split_layer or uses unsupported input/output shapes. + """ + if not isinstance(model, tf.keras.Model): + raise ValueError("model must be a tf.keras.Model") + if not isinstance(split_layer, int) or isinstance(split_layer, bool): + raise ValueError("split_layer must be an integer layer index") + + layers = list(model.layers) + if not layers: + raise ValueError("model must contain at least one layer") + if not -len(layers) <= split_layer < len(layers): + raise ValueError( + f"split_layer must be between {-len(layers)} and {len(layers) - 1}, " + f"got {split_layer}" + ) + + try: + model_inputs = tf.nest.flatten(model.inputs) + model_outputs = tf.nest.flatten(model.outputs) + except (AttributeError, ValueError) as error: + raise ValueError("model must be built before creating a latent extractor") from error + if len(model_inputs) != 1 or len(model_outputs) != 1: + raise ValueError( + "LayeredModelExtractorBuilder only supports single-input, single-output models" + ) + + split_layer_obj = layers[split_layer] + split_outputs = tf.nest.flatten(split_layer_obj.output) + if len(split_outputs) != 1: + raise ValueError("split_layer must produce exactly one tensor") + split_output = split_outputs[0] + + # Every path to the output must pass through split_output. Keras otherwise + # accepts a bypassed branch and silently reconnects it to the new h input. + pending_tensors = [model_outputs[0]] + visited_tensors = set() + while pending_tensors: + tensor = pending_tensors.pop() + if tensor is split_output: + continue + if id(tensor) in visited_tensors: + continue + visited_tensors.add(id(tensor)) + + history = getattr(tensor, "_keras_history", None) + if history is None: + continue + operation = history.operation + node = operation._inbound_nodes[history.node_index] + parent_tensors = tf.nest.flatten(node.input_tensors) + if not parent_tensors: + raise ValueError( + "split_layer must form a graph cut: model outputs must depend only on " + "the selected layer output" + ) + pending_tensors.extend(parent_tensors) + + try: + # Reuse the Functional graph rather than replaying layers in sequence. This + # preserves branches and merge layers such as residual Add connections. + g_model = tf.keras.Model(inputs=model_inputs[0], outputs=split_output) + h_model = tf.keras.Model(inputs=split_output, outputs=model_outputs[0]) + except ValueError as error: + raise ValueError( + "split_layer must form a graph cut: model outputs must depend only on " + "the selected layer output" + ) from error + + def g(images: tf.Tensor) -> LayeredLatentData: + """ + Extract activations from the split layer (bottleneck features). + + Parameters + ---------- + images + Input images tensor of shape (batch, height, width, 3). + + Returns + ------- + latent_data + LayeredLatentData containing split layer activations. + """ + activations = g_model(images) + return LayeredLatentData(activations) + + def h(latent_data: LayeredLatentData) -> tf.Tensor: + """ + Process latent activations through remaining layers to get logits. + + Parameters + ---------- + latent_data + LayeredLatentData containing split layer activations. + + Returns + ------- + logits + Classification logits tensor of shape (batch, num_classes). + """ + return h_model(latent_data.activations) + + latent_extractor = TfLatentExtractor( + model, + g, + h, + latent_data_class=LayeredLatentData, + output_formatter=TfClassifierTensor.from_predictions, + batch_size=batch_size, + ) + + return latent_extractor diff --git a/xplique/concepts/torch/__init__.py b/xplique/concepts/torch/__init__.py new file mode 100644 index 00000000..fc32121b --- /dev/null +++ b/xplique/concepts/torch/__init__.py @@ -0,0 +1,14 @@ +""" +PyTorch-specific latent extractor implementations +""" + +from .factorizer import TorchSklearnNMFFactorizer +from .holistic_craft import HolisticCraftTorch +from .latent_extractor import TorchLatentData, TorchLatentExtractor + +__all__ = [ + "TorchLatentData", + "TorchLatentExtractor", + "HolisticCraftTorch", + "TorchSklearnNMFFactorizer", +] diff --git a/xplique/concepts/torch/factorizer.py b/xplique/concepts/torch/factorizer.py new file mode 100644 index 00000000..eafce7d5 --- /dev/null +++ b/xplique/concepts/torch/factorizer.py @@ -0,0 +1,274 @@ +""" +PyTorch-specific factorizer implementations +""" + +import numpy as np +import torch + +from ..factorizer import ConceptFactorizer, SklearnNMFFactorizer + + +class TorchSklearnNMFFactorizer(SklearnNMFFactorizer): + """ + PyTorch-compatible sklearn NMF factorizer with differentiable encoding. + """ + + def encode_differentiable(self, activations: torch.Tensor) -> torch.Tensor: + """ + Encode activations with a differentiable non-negative solver. + + Solves the fixed-dictionary NMF subproblem with restartable FISTA so + gradients still flow through ``activations``. + + Parameters + ---------- + activations : torch.Tensor + Activations to encode, shape (n_samples, n_features) + + Returns + ------- + torch.Tensor + Coefficients, shape (n_samples, n_concepts) + """ + if self._concept_bank_w is None: + raise ValueError("Factorizer must be fitted before encoding") + + if not isinstance(activations, torch.Tensor): + activations = torch.as_tensor(activations) + + beta_loss = self.nmf_kwargs.get("beta_loss", "frobenius") + if beta_loss != "frobenius": + raise NotImplementedError( + "PyTorch differentiable NMF encoding only supports beta_loss='frobenius'" + ) + + if torch.any(activations < 0).item(): + raise ValueError("NMF requires non-negative activations") + + concept_bank_tensor = torch.tensor( + self._concept_bank_w, dtype=activations.dtype, device=activations.device + ) + dtype = activations.dtype + device = activations.device + eps = torch.tensor(1e-8, dtype=dtype, device=device) + one = torch.tensor(1.0, dtype=dtype, device=device) + four = torch.tensor(4.0, dtype=dtype, device=device) + + alpha_w = torch.tensor(self.nmf_kwargs.get("alpha_W", 0.0), dtype=dtype, device=device) + l1_ratio = torch.tensor(self.nmf_kwargs.get("l1_ratio", 0.0), dtype=dtype, device=device) + n_features = torch.tensor(activations.shape[1], dtype=dtype, device=device) + l1_reg = n_features * alpha_w * l1_ratio + l2_reg = n_features * alpha_w * (1.0 - l1_ratio) + + max_iter = int(self.nmf_kwargs.get("max_iter", 200)) + tol = float(self.nmf_kwargs.get("tol", 1e-4)) + + gram = concept_bank_tensor @ concept_bank_tensor.T + cross = activations @ concept_bank_tensor.T + diagonal = torch.diagonal(gram) + coefficients = torch.relu(cross / (diagonal.unsqueeze(0) + l2_reg + eps)) + + lipschitz = torch.linalg.norm(gram, ord="fro") + l2_reg + eps + step = one / lipschitz + + def fista_step(coeffs, extrapolated_coeffs, momentum): + gradient = extrapolated_coeffs @ gram - cross + l2_reg * extrapolated_coeffs + updated = torch.relu(extrapolated_coeffs - step * gradient - step * l1_reg) + coeff_delta = torch.linalg.norm(updated - coeffs) / (torch.linalg.norm(updated) + eps) + + next_momentum = (one + torch.sqrt(one + four * torch.square(momentum))) / 2.0 + accelerated = updated + ((momentum - one) / next_momentum) * (updated - coeffs) + restart = torch.sum((extrapolated_coeffs - updated) * (updated - coeffs)) > 0 + + next_extrapolated = torch.where(restart, updated, accelerated) + next_momentum = torch.where(restart, one, next_momentum) + return updated, next_extrapolated, next_momentum, coeff_delta + + extrapolated_coeffs = coefficients + momentum = one + + if tol > 0.0: + for _ in range(max_iter): + coefficients, extrapolated_coeffs, momentum, delta = fista_step( + coefficients, extrapolated_coeffs, momentum + ) + if delta.item() <= tol: + break + else: + for _ in range(max_iter): + coefficients, extrapolated_coeffs, momentum, _ = fista_step( + coefficients, extrapolated_coeffs, momentum + ) + + return coefficients + + def decode(self, coefficients: torch.Tensor) -> torch.Tensor: + """ + Decode coefficients to activations via matrix multiplication. + + Parameters + ---------- + coefficients : torch.Tensor + Coefficients to decode, shape (n_samples, n_concepts) + + Returns + ------- + torch.Tensor + Reconstructed activations, shape (n_samples, n_features) + """ + if self._concept_bank_w is None: + raise ValueError("Factorizer must be fitted before decoding") + + if isinstance(coefficients, np.ndarray): + return coefficients @ self._concept_bank_w + + concept_bank_tensor = torch.tensor( + self._concept_bank_w, dtype=coefficients.dtype, device=coefficients.device + ) + + return coefficients @ concept_bank_tensor + + +class OvercompleteFactorizer(ConceptFactorizer): + """ + Factorizer wrapper for overcomplete optimization methods. + """ + + def __init__(self, optimizer_class, nb_concepts, device="cuda", **kwargs): + """ + Initialize the overcomplete factorizer. + + Parameters + ---------- + optimizer_class : class + The NMF Optimizer class to use (e.g., SemiNMF) + nb_concepts : int + Number of concepts to extract + device : str + Device to use for computation ('cuda' or 'cpu') + **kwargs + Additional arguments passed to the optimizer + """ + self.concept_model = optimizer_class(nb_concepts=nb_concepts, device=device, **kwargs) + self.device = device + + def fit(self, activations: np.ndarray): + """ + Fit the factorizer on activations. + + Parameters + ---------- + activations : np.ndarray + Activations to factorize + + Returns + ------- + tuple + Concept bank and coefficients + """ + activations_torch = torch.tensor(activations, device=self.device) + result = self.concept_model.fit(activations_torch) + # Handle both tuple return (Z, D) and single tensor return + if isinstance(result, tuple): + coeffs_torch, dictionary_torch = result[0], result[1] + else: + coeffs_torch = result + dictionary_torch = self.concept_model.get_dictionary() + + concept_bank_w = dictionary_torch.detach().cpu().numpy() + coeffs_u = coeffs_torch.detach().cpu().numpy() + return concept_bank_w, coeffs_u + + def encode(self, activations: np.ndarray) -> np.ndarray: + """ + Encode activations to coefficients. + + Parameters + ---------- + activations : np.ndarray + Activations to encode + + Returns + ------- + np.ndarray + Coefficients + """ + activations_torch = torch.tensor(activations, device=self.device) + result = self.concept_model.encode(activations_torch) + return result.detach().cpu().numpy() + + def encode_differentiable(self, activations: torch.Tensor) -> torch.Tensor: + """ + Encode activations using differentiable operations. + + Parameters + ---------- + activations : torch.Tensor + Activations to encode + + Returns + ------- + torch.Tensor + Coefficients + """ + return self.concept_model.encode(activations) + + def decode(self, coefficients): + """ + Decode coefficients to activations. + + Parameters + ---------- + coefficients : np.ndarray or torch.Tensor + Coefficients to decode + + Returns + ------- + np.ndarray or torch.Tensor + Reconstructed activations + """ + if isinstance(coefficients, np.ndarray): + coefficients_torch = torch.tensor(coefficients, device=self.device) + result = self.concept_model.decode(coefficients_torch) + return result.detach().cpu().numpy() + + return self.concept_model.decode(coefficients) + + def get_concept_bank(self) -> np.ndarray: + """ + Get the concept bank (dictionary). + + Returns + ------- + np.ndarray + Concept bank + """ + dictionary = self.concept_model.get_dictionary() + return dictionary.detach().cpu().numpy() + + @property + def is_fitted(self) -> bool: + """ + Check if the factorizer has been fitted. + + Returns + ------- + bool + True if fitted, False otherwise + """ + return self.concept_model.fitted + + @property + def requires_positive_activations(self) -> bool: + """ + Check if positive activations are required. + + Returns + ------- + bool + True if positive activations are required + """ + # pylint: disable=import-outside-toplevel + from overcomplete.optimization import SemiNMF + + return not isinstance(self.concept_model, SemiNMF) diff --git a/xplique/concepts/torch/holistic_craft.py b/xplique/concepts/torch/holistic_craft.py new file mode 100644 index 00000000..9c718d1a --- /dev/null +++ b/xplique/concepts/torch/holistic_craft.py @@ -0,0 +1,235 @@ +"""PyTorch-specific wrapper for HolisticCraft.""" + +from typing import Any, Optional, Union + +import numpy as np +import torch +from torch import nn + +from xplique.concepts.factorizer import ConceptFactorizer +from xplique.utils_functions.object_detection.torch.box_model_wrapper import ( + _pad_and_stack_box_predictions, +) +from xplique.wrappers import TorchWrapper + +from ..holistic_craft import ConceptDecoder, HolisticCraft +from ..latent_extractor import LatentData +from .factorizer import TorchSklearnNMFFactorizer +from .latent_extractor import TorchLatentExtractor as LatentExtractor + + +class HolisticCraftTorch(HolisticCraft): + """ + PyTorch-specific implementation of CRAFT for holistic model explanations. + + This class is a thin wrapper around the framework-agnostic base class. + All core functionality is inherited from HolisticCraft. + + Parameters + ---------- + latent_extractor + PyTorch latent extractor for the model + number_of_concepts + Number of concepts to extract (default: 20) + device + PyTorch device. If None, uses the latent extractor device. + factorizer + Optional factorizer instance. If None, creates a TorchSklearnNMFFactorizer + with alpha_W=1e-2 and max_iter=200 + """ + + def __init__( + self, + latent_extractor: LatentExtractor, + number_of_concepts: int = 20, + device: Optional[Union[str, torch.device]] = None, + factorizer: Optional[ConceptFactorizer] = None, + ) -> None: + """ + Initialize the PyTorch CRAFT wrapper. + + Parameters + ---------- + latent_extractor + PyTorch latent extractor for the model + number_of_concepts + Number of concepts to extract (default: 20) + device + PyTorch device. If None, uses the latent extractor device. + factorizer + Optional factorizer instance. If None, creates a TorchSklearnNMFFactorizer + with alpha_W=1e-2 and max_iter=200 + """ + # Create PyTorch-specific factorizer if none provided + if factorizer is None: + factorizer = TorchSklearnNMFFactorizer( + n_components=number_of_concepts, alpha_W=1e-2, max_iter=200 + ) + + if device is None: + device = latent_extractor.device + super().__init__(latent_extractor, number_of_concepts, device, factorizer) + self.framework = "torch" + self._framework_module = torch + + def latent_to_concept_differentiable(self, latent_data: LatentData) -> torch.Tensor: + """ + Transform latent data to concept coefficients with gradient preservation. + + PyTorch-specific implementation using a differentiable non-negative + optimization solver. Ensures gradients flow through the concept + projection for attribution methods. + + Parameters + ---------- + latent_data + Single image's latent representation containing activations + + Returns + ------- + coeffs_u + Concept coefficients as PyTorch tensor with gradients preserved + + Raises + ------ + ValueError + If latent_data is not a single LatentData instance + NotFittedError + If fit() has not been called yet + """ + if not isinstance(latent_data, LatentData): + raise ValueError( + f"latent_to_concept_differentiable() only accepts a single " + f"LatentData as input, got {type(latent_data)}" + ) + self.check_if_fitted() + + # Get activations as tensors with gradients preserved + activations = latent_data.get_activations(as_numpy=False, keep_gradients=True) + + # Ensure we have PyTorch tensors with gradients enabled + if not isinstance(activations, torch.Tensor): + activations = torch.tensor(activations, requires_grad=True) + elif not activations.requires_grad: + activations = activations.clone().detach().requires_grad_(True) + + activations_original_shape = activations.shape[:-1] + activations_flat = activations.reshape(-1, activations.shape[-1]) + + # Use factorizer's differentiable encoding + coeffs_u = self.factorizer.encode_differentiable(activations_flat) + + coeffs_u = coeffs_u.reshape(*activations_original_shape, -1) + return coeffs_u + + def _to_numpy(self, tensor: Any) -> np.ndarray: + """ + Convert PyTorch tensor to numpy array. + + Parameters + ---------- + tensor + PyTorch tensor, TensorFlow tensor, or numpy array + + Returns + ------- + array + Numpy array detached from computational graph + """ + if isinstance(tensor, np.ndarray): + return tensor + if isinstance(tensor, torch.Tensor): + return tensor.detach().cpu().numpy() + if hasattr(tensor, "numpy"): + return tensor.numpy() + return np.asarray(tensor) + + def _to_tensor(self, array: np.ndarray, dtype: Optional[torch.dtype] = None) -> torch.Tensor: + """ + Convert numpy array to PyTorch tensor. + + Parameters + ---------- + array + Numpy array to convert + dtype + Target PyTorch dtype (e.g., torch.float32) + + Returns + ------- + tensor + PyTorch tensor on the specified device + """ + kwargs = {"device": self.device} + if dtype is not None: + kwargs["dtype"] = dtype + return torch.tensor(array, **kwargs) + + def make_concept_decoder(self, latent_data: LatentData) -> TorchWrapper: + """ + Creates a PyTorch concept decoder for gradient-based attribution. + + The decoder is a PyTorch nn.Module that accepts concept coefficients and + returns detection predictions. It maintains a reference to the latent_data + to reconstruct activations during the forward pass. + + Parameters + ---------- + latent_data + Single image's latent representation + + Returns + ------- + decoder + ConceptDecoderTorch instance (nn.Module) with forward method, wrapped in TorchWrapper + """ + + torch_decoder = ConceptDecoderTorch(self, latent_data) + wrapped_decoder = TorchWrapper( + torch_decoder.eval(), device=self.device, is_channel_first=False + ) + return wrapped_decoder + + +class ConceptDecoderTorch(nn.Module, ConceptDecoder): + """ + PyTorch concept decoder module. + + Converts concept coefficients back to object detection predictions by + reconstructing activations and passing them through the decoder network. + + Parameters + ---------- + latent_data + Image-specific latent representation to use for decoding + """ + + def __init__(self, parent_craft: HolisticCraft, latent_data: LatentData) -> None: + super().__init__() + self.parent_craft = parent_craft + self.latent_data = latent_data + + def forward(self, coeffs_u: torch.Tensor) -> torch.Tensor: + """ + Decode concept coefficients to predictions. + + Parameters + ---------- + coeffs_u + Batched concept coefficients. + + Returns + ------- + logits + Predictions as a dense batched tensor. Object detections are zero-padded + to the largest number of boxes in the batch. + """ + + return self._decode(coeffs_u) + + def _predictions_to_tensor(self, predictions) -> torch.Tensor: + if isinstance(predictions, (list, tuple)): + return _pad_and_stack_box_predictions(predictions) + if hasattr(predictions, "to_batched_tensor"): + return predictions.to_batched_tensor() + return predictions diff --git a/xplique/concepts/torch/latent_extractor.py b/xplique/concepts/torch/latent_extractor.py new file mode 100644 index 00000000..ba362fb5 --- /dev/null +++ b/xplique/concepts/torch/latent_extractor.py @@ -0,0 +1,318 @@ +""" +PyTorch-specific latent extractor implementations for object detection models. +""" + +from abc import abstractmethod +from contextlib import nullcontext +from typing import Callable, Generator, List, Optional, Union + +import torch + +from xplique.utils_functions.object_detection.base.box_formatter import ( + BaseBoxFormatter, +) +from xplique.utils_functions.object_detection.torch.box_model_wrapper import ( + _pad_and_stack_box_predictions, +) +from xplique.utils_functions.object_detection.torch.multi_box_tensor import TorchMultiBoxTensor +from xplique.utils_functions.output_as_list_mixin import OutputAsListMixin + +from ..latent_extractor import LatentData, LatentExtractor + + +class TorchLatentData(LatentData): + """ + Base class for PyTorch-based latent representations. + + This abstract class provides a common interface for storing intermediate + activations and positional encodings from PyTorch object detection models. + Subclasses must implement the detach method for gradient management. + + Attributes + ---------- + features + List of feature tensors from the model. + pos + List of positional encoding tensors. + """ + + def __init__(self, features: List, pos: List[torch.Tensor]): + """ + Initialize PyTorch latent data with features and positional encodings. + + Parameters + ---------- + features + List of feature tensors from the model. + pos + List of positional encoding tensors. + """ + self.features = features + self.pos = pos + + @abstractmethod + def detach(self) -> "TorchLatentData": + """ + Detach all tensors from the computation graph. + + This method must be implemented by subclasses to detach features + and positional encodings, preventing gradient computation. + + Returns + ------- + latent_data + Self reference after detaching tensors. + + Raises + ------ + NotImplementedError + If not implemented by subclass. + """ + raise NotImplementedError("detach method must be implemented by subclasses") + + +class TorchLatentExtractor(OutputAsListMixin, LatentExtractor): + """ + PyTorch-specific latent extractor for object detection models. + + This class provides PyTorch-specific implementations for extracting and processing + latent representations from object detection models. It handles device management, + batching, and gradient control for PyTorch tensors. + + Attributes + ---------- + model + PyTorch object detection model. + device + Device for computation ('cuda' or 'cpu'). + training + Training mode flag from the model. + output_as_list + If True, return outputs as list; if False, stack as tensor. + """ + + def __init__( + self, + model: Callable, + input_to_latent_model: Callable, + latent_to_logit_model: Callable, + latent_data_class=LatentData, + output_formatter: Optional[BaseBoxFormatter] = None, + batch_size: int = 8, + device: Optional[Union[str, torch.device]] = None, + ): + """ + Initialize PyTorch latent extractor with model and configuration. + + Parameters + ---------- + model + PyTorch object detection model. + input_to_latent_model + Function (g) that extracts latent representations from inputs. + latent_to_logit_model + Function (h) that processes latent data to predictions. + latent_data_class + Class for storing latent data. Default is LatentData. + output_formatter + Optional formatter for model outputs. Default is None. + batch_size + Batch size for processing. Default is 8. + device + Device for computation. If None, CUDA is used when available and CPU otherwise. + """ + if not isinstance(model, torch.nn.Module): + raise TypeError("model must be a torch.nn.Module") + if not isinstance(batch_size, int) or isinstance(batch_size, bool) or batch_size <= 0: + raise ValueError("batch_size must be a positive integer") + super().__init__( + model, + input_to_latent_model, + latent_to_logit_model, + latent_data_class, + output_formatter, + batch_size, + ) + self.device = self._resolve_device(device) + self.model = self.model.to(self.device) + self.output_as_list = True + + @staticmethod + def _resolve_device(device: Optional[Union[str, torch.device]]) -> torch.device: + """Select the default device and reject unavailable CUDA targets early.""" + if device is None: + device = "cuda" if torch.cuda.is_available() else "cpu" + try: + resolved_device = torch.device(device) + except (TypeError, RuntimeError) as error: + raise ValueError(f"Invalid PyTorch device: {device!r}") from error + + if resolved_device.type == "cuda": + if not torch.cuda.is_available(): + raise ValueError("CUDA was requested but is not available") + if ( + resolved_device.index is not None + and resolved_device.index >= torch.cuda.device_count() + ): + raise ValueError(f"CUDA device index {resolved_device.index} is not available") + return resolved_device + + @staticmethod + def _prepare_inputs(inputs: torch.Tensor) -> torch.Tensor: + """Validate image inputs and add a batch dimension to one image.""" + if not isinstance(inputs, torch.Tensor): + raise TypeError("inputs must be a torch.Tensor") + if inputs.ndim not in (3, 4): + raise ValueError( + "inputs must have rank 3 (a single image) or rank 4 (a batch of images)" + ) + if inputs.numel() == 0: + raise ValueError("inputs must contain at least one value") + if inputs.ndim == 3: + inputs = inputs.unsqueeze(0) + return inputs + + @property + def training(self) -> bool: + """Reflect the underlying model's training mode.""" + return self.model.training + + def eval(self) -> "TorchLatentExtractor": + """ + Set model to evaluation mode. + + Returns + ------- + self + Self reference for method chaining. + """ + self.model.eval() + return self + + def to(self, device: Union[str, torch.device]) -> "TorchLatentExtractor": + """ + Move model to specified device. + + Parameters + ---------- + device + Target device (e.g., 'cuda', 'cpu'). + + Returns + ------- + self + Self reference for method chaining. + """ + self.device = self._resolve_device(device) + self.model.to(self.device) + return self + + def zero_grad(self) -> "TorchLatentExtractor": + """ + Zero out all gradients in the model. + + Returns + ------- + self + Self reference for method chaining. + """ + self.model.zero_grad() + return self + + def forward(self, samples: torch.Tensor) -> Union[List[TorchMultiBoxTensor], torch.Tensor]: + """ + Run full forward pass from inputs to predictions. + + Parameters + ---------- + samples + Input images as PyTorch tensors. + + Returns + ------- + outputs + Model predictions, formatted and optionally stacked based on output_as_list setting. + """ + latent_data = self.input_to_latent(samples) + outputs = self.latent_to_logit_model(latent_data) + if self.output_formatter: + outputs = self.output_formatter(outputs) + if not self.output_as_list: + if isinstance(outputs, (list, tuple)): + outputs = _pad_and_stack_box_predictions(outputs) + elif hasattr(outputs, "to_batched_tensor"): + outputs = outputs.to_batched_tensor() + return outputs + + def input_to_latent(self, inputs: torch.Tensor) -> LatentData: + """ + Extract latent representations from inputs. + + Parameters + ---------- + inputs + Input images as PyTorch tensors (3D or 4D). + + Returns + ------- + latent_data + Latent representations extracted by input_to_latent_model. + """ + inputs = self._prepare_inputs(inputs).to(self.device) + latent_data = self.input_to_latent_model(inputs) + return latent_data + + def input_to_latent_generator( + self, inputs: torch.Tensor, resize: Optional[tuple] = None, keep_gradients: bool = False + ) -> Generator[LatentData, None, None]: + """ + Generator that yields latent data batch by batch. + + Parameters + ---------- + inputs + Input images as PyTorch tensors. + resize + Optional target size for resizing inputs. Default is None. + keep_gradients + If True, preserve gradients during processing. Default is False. + + Yields + ------ + latent_data + LatentData object for each batch, with automatic memory management. + """ + inputs = self._prepare_inputs(inputs) + + for i in range(0, inputs.shape[0], self.batch_size): + i_end = min(i + self.batch_size, inputs.shape[0]) + with nullcontext() if keep_gradients else torch.no_grad(): + batch = inputs[i:i_end].to(self.device) + + if resize: + batch = torch.nn.functional.interpolate( + batch, size=resize, mode="bilinear", align_corners=False + ) + + latent_data = self.input_to_latent_model(batch) + del batch + yield latent_data + + def latent_to_logit(self, latent_data: LatentData) -> List[TorchMultiBoxTensor]: + """ + Process latent data to model predictions. + + Parameters + ---------- + latent_data + Latent representations to process. + + Returns + ------- + output + Model predictions (boxes, scores, labels), optionally formatted. + """ + output = self.latent_to_logit_model(latent_data) + if self.output_formatter: + output = self.output_formatter(output) + return output diff --git a/xplique/concepts/torch/layered_model_latent_extractor.py b/xplique/concepts/torch/layered_model_latent_extractor.py new file mode 100644 index 00000000..f56bee6b --- /dev/null +++ b/xplique/concepts/torch/layered_model_latent_extractor.py @@ -0,0 +1,284 @@ +""" +PyTorch latent data and extractor builder for layered models. +""" + +from typing import Optional, Union + +import numpy as np +import torch + +from xplique.utils_functions.classification.torch import TorchClassifierTensor + +from ..latent_extractor import LatentData, LatentExtractorBuilder +from .latent_extractor import TorchLatentExtractor + + +class LayeredLatentData(LatentData): + """ + Stores latent representations (activations) from a layered PyTorch model. + + This class encapsulates intermediate activations from any layered model + (ResNet, VGG, DenseNet, etc.) used for classification tasks. It stores + activations from a single intermediate layer of interest. + + Attributes + ---------- + activations + Tensor of intermediate activations from the model. Expected shape depends on + the extraction layer (e.g., (batch, channels, height, width) for conv layers, + or (batch, features) for fully connected layers). + """ + + def __init__(self, activations: torch.Tensor): + """ + Initialize layered model latent data with activations. + + Parameters + ---------- + activations + Intermediate activations tensor from the model. + """ + self.activations = activations + + def __len__(self) -> int: + """ + Return the batch size from the activations. + + Returns + ------- + batch_size + Number of samples in the batch. + """ + return self.activations.shape[0] + + def __getitem__(self, indices: Union[int, slice]) -> "LayeredLatentData": + """ + Get a subset of the latent data by indexing. + + Parameters + ---------- + indices + Indices or slice to extract from the batch. + + Returns + ------- + latent_data + New LayeredLatentData instance with selected samples. + """ + return LayeredLatentData(self.activations[indices]) + + def detach(self) -> "LayeredLatentData": + """ + Detach activations tensor from the computation graph. + + This method detaches the activations, preventing gradient computation + through this tensor. + + Returns + ------- + self + Self reference for method chaining. + """ + self.activations = self.activations.detach() + return self + + def get_activations( + self, as_numpy: bool = True, keep_gradients: bool = False + ) -> Union[np.ndarray, torch.Tensor]: + """ + Extract activations as a numpy array or tensor. + + Parameters + ---------- + as_numpy + If True, convert tensors to numpy arrays. Default is True. + keep_gradients + If True, preserve gradient information. Default is False. + + Returns + ------- + activations + Activations as numpy array or PyTorch tensor. + """ + activations = self.activations + + if not keep_gradients: + activations = activations.detach() + + is_4d = len(activations.shape) == 4 + if is_4d: + activations = activations.permute(0, 2, 3, 1) + + if as_numpy: + activations = activations.cpu().numpy() + + return activations + + def set_activations(self, values: torch.Tensor) -> None: + """ + Update activations with new values. + + Parameters + ---------- + values + New activation tensor (torch.Tensor). If 4D, expected format is + (N, H, W, C), which will be converted to PyTorch's (N, C, H, W) format. + """ + is_4d = len(values.shape) == 4 + if is_4d: + values = values.permute(0, 3, 1, 2) + self.activations = values + + def to(self, device: torch.device) -> "LayeredLatentData": + """ + Move all data to the specified device. + + Parameters + ---------- + device + Target device (e.g., torch.device('cuda') or torch.device('cpu')). + + Returns + ------- + latent_data + New LayeredLatentData instance with data on the target device. + """ + return LayeredLatentData(self.activations.to(device)) + + +class LayeredModelExtractorBuilder(LatentExtractorBuilder): + """ + Builder for creating LatentExtractor instances for generic layered PyTorch models. + + This class provides methods to construct a TorchLatentExtractor for any layered + model (ResNet, VGG, DenseNet, etc.) by specifying a split layer. It automatically + splits the model's forward pass into feature extraction (g) and classification (h). + """ + + # pylint: disable=arguments-differ + @classmethod + def build( + cls, + model: torch.nn.Module, + split_layer: int, + device: Optional[Union[str, torch.device]] = None, + batch_size: int = 1, + **kwargs, + ) -> "TorchLatentExtractor": + """ + Build a LatentExtractor for a generic layered classifier model. + + This method creates custom g and h functions that split the model's forward pass + at a specified layer: g extracts features up to and including the split layer, + and h processes them through the remaining layers to produce predictions. + + Parameters + ---------- + model + PyTorch model instance with a sequential structure (via named_modules or + direct access to layers). + split_layer + Integer index of the layer to split at. Supports negative indexing + (e.g., -1 for the last layer, -2 for the second-to-last). The split + targets the layer at this index, and h processes the remaining layers. + device + Device to run computations on. If None, CUDA is used when available and CPU otherwise. + batch_size + Batch size for processing. Default is 1. + **kwargs + Additional keyword arguments (ignored, for compatibility). + + Returns + ------- + latent_extractor + Configured TorchLatentExtractor instance for the model. + + Raises + ------ + ValueError + If split_layer is not found in the model or is an invalid type. + """ + # Get all model children (sequential layers) + children_list = list(model.children()) + if not children_list: + raise ValueError("model must contain at least one child layer") + if not isinstance(split_layer, int) or isinstance(split_layer, bool): + raise ValueError("split_layer must be an integer layer index") + if not -len(children_list) <= split_layer < len(children_list): + raise ValueError( + f"split_layer must be between {-len(children_list)} and " + f"{len(children_list) - 1}, got {split_layer}" + ) + split_index = split_layer % len(children_list) + + def g(images: torch.Tensor) -> LayeredLatentData: + """ + Extract activations from the split layer (bottleneck features). + + Parameters + ---------- + images + Input images tensor of shape (batch, 3, height, width). + + Returns + ------- + latent_data + LayeredLatentData containing split layer activations. + """ + x = images + for layer in children_list[: split_index + 1]: + x = layer(x) + + # Extract activations at split layer + activations = x + + return LayeredLatentData(activations) + + def h(latent_data: LayeredLatentData) -> torch.Tensor: + """ + Process latent activations through remaining layers to get logits. + + Parameters + ---------- + latent_data + LayeredLatentData containing split layer activations. + + Returns + ------- + logits + Classification logits tensor of shape (batch, num_classes). + """ + x = latent_data.activations + + # Process through remaining layers after split + for layer in children_list[split_index + 1 :]: + # Special handling for Sequential containers - check first child + # This handles cases like VGG's classifier which is a Sequential + first_child = None + if isinstance(layer, torch.nn.Sequential): + # Get the first child layer to check if it's Linear + children_of_seq = list(layer.children()) + if children_of_seq: + first_child = children_of_seq[0] + elif isinstance(layer, torch.nn.Linear): + first_child = layer + + # Flatten before FC layers if needed + if first_child is not None and isinstance(first_child, torch.nn.Linear): + if len(x.shape) > 2: + x = torch.flatten(x, 1) + x = layer(x) + + return x + + latent_extractor = TorchLatentExtractor( + model, + g, + h, + latent_data_class=LayeredLatentData, + output_formatter=TorchClassifierTensor.from_predictions, + batch_size=batch_size, + device=device, + ) + + return latent_extractor diff --git a/xplique/features_visualizations/preconditioning.py b/xplique/features_visualizations/preconditioning.py index 3a6a4f46..6495295b 100644 --- a/xplique/features_visualizations/preconditioning.py +++ b/xplique/features_visualizations/preconditioning.py @@ -6,14 +6,14 @@ Credit is due to the original Lucid authors. """ +from pathlib import Path + import numpy as np import tensorflow as tf from ..types import Callable, Optional, Tuple, Union -IMAGENET_SPECTRUM_URL = ( - "https://storage.googleapis.com/serrelab/loupe/spectrums/imagenet_decorrelated.npy" -) +IMAGENET_SPECTRUM_PATH = Path(__file__).with_name("spectrum_decorrelated.npy") def recorrelate_colors(images: tf.Tensor) -> tf.Tensor: @@ -268,10 +268,7 @@ def init_maco_buffer(image_shape, dataset: Optional = None, std=1.0): # init randomly the phase and load the constrained spectrum (average spectrum) phase = np.random.normal(size=(3, *spectrum_shape), scale=std).astype(np.float32) - magnitude_path = tf.keras.utils.get_file( - "spectrum_decorrelated.npy", IMAGENET_SPECTRUM_URL, cache_subdir="spectrums" - ) - magnitude = np.load(magnitude_path) + magnitude = np.load(IMAGENET_SPECTRUM_PATH) magnitude = tf.image.resize(np.moveaxis(magnitude, 0, -1), spectrum_shape).numpy() magnitude = np.moveaxis(magnitude, -1, 0) else: diff --git a/xplique/features_visualizations/spectrum_decorrelated.npy b/xplique/features_visualizations/spectrum_decorrelated.npy new file mode 100644 index 00000000..00f9e418 Binary files /dev/null and b/xplique/features_visualizations/spectrum_decorrelated.npy differ diff --git a/xplique/plots/__init__.py b/xplique/plots/__init__.py index 93c75520..509c828e 100644 --- a/xplique/plots/__init__.py +++ b/xplique/plots/__init__.py @@ -2,7 +2,8 @@ Utility functions to visualize explanations """ -from .image import plot_attribution, plot_attributions, plot_examples, plot_maco +from .image import generate_heatmap, plot_attribution, plot_attributions, plot_examples, plot_maco +from .object_detection import plot_image_detections, plot_images_detections from .tabular import plot_feature_impact, plot_mean_feature_impact, summary_plot_tabular from .timeseries import plot_timeseries_attributions @@ -11,8 +12,11 @@ "plot_attributions", "plot_examples", "plot_maco", + "generate_heatmap", "plot_feature_impact", "plot_mean_feature_impact", "summary_plot_tabular", "plot_timeseries_attributions", + "plot_image_detections", + "plot_images_detections", ] diff --git a/xplique/plots/image.py b/xplique/plots/image.py index 5602d8a8..5a85da21 100644 --- a/xplique/plots/image.py +++ b/xplique/plots/image.py @@ -63,6 +63,7 @@ def _clip_normalize( explanation: Union[tf.Tensor, np.ndarray], clip_percentile: Optional[float] = 0.1, absolute_value: bool = False, + normalize: bool = True, ) -> Union[tf.Tensor, np.ndarray]: if absolute_value: explanation = np.abs(explanation) @@ -70,11 +71,61 @@ def _clip_normalize( if clip_percentile: explanation = _clip_percentile(explanation, clip_percentile) - explanation = _normalize(explanation) + if normalize: + explanation = _normalize(explanation) return explanation +def generate_heatmap( + explanation, + size: tuple, + clip_percentile: Optional[float] = 0.1, + absolute_value: bool = False, + normalize_value: bool = True, +) -> np.ndarray: + """ + Generate a heatmap from the explanation to a specified 2d size. + + Parameters + ---------- + explanation + Attribution / heatmap to plot. + size + Target size of the heatmap (height, width). + clip_percentile + Percentile value to use if clipping is needed, e.g a value of 1 will perform a clipping + between percentile 1 and 99. This parameter allows to avoid outliers in case of too + extreme values. + absolute_value + Whether an absolute value is applied to the explanations. + normalize_value + Whether an normalization is applied to the explanations. + + Returns + ------- + heatmap + The generated heatmap as a numpy array. + """ + if len(explanation.shape) == 4: + raise ValueError( + "Explanation should be 2D or 3D (with channels reduced), " + f"got shape {explanation.shape}." + ) + + heatmap = _clip_normalize(explanation, clip_percentile, absolute_value, normalize_value) + + # resize the explanation to match the image size + if size is not None and size != heatmap.shape[:2]: + if len(heatmap.shape) == 2: + heatmap = tf.image.resize( + heatmap[..., np.newaxis], size, method=tf.image.ResizeMethod.BILINEAR + )[..., 0].numpy() + else: + heatmap = tf.image.resize(heatmap, size, method=tf.image.ResizeMethod.BILINEAR).numpy() + return heatmap + + def plot_attribution( explanation, image: Optional[np.ndarray] = None, @@ -82,6 +133,7 @@ def plot_attribution( alpha: float = 0.5, clip_percentile: Optional[float] = 0.1, absolute_value: bool = False, + normalize_value: bool = True, **plot_kwargs, ): """ @@ -104,17 +156,28 @@ def plot_attribution( extreme values. absolute_value Whether an absolute value is applied to the explanations. + normalize_value + Whether an normalization is applied to the explanations. plot_kwargs Additional parameters passed to `plt.imshow()`. """ if image is not None: image = _normalize(image) - plt.imshow(image) + if image.shape[-1] == 1: + plt.imshow(image[:, :, 0], cmap="Greys") + else: + plt.imshow(image) if len(explanation.shape) == 4: # images channel are reduced explanation = np.mean(explanation, -1) - explanation = _clip_normalize(explanation, clip_percentile, absolute_value) + explanation = generate_heatmap( + explanation, + size=image.shape[:2] if image is not None else None, + clip_percentile=clip_percentile, + absolute_value=absolute_value, + normalize_value=normalize_value, + ) plt.imshow(explanation, cmap=cmap, alpha=alpha, **plot_kwargs) plt.axis("off") @@ -196,16 +259,9 @@ def plot_attributions( for i, explanation in enumerate(explanations): plt.subplot(rows, cols, i + 1) - - if images is not None: - img = _normalize(images[i]) - if img.shape[-1] == 1: - plt.imshow(img[:, :, 0], cmap="Greys") - else: - plt.imshow(img) - plot_attribution( explanation, + image=images[i] if images is not None else None, cmap=cmap, alpha=alpha, clip_percentile=clip_percentile, diff --git a/xplique/plots/metrics.py b/xplique/plots/metrics.py index 3363bd7a..4c5a779e 100644 --- a/xplique/plots/metrics.py +++ b/xplique/plots/metrics.py @@ -66,10 +66,10 @@ def barplot( # either None or string if methods_colors is None: # default cmap - cmap = matplotlib.cm.get_cmap("Set3") + cmap = matplotlib.colormaps["Set3"] else: # methods_color is a string linking to a cmap - cmap = matplotlib.cm.get_cmap(methods_colors) + cmap = matplotlib.colormaps[methods_colors] methods_colors = {methods[i]: cmap((i + 1) / len(methods)) for i in range(len(methods))} @@ -135,10 +135,10 @@ def fidelity_curves( # either None or string if methods_colors is None: # default cmap - cmap = matplotlib.cm.get_cmap("Set3") + cmap = matplotlib.colormaps["Set3"] else: # methods_color is a string linking to a cmap - cmap = matplotlib.cm.get_cmap(methods_colors) + cmap = matplotlib.colormaps[methods_colors] methods_colors = {methods[i]: cmap((i + 1) / len(methods)) for i in range(len(methods))} diff --git a/xplique/plots/object_detection.py b/xplique/plots/object_detection.py new file mode 100644 index 00000000..c2273766 --- /dev/null +++ b/xplique/plots/object_detection.py @@ -0,0 +1,330 @@ +""" +Utilities for displaying images with bounding boxes and optional heatmap overlays. +""" + +import warnings + +import matplotlib.patches as mpatches +import matplotlib.pyplot as plt +import numpy as np +from PIL.Image import Image as PILImage + +from xplique.utils_functions.object_detection.base.box_manager import ( + BaseBoxCoordinatesTranslator, + BoxFormat, + BoxType, + NumpyBoxCoordinatesTranslator, +) +from xplique.utils_functions.object_detection.base.multi_box_tensor import BaseMultiBoxTensor +from xplique.utils_functions.object_detection.tf.box_manager import TfBoxCoordinatesTranslator +from xplique.utils_functions.object_detection.tf.multi_box_tensor import TfMultiBoxTensor + +from ..types import Dict, List, Optional, Tuple, Union + +try: + from xplique.utils_functions.object_detection.torch.box_manager import ( + TorchBoxCoordinatesTranslator, + ) + from xplique.utils_functions.object_detection.torch.multi_box_tensor import TorchMultiBoxTensor +except ImportError: + # If PyTorch is not installed, we can still use the plotting utilities without the + # torch-specific translator and tensor. + TorchMultiBoxTensor = None + TorchBoxCoordinatesTranslator = None + + +_TARGET_BOX_TYPE = BoxType(BoxFormat.XYXY, is_normalized=False) + + +def _get_image_size(image) -> Tuple[int, int]: + """Return (width, height) for both PIL images and NumPy arrays.""" + if isinstance(image, np.ndarray): + return (image.shape[1], image.shape[0]) # (W, H) from (H, W, C) + return image.size # PIL Image already returns (W, H) + + +def _make_translator(multibox_results, box_type) -> BaseBoxCoordinatesTranslator: + """Instantiate the right BoxCoordinatesTranslator based on tensor framework.""" + if TorchMultiBoxTensor is not None and isinstance(multibox_results, TorchMultiBoxTensor): + return TorchBoxCoordinatesTranslator(box_type, _TARGET_BOX_TYPE) + if isinstance(multibox_results, TfMultiBoxTensor): + return TfBoxCoordinatesTranslator(box_type, _TARGET_BOX_TYPE) + return NumpyBoxCoordinatesTranslator(box_type, _TARGET_BOX_TYPE) + + +def _draw_boxes_on_ax( + ax: plt.Axes, + image: Union[np.ndarray, PILImage], + multibox_results: BaseMultiBoxTensor, + classes_labels: List[str], + label_to_color: Dict[str, str], + box_translator: BaseBoxCoordinatesTranslator, + heatmap: Optional[np.ndarray] = None, + cmap: Optional[str] = "viridis", + alpha: Optional[float] = 0.5, + title: Optional[str] = None, + verbose: bool = False, +) -> None: + """Draw a single image with boxes onto an existing Axes using a pre-built translator.""" + class_id_to_label = {i: classes_labels[i] for i in range(len(classes_labels))} + + ax.imshow(image) + + if heatmap is not None: + ax.imshow(heatmap, cmap=cmap, alpha=alpha, extent=ax.images[0].get_extent()) + + boxes = multibox_results.boxes() + scores = multibox_results.scores() + probas = multibox_results.probas() + + image_size = _get_image_size(image) + img_w, img_h = image_size + + _coordinate_upper_scale_limit = 3 + _coordinate_lower_limit = 1.0 + found_labels = set() + for box_coords, score, proba in zip(boxes, scores, probas): + if ( + not box_translator.input_box_type.is_normalized + and max(float(c) for c in box_coords) < _coordinate_lower_limit + ): + warnings.warn( + f"Box coordinates {[float(c) for c in box_coords]} are all below " + f"{_coordinate_lower_limit} but is_normalized=False was declared. " + f"If boxes are in [0, 1] range use is_normalized=True.", + UserWarning, + stacklevel=2, + ) + + translated = box_translator.translate(box_coords[np.newaxis], image_size=image_size) + xmin, ymin, xmax, ymax = box_translator.box_manager.to_numpy_tuple(*translated[0]) + + if box_translator.input_box_type.is_normalized: + if ( + xmax > img_w * _coordinate_upper_scale_limit + or ymax > img_h * _coordinate_upper_scale_limit + ): + raise ValueError( + f"Translated box coordinates ({xmin:.1f}, {ymin:.1f}, {xmax:.1f}, {ymax:.1f}) " + f"far exceed image dimensions ({img_w}x{img_h}). " + f"This usually means box_type.is_normalized does not match the actual boxes. " + f"If boxes are pixel coordinates use is_normalized=False (default); " + f"if boxes are in [0, 1] range use is_normalized=True." + ) + cl = box_translator.box_manager.probas_argmax(proba) + color = label_to_color.get(classes_labels[cl]) + if color is None and verbose: + print( + f"Warning: No color defined for class '{classes_labels[cl]}'. " + f"Using default color 'black'." + ) + name = class_id_to_label.get(cl, "unknown") + found_labels.add(name) + if verbose: + print( + f"cl:{cl}, Drawing box for {name} with color {color} at coords " + f"({xmin}, {ymin}, {xmax}, {ymax}) with score {score:.2f}" + ) + ax.add_patch( + plt.Rectangle( + (xmin, ymin), xmax - xmin, ymax - ymin, fill=False, color=color, linewidth=2 + ) + ) + ax.text( + xmin, + ymin - 15, + f"{score:.2f}", + color=color, + fontsize=10, + bbox={"facecolor": "white", "alpha": 0.0}, + ) + + handles = [ + mpatches.Patch(color=color, label=label) + for label, color in label_to_color.items() + if label in found_labels + ] + ax.legend(handles=handles) + if title is not None: + ax.set_title(title) + + +def plot_image_detections( + image: Union[np.ndarray, PILImage], + multibox_results: BaseMultiBoxTensor, + classes_labels: List[str], + label_to_color: Dict[str, str], + box_type: BoxType = BoxType(BoxFormat.XYXY, is_normalized=False), + heatmap: Optional[np.ndarray] = None, + cmap: Optional[str] = "viridis", + alpha: Optional[float] = 0.5, + title: Optional[str] = None, + ax: Optional[plt.Axes] = None, + verbose: bool = False, +): + """ + Display an image with bounding boxes, optionally with a heatmap overlay. + + Parameters + ---------- + image + Input image to display. Accepts either a PIL ``Image`` or a NumPy array + of shape (H, W, C). + multibox_results + Bounding box annotations. + classes_labels + List of class labels. + label_to_color + Dictionary mapping labels to colors. + box_type + Box type describing the coordinate format and normalization of the boxes + stored in multibox_results. Defaults to XYXY non-normalized (pixel coordinates). + heatmap + Optional explanation heatmap (2D array) to overlay on the image. + cmap + Optional Matplotlib colormap for the explanation heatmap. + alpha + Optional Alpha transparency for the explanation heatmap overlay. + title + Optional title for the plot. + ax + Optional Matplotlib Axes object to plot on. If None, a new figure and axes will be created. + verbose + Whether to print debug information. + + Returns + ------- + fig + Matplotlib figure object. + """ + if ax is None: + fig, ax = plt.subplots() + else: + fig = ax.figure + + box_translator = _make_translator(multibox_results, box_type) + _draw_boxes_on_ax( + ax, + image, + multibox_results, + classes_labels, + label_to_color, + box_translator, + heatmap, + cmap, + alpha, + title, + verbose, + ) + return fig + + +def plot_images_detections( + images: List[Union[np.ndarray, PILImage]], + multibox_results_list: List[BaseMultiBoxTensor], + classes_labels: List[str], + label_to_color: Dict[str, str], + box_type: BoxType = BoxType(BoxFormat.XYXY, is_normalized=False), + heatmaps: Optional[List[np.ndarray]] = None, + cmap: Optional[str] = "viridis", + alpha: Optional[float] = 0.5, + titles: Optional[List[str]] = None, + num_cols: int = 5, + verbose: bool = False, +): + """ + Display multiple images with bounding boxes, optionally with heatmap overlays. + + Parameters + ---------- + images + List of images to display. + Each image can be a PIL ``Image`` or a NumPy array of shape (H, W, C). + multibox_results_list + List of bounding box annotations for each image. + classes_labels + List of class labels. + label_to_color + Dictionary mapping labels to colors. + box_type + Box type describing the coordinate format and normalization of the boxes + stored in multibox_results. Defaults to XYXY non-normalized (pixel coordinates). + heatmaps + Optional list of explanation heatmaps (2D arrays) to overlay on the images. + cmap + Optional Matplotlib colormap for the explanation heatmaps. + alpha + Optional Alpha transparency for the explanation heatmap overlays. + titles + Optional list of titles for each subplot. + num_cols + Number of columns in the subplot grid. Default is 5. + verbose + Whether to print debug information. + + Returns + ------- + fig + Matplotlib figure object. + + Raises + ------ + ValueError + If inputs are empty, or if the number of images and multibox_results are not equal. + """ + # Validate inputs + if len(images) == 0 or len(multibox_results_list) == 0: + raise ValueError("images and multibox_results_list must contain at least one element") + + if len(images) != len(multibox_results_list): + raise ValueError( + f"Number of images ({len(images)}) must match number of multibox_results " + f"({len(multibox_results_list)})" + ) + + if heatmaps is not None and len(heatmaps) != len(images): + raise ValueError( + f"Number of heatmaps ({len(heatmaps)}) must match number of images ({len(images)})" + ) + + if titles is not None and len(titles) != len(images): + raise ValueError( + f"Number of titles ({len(titles)}) must match number of images ({len(images)})" + ) + + num_images = len(images) + + # Build translator once for the first element (framework is the same for all) + box_translator = _make_translator(multibox_results_list[0], box_type) + + num_cols = min(num_cols, num_images) + num_rows = int(np.ceil(num_images / num_cols)) + + fig, axes = plt.subplots(num_rows, num_cols, figsize=(5 * num_cols, 5 * num_rows)) + + # Handle case where we have only one subplot + if num_images == 1: + axes = np.array([axes]) + axes = axes.flatten() + + for idx in range(num_images): + _draw_boxes_on_ax( + axes[idx], + images[idx], + multibox_results_list[idx], + classes_labels, + label_to_color, + box_translator, + heatmaps[idx] if heatmaps is not None else None, + cmap, + alpha, + titles[idx] if titles is not None else None, + verbose, + ) + + # Hide unused subplots + for idx in range(num_images, len(axes)): + axes[idx].axis("off") + + plt.tight_layout() + return fig diff --git a/xplique/utils_functions/__init__.py b/xplique/utils_functions/__init__.py index a1dba97f..ff71acf5 100644 --- a/xplique/utils_functions/__init__.py +++ b/xplique/utils_functions/__init__.py @@ -2,6 +2,7 @@ Functions to ease attributions """ +from .object_detection import object_detection_operators from .segmentation import get_class_zone, get_common_border, get_connected_zone, get_in_out_border __all__ = [ @@ -9,4 +10,5 @@ "get_connected_zone", "get_common_border", "get_in_out_border", + "object_detection_operators", ] diff --git a/xplique/utils_functions/classification/__init__.py b/xplique/utils_functions/classification/__init__.py new file mode 100644 index 00000000..128300f6 --- /dev/null +++ b/xplique/utils_functions/classification/__init__.py @@ -0,0 +1,3 @@ +""" +Classification utilities for Xplique. +""" diff --git a/xplique/utils_functions/classification/base/__init__.py b/xplique/utils_functions/classification/base/__init__.py new file mode 100644 index 00000000..e6750b6d --- /dev/null +++ b/xplique/utils_functions/classification/base/__init__.py @@ -0,0 +1,5 @@ +""" +Base classification utilities. +""" + +__all__ = [] diff --git a/xplique/utils_functions/classification/tf/__init__.py b/xplique/utils_functions/classification/tf/__init__.py new file mode 100644 index 00000000..eddc200e --- /dev/null +++ b/xplique/utils_functions/classification/tf/__init__.py @@ -0,0 +1,7 @@ +""" +TensorFlow classification utilities. +""" + +from .classifier_tensor import TfClassifierTensor + +__all__ = ["TfClassifierTensor"] diff --git a/xplique/utils_functions/classification/tf/classifier_tensor.py b/xplique/utils_functions/classification/tf/classifier_tensor.py new file mode 100644 index 00000000..476a926a --- /dev/null +++ b/xplique/utils_functions/classification/tf/classifier_tensor.py @@ -0,0 +1,178 @@ +""" +TensorFlow implementation of ClassifierTensor for classification predictions. + +This module provides a TensorFlow wrapper for classification predictions +with a unified format compatible with the StructuredPrediction protocol. +""" + +import warnings +from numbers import Integral + +import tensorflow as tf + +from xplique.commons.prediction_types import StructuredPrediction + + +class TfClassifierTensor(StructuredPrediction): + """ + TensorFlow wrapper for classification predictions. + + This class wraps TensorFlow tensors from classification models to provide + the same interface as MultiBoxTensor, allowing polymorphic handling of + both object detection and classification predictions. + + Parameters + ---------- + tensor + TensorFlow tensor containing classifier predictions (logits or probabilities) + """ + + def __init__(self, tensor: tf.Tensor): + tensor = tf.convert_to_tensor(tensor) + if tensor.shape.rank not in (1, 2): + raise ValueError("Classifier predictions must have rank 1 or 2.") + self.tensor = tensor + + @classmethod + def from_predictions(cls, predictions): + """Wrap raw classifier predictions unless they are already formatted. + + Raises + ------ + ValueError + If predictions rank is not 1 or 2. + """ + if isinstance(predictions, cls): + return predictions + return cls(predictions) + + @property + def shape(self): + """Return the shape of the underlying tensor for compatibility.""" + return self.tensor.shape + + @staticmethod + def _dim_as_int(tensor: tf.Tensor, axis: int) -> int: + """Return static dimension when available, else evaluate dynamic size eagerly.""" + dim = tensor.shape[axis] + if dim is not None: + return int(dim) + return int(tf.shape(tensor)[axis]) + + @property + def num_classes(self) -> int: + """Number of classes in the prediction tensor.""" + return self._dim_as_int(self.tensor, -1) + + @property + def batch_size(self) -> int: + """Batch size of predictions (1 for rank-1 single predictions).""" + if self.tensor.shape.rank == 1: + return 1 + return self._dim_as_int(self.tensor, 0) + + @property + def is_empty(self) -> bool: + """Whether there are no predictions to explain.""" + return self.batch_size == 0 or self.num_classes == 0 + + def __len__(self): + """ + Deprecated length of the classifier tensor. + + Use ``num_classes``, ``batch_size``, or ``is_empty`` instead. + + Returns + ------- + length + Number of classes in the classification output. + """ + warnings.warn( + "len(TfClassifierTensor) is ambiguous and deprecated; use " + "num_classes, batch_size, or is_empty instead.", + DeprecationWarning, + stacklevel=2, + ) + return self.num_classes + + def __tf_tensor__(self, dtype=None, name=None): + """ + Convert to TensorFlow tensor for use in TF operations. + + This method enables ClassifierTensor to be used directly in TensorFlow + operations like tf.stack(), tf.expand_dims(), etc. without explicitly + accessing the .tensor attribute. + + Parameters + ---------- + dtype + Optional dtype to convert to. + name + Optional name for the operation. + + Returns + ------- + tensor + The underlying TensorFlow tensor. + """ + return tf.convert_to_tensor(self.tensor, dtype=dtype, name=name) + + def to_batched_tensor(self) -> tf.Tensor: + """ + Ensure tensor has batch dimension. + + For classifiers, if the tensor is 1D (single prediction), adds a batch + dimension. If already 2D or higher, returns as-is. + + Returns + ------- + batched_tensor + Tensor with batch dimension: (1, num_classes) or (batch, num_classes) + """ + if self.tensor.shape.rank == 1: + return tf.convert_to_tensor(tf.expand_dims(self.tensor, axis=0)) + return tf.convert_to_tensor(self.tensor) + + def filter(self, class_id=None, confidence=None): + """ + No-op for classifiers. + + Classifiers do not have multiple detections to filter. Returns self + unchanged for interface compatibility with object detection types. + """ + return self + + def to_attribution_target(self, class_id=None): + """ + Build a one-hot attribution target for a selected class. + + The actual logit/probability values are not used by the attribution + method — only the class index and output shape matter. ``class_id`` + selects which output neuron to differentiate through. + + Parameters + ---------- + class_id + Class to target. If None, returns self (raw model output as target). + confidence + Ignored for classifiers. + + Returns + ------- + target + One-hot tensor for ``class_id``, or self when class_id is None. + """ + if class_id is None: + return self + + if not isinstance(class_id, Integral) or isinstance(class_id, bool): + raise ValueError("class_id must be an integer.") + + class_id = int(class_id) + num_classes = self.tensor.shape[-1] + if class_id < 0 or (num_classes is not None and class_id >= num_classes): + raise ValueError(f"class_id must be in [0, {num_classes}).") + + target = tf.one_hot(class_id, tf.shape(self.tensor)[-1], dtype=self.tensor.dtype) + target = tf.broadcast_to(target, tf.shape(self.tensor)) + return TfClassifierTensor(target) diff --git a/xplique/utils_functions/classification/torch/__init__.py b/xplique/utils_functions/classification/torch/__init__.py new file mode 100644 index 00000000..3d763233 --- /dev/null +++ b/xplique/utils_functions/classification/torch/__init__.py @@ -0,0 +1,7 @@ +""" +PyTorch classification utilities. +""" + +from .classifier_tensor import TorchClassifierTensor + +__all__ = ["TorchClassifierTensor"] diff --git a/xplique/utils_functions/classification/torch/classifier_tensor.py b/xplique/utils_functions/classification/torch/classifier_tensor.py new file mode 100644 index 00000000..62444984 --- /dev/null +++ b/xplique/utils_functions/classification/torch/classifier_tensor.py @@ -0,0 +1,156 @@ +""" +PyTorch implementation of ClassifierTensor for classification predictions. + +This module provides a PyTorch tensor subclass for classification predictions +with a unified format. Due to metaclass conflicts with torch.Tensor, this class +cannot explicitly inherit from the StructuredPrediction protocol but implements its +interface via structural typing (duck typing). +""" + +import warnings +from numbers import Integral + +import torch + +from xplique.commons.prediction_types import StructuredPrediction + + +class TorchClassifierTensor(torch.Tensor): + """ + Tensor representation for classification predictions. + + This class extends torch.Tensor to represent classification model outputs + (logits or probabilities) with a shape of (num_classes,) for single predictions + or (batch_size, num_classes) for batched predictions. + + Note: This class implements the StructuredPrediction protocol (see + xplique.commons.prediction_types.StructuredPrediction) via structural typing. + The class complies with the protocol by implementing: + - to_batched_tensor(): Adds batch dimension if needed + - filter(class_id, confidence): Creates a one-hot target for a selected class + """ + + @classmethod + def __torch_function__(cls, func, types, args=(), kwargs=None): + """Delegate torch operations while preserving Tensor subclass semantics. + + PyTorch does not expose an equivalent of TensorFlow's ``__tf_tensor__`` + conversion hook, so this class intentionally subclasses ``torch.Tensor`` + to keep autograd and tensor operations available on formatted outputs. + Delegating here keeps the default PyTorch subclass behavior explicit. + """ + if kwargs is None: + kwargs = {} + return super().__torch_function__(func, types, args, kwargs) + + @classmethod + def from_predictions(cls, predictions): + """Wrap raw classifier predictions unless they are already formatted. + + Raises + ------ + ValueError + If predictions rank is not 1 or 2. + """ + if isinstance(predictions, cls): + return predictions + tensor = torch.as_tensor(predictions) + if tensor.ndim not in (1, 2): + raise ValueError("Classifier predictions must have rank 1 or 2.") + return tensor.as_subclass(cls) + + @property + def num_classes(self) -> int: + """Number of classes in the prediction tensor.""" + return int(self.shape[-1]) + + @property + def batch_size(self) -> int: + """Batch size of predictions (1 for rank-1 single predictions).""" + if self.ndim == 1: + return 1 + return int(self.shape[0]) + + @property + def is_empty(self) -> bool: + """Whether there are no predictions to explain.""" + return self.batch_size == 0 or self.num_classes == 0 + + def __len__(self) -> int: + """Deprecated length accessor for classifier tensors.""" + warnings.warn( + "len(TorchClassifierTensor) is ambiguous and deprecated; use " + "num_classes, batch_size, or is_empty instead.", + DeprecationWarning, + stacklevel=2, + ) + return super().__len__() + + def to_batched_tensor(self) -> torch.Tensor: + """ + Ensure tensor has batch dimension. + + For classifiers, if the tensor is 1D (single prediction), adds a batch + dimension. If already 2D or higher, returns as-is. + + Returns + ------- + batched_tensor + Tensor with batch dimension: (1, num_classes) or (batch, num_classes) + """ + if len(self.shape) == 1: + return torch.unsqueeze(self, 0) + return self + + def filter(self, class_id=None, confidence=None): + """ + No-op for classifiers. + + Classifiers do not have multiple detections to filter. Returns self + unchanged for interface compatibility with object detection types. + """ + return self + + def to_attribution_target(self, class_id=None): + """ + Build a one-hot attribution target for a selected class. + + The actual logit/probability values are not used by the attribution + method — only the class index and output shape matter. ``class_id`` + selects which output neuron to differentiate through. + + Parameters + ---------- + class_id + Class to target. If None, returns self (raw model output as target). + + Returns + ------- + target + One-hot tensor for ``class_id``, or self when class_id is None. + """ + if class_id is None: + return self + + if not isinstance(class_id, Integral) or isinstance(class_id, bool): + raise ValueError("class_id must be an integer.") + + class_id = int(class_id) + num_classes = self.shape[-1] + if class_id < 0 or class_id >= num_classes: + raise ValueError(f"class_id must be in [0, {num_classes}).") + + target = torch.zeros_like(self) + target[..., class_id] = 1 + # Preserve subclass identity regardless of upstream torch.zeros_like semantics. + if not isinstance(target, type(self)): + target = target.as_subclass(type(self)) + return target + + +# Verify structural compliance with StructuredPrediction protocol at import time. +# TorchClassifierTensor cannot explicitly inherit from StructuredPrediction due to a +# metaclass conflict between torch.Tensor (_TensorMeta) and Protocol (_ProtocolMeta). +assert issubclass(TorchClassifierTensor, StructuredPrediction), ( + "TorchClassifierTensor must structurally satisfy the StructuredPrediction protocol" +) diff --git a/xplique/utils_functions/common/__init__.py b/xplique/utils_functions/common/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/xplique/utils_functions/common/tf/__init__.py b/xplique/utils_functions/common/tf/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/xplique/utils_functions/common/tf/gradients_check.py b/xplique/utils_functions/common/tf/gradients_check.py new file mode 100644 index 00000000..9782f67d --- /dev/null +++ b/xplique/utils_functions/common/tf/gradients_check.py @@ -0,0 +1,151 @@ +""" +TensorFlow gradient checking utilities for object detection models. +""" + +from typing import Any, List + +import tensorflow as tf + + +def _extract_all_tensors(obj: Any) -> List[tf.Tensor]: + """ + Recursively extract all TensorFlow tensors from a nested structure. + + Handles nested combinations of dicts, lists, tuples, and tensors. + Special handling for MultiBoxTensor objects. + + Parameters + ---------- + obj + Object to extract tensors from (can be tensor, dict, list, tuple, or nested combinations). + + Returns + ------- + tensors + List of all tensors found in the structure. + """ + if isinstance(obj, tf.Tensor): + return [obj] + if hasattr(obj, "tensor"): + # Handle MultiBoxTensor or similar objects with a .tensor attribute + return [obj.tensor] + if isinstance(obj, dict): + tensors = [] + for value in obj.values(): + tensors.extend(_extract_all_tensors(value)) + return tensors + if isinstance(obj, (list, tuple)): + tensors = [] + for item in obj: + tensors.extend(_extract_all_tensors(item)) + return tensors + # Not a tensor or container, return empty list + return [] + + +def _vjp_probes(tensor: tf.Tensor) -> List[tf.Tensor]: + """Return deterministic probes that do not reduce every output to its sum.""" + alternating = tf.range(tf.size(tensor)) + alternating = tf.cast(2 * tf.math.floormod(alternating, 2) - 1, tensor.dtype) + return [tf.ones_like(tensor), tf.reshape(alternating, tf.shape(tensor))] + + +def check_model_gradients(func: Any, input_tensor: tf.Tensor, verbose: bool = False) -> bool: + """ + Test gradients in both Eager and Graph modes for an object detection model. + + This function validates that gradients can be computed through the model in both + TensorFlow execution modes (eager and graph). It handles various output formats + including dictionaries, lists, MultiBoxTensor objects, and raw tensors using + recursive tensor extraction. + + Parameters + ---------- + func + Callable model or function to test. Should accept input_tensor and return + predictions in dict, list, or tensor format. + input_tensor + Input tensor to use for gradient computation testing. + verbose + If True, print information about gradient computation. Default is False. + + Returns + ------- + success + True if gradients can be computed successfully in at least one mode (eager or graph), + False otherwise. + """ + + def _test_gradients_single_mode(mode_name: str, func_to_call): + """ + Test gradients using the provided callable. + + Parameters + ---------- + mode_name + Name of the mode for logging ("Eager" or "Graph"). + func_to_call + The callable to invoke — either the original function (eager) + or a tf.function-wrapped version (graph). + + Returns + ------- + success + True if at least one gradient computation succeeded, False otherwise. + """ + if verbose: + print(f"\n--- Testing {mode_name} mode ---") + + try: + with tf.GradientTape(persistent=True) as tape: + tape.watch(input_tensor) + predictions = func_to_call(input_tensor) + + # Extract all tensors recursively from the output structure + tensors = _extract_all_tensors(predictions) + + if not tensors: + if verbose: + print("No tensors found in outputs") + return False + + # Probe each output independently: summing outputs can cancel gradients, + # for example when a model returns probabilities that sum to one. + for tensor in tensors: + if not (tensor.dtype.is_floating or tensor.dtype.is_complex): + continue + for probe in _vjp_probes(tensor): + gradients = tape.gradient(tensor, input_tensor, output_gradients=probe) + if gradients is None: + continue + is_finite = bool(tf.reduce_all(tf.math.is_finite(gradients)).numpy()) + is_nonzero = bool(tf.reduce_any(tf.not_equal(gradients, 0)).numpy()) + if is_finite and is_nonzero: + if verbose: + grad_sum = tf.reduce_sum(tf.abs(gradients)) + print(f"Gradients OK - sum={grad_sum.numpy():.6f}") + return True + + if verbose: + print("No finite, non-zero gradients") + return False + + # pylint: disable=broad-exception-caught + except Exception as e: + if verbose: + print(f"{mode_name} mode failed completely: {e}") + return False + + # Test both modes — no global state is mutated. + # Eager: call func directly (TF default behaviour). + # Graph: tf.function forces graph compilation regardless of any global flag. + eager_result = _test_gradients_single_mode("Eager", func) + graph_result = _test_gradients_single_mode("Graph", tf.function(func)) + + # Summary + if verbose: + print("\n=== SUMMARY ===") + print(f"Eager mode: {'OK' if eager_result else 'FAIL'}") + print(f"Graph mode: {'OK' if graph_result else 'FAIL'}") + + return eager_result or graph_result diff --git a/xplique/utils_functions/common/torch/__init__.py b/xplique/utils_functions/common/torch/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/xplique/utils_functions/common/torch/gradients_check.py b/xplique/utils_functions/common/torch/gradients_check.py new file mode 100644 index 00000000..2999c2c1 --- /dev/null +++ b/xplique/utils_functions/common/torch/gradients_check.py @@ -0,0 +1,154 @@ +""" +Utilities for checking gradient propagation in PyTorch models. +""" + +# pylint: disable=duplicate-code +from typing import Any, Callable, List, Union + +import torch + + +def _extract_all_tensors(obj: Any) -> List[torch.Tensor]: + """ + Recursively extract all PyTorch tensors from a nested structure. + + Handles nested combinations of dicts, lists, tuples, and tensors. + + Parameters + ---------- + obj + Object to extract tensors from (can be tensor, dict, list, tuple, or nested combinations). + + Returns + ------- + tensors + List of all tensors found in the structure. + """ + if isinstance(obj, torch.Tensor): + return [obj] + if isinstance(obj, dict): + tensors = [] + for value in obj.values(): + tensors.extend(_extract_all_tensors(value)) + return tensors + if isinstance(obj, (list, tuple)): + tensors = [] + for item in obj: + tensors.extend(_extract_all_tensors(item)) + return tensors + # Not a tensor or container, return empty list + return [] + + +def _vjp_probes(tensor: torch.Tensor) -> List[torch.Tensor]: + """Return deterministic probes that do not reduce every output to its sum.""" + alternating = torch.arange(tensor.numel(), device=tensor.device) + alternating = 2 * torch.remainder(alternating, 2) - 1 + return [torch.ones_like(tensor), alternating.to(tensor.dtype).reshape_as(tensor)] + + +def check_model_gradients( + func: Union[Callable, torch.nn.Module], input_tensor: torch.Tensor, verbose: bool = False +) -> bool: + """ + Checks if gradients are propagated to the inputs of a PyTorch model. + + Parameters + ---------- + func + A PyTorch model (nn.Module) or a callable function. + input_tensor + The input tensor. + verbose + If True, print information about gradient computation. Default is False. + + Returns + ------- + bool + True if non-zero gradients are propagated to the input, False otherwise. + """ + if not isinstance(input_tensor, torch.Tensor): + raise TypeError("input_tensor must be a torch.Tensor") + if not (input_tensor.is_floating_point() or input_tensor.is_complex()): + raise TypeError("input_tensor must have a floating-point or complex dtype") + + module_states = [] + in_place_warning_printed = False + if isinstance(func, torch.nn.Module): + module_states = [(module, module.training) for module in func.modules()] + + try: + # This context restores the caller's grad-mode state on every exit path. + with torch.enable_grad(): + device = None + if isinstance(func, torch.nn.Module): + func.eval() + + # Print a warning about in-place operations in ReLU, etc... + for module in func.modules(): + if hasattr(module, "inplace") and module.inplace: + if not in_place_warning_printed and verbose: + print( + f"Warning: In-place operation found in {type(module)}. " + f"This may cause issues with gradient computation." + ) + in_place_warning_printed = True + + try: + device = next(func.parameters()).device + except StopIteration: + try: + device = next(func.buffers()).device + except StopIteration: + pass + + # Transfer before enabling gradients on the input so it remains a leaf. + x = input_tensor.detach() + if device is not None: + x = x.to(device) + x = x.clone().requires_grad_(True) + outputs = func(x) + tensors = _extract_all_tensors(outputs) + + if not tensors: + if verbose: + print("No tensor found in outputs") + return False + + # Probe each output independently: summing outputs can cancel gradients, + # for example when a model returns probabilities that sum to one. + for tensor in tensors: + if ( + not (tensor.is_floating_point() or tensor.is_complex()) + or not tensor.requires_grad + ): + continue + for probe in _vjp_probes(tensor): + (gradients,) = torch.autograd.grad( + tensor, + x, + grad_outputs=probe, + retain_graph=True, + allow_unused=True, + ) + if gradients is None: + continue + is_finite = bool(torch.all(torch.isfinite(gradients)).item()) + is_nonzero = bool(torch.any(gradients != 0).item()) + if is_finite and is_nonzero: + if verbose: + print(f"Gradients OK - sum={gradients.abs().sum().item():.6f}") + return True + + if verbose: + print("No finite, non-zero gradients") + return False + + # pylint: disable=broad-exception-caught + except Exception as e: + if verbose: + print(f"Error: {str(e)}") + return False + finally: + for module, training in module_states: + module.training = training diff --git a/xplique/utils_functions/object_detection/__init__.py b/xplique/utils_functions/object_detection/__init__.py new file mode 100644 index 00000000..f3c84730 --- /dev/null +++ b/xplique/utils_functions/object_detection/__init__.py @@ -0,0 +1,11 @@ +""" +Object detection utilities +""" + +from .object_detection_operators import _EPSILON, _box_iou, _format_objects + +__all__ = [ + "_box_iou", + "_format_objects", + "_EPSILON", +] diff --git a/xplique/utils_functions/object_detection/base/__init__.py b/xplique/utils_functions/object_detection/base/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/xplique/utils_functions/object_detection/base/box_formatter.py b/xplique/utils_functions/object_detection/base/box_formatter.py new file mode 100644 index 00000000..f3c4e724 --- /dev/null +++ b/xplique/utils_functions/object_detection/base/box_formatter.py @@ -0,0 +1,89 @@ +""" +Base classes for formatting object detection predictions. +""" + +from abc import ABC, abstractmethod +from typing import Any + +from .box_manager import BoxFormat, BoxType +from .multi_box_tensor import BaseMultiBoxTensor + + +class BaseBoxFormatter(ABC): + """ + Abstract base class for formatting object detection predictions. + + This class provides a common interface for converting model predictions into + a standardized format with configurable box coordinate representations. + + Parameters + ---------- + input_box_type + The box type (format and normalization) of the input predictions. + output_box_type + The desired box type (format and normalization) for the output. + Default is XYXY format with normalized coordinates. + """ + + def __init__( + self, + input_box_type: BoxType, + output_box_type: BoxType = BoxType(BoxFormat.XYXY, is_normalized=True), + ) -> None: + super().__init__() + self.input_box_type = input_box_type + self.output_box_type = output_box_type + + def __call__(self, predictions: Any) -> BaseMultiBoxTensor: + """ + Callable interface for formatting predictions. + + Parameters + ---------- + predictions + Raw model predictions to format. + + Returns + ------- + formatted_predictions + Formatted predictions in the standardized format. + """ + return self.forward(predictions) + + @abstractmethod + def forward(self, predictions: Any) -> BaseMultiBoxTensor: + """ + Forward pass to format predictions. + + This method must be implemented by framework-specific subclasses. + + Parameters + ---------- + predictions + Raw model predictions to format. + + Returns + ------- + formatted_predictions + Formatted predictions in the standardized format. + """ + raise NotImplementedError("This method should be implemented in the subclass") + + def format_predictions(self, predictions: Any) -> BaseMultiBoxTensor: + """ + Format predictions into the standardized MultiBoxTensor format. + + Subclasses may override this helper when they expose a separate + formatting step in addition to ``forward``. + + Parameters + ---------- + predictions + Dictionary or tensor containing boxes, scores, and class probabilities. + + Returns + ------- + MultiBoxTensor + Formatted predictions with boxes, scores, and class probabilities. + """ + raise NotImplementedError("format_predictions is optional and not implemented") diff --git a/xplique/utils_functions/object_detection/base/box_manager.py b/xplique/utils_functions/object_detection/base/box_manager.py new file mode 100644 index 00000000..f64ef0df --- /dev/null +++ b/xplique/utils_functions/object_detection/base/box_manager.py @@ -0,0 +1,403 @@ +""" +Base classes for managing bounding box operations across different frameworks. +""" + +from abc import ABC, abstractmethod +from dataclasses import dataclass +from enum import Enum +from typing import Tuple + +import numpy as np + + +class BoxFormat(Enum): + """ + Enumeration of supported bounding box coordinate formats. + + Attributes + ---------- + CXCYWH + Center x, center y, width, height format. + XYWH + Top-left x, top-left y, width, height format. + XYXY + Top-left x, top-left y, bottom-right x, bottom-right y format. + """ + + CXCYWH = "CXCYWH" + XYWH = "XYWH" + XYXY = "XYXY" + + +@dataclass +class BoxType: + """ + Data class representing the type and coordinate system of bounding boxes. + + Attributes + ---------- + format + The coordinate format of the boxes (CXCYWH, XYWH, or XYXY). + is_normalized + Whether coordinates are normalized to [0, 1] range or in pixel units. + """ + + format: BoxFormat + is_normalized: bool + + +class BoxManager(ABC): + """ + Abstract base class for managing bounding box operations. + + This class defines the interface for box management operations across different + frameworks (NumPy, TensorFlow, PyTorch). Subclasses should implement framework-specific + box transformations and conversions. + """ + + @staticmethod + @abstractmethod + def to_numpy_tuple(*arrays) -> Tuple: + """Convert framework tensors/arrays to a tuple of NumPy arrays.""" + + @staticmethod + @abstractmethod + def probas_argmax(proba) -> int: + """Return the class ID with the highest probability as a Python int.""" + + +class BaseBoxCoordinatesTranslator(ABC): + """ + Abstract base class for box coordinates translators. + + Defines the common interface for all framework-specific translators. + Each subclass must expose a ``box_manager`` instance providing + framework-specific utility operations (e.g. tensor-to-numpy conversion). + """ + + @property + @abstractmethod + def box_manager(self) -> "BoxManager": + """Framework-specific BoxManager instance for utility operations.""" + + @abstractmethod + def translate(self, box, image_size=None): + """Translate boxes from input format/scale to output format/scale.""" + + +class NumpyBoxManager(BoxManager): + """ + NumPy-based implementation of box management operations. + + This class provides methods for converting between different bounding box formats, + normalizing/denormalizing coordinates, and performing transformations using NumPy. + """ + + @staticmethod + def _as_floating(boxes: np.ndarray) -> np.ndarray: + """Return boxes in a floating dtype suitable for coordinate arithmetic.""" + boxes = np.asarray(boxes) + if not np.issubdtype(boxes.dtype, np.floating): + boxes = boxes.astype(np.float32) + return boxes + + @staticmethod + def _coordinate_scale(boxes: np.ndarray, size: Tuple[int, int]) -> np.ndarray: + """Build a scale that leaves prediction fields after coordinates unchanged.""" + if boxes.shape[-1] < 4: + raise ValueError("Boxes must contain at least four coordinate columns.") + size = np.asarray(size, dtype=boxes.dtype).reshape(-1) + if size.size != 2: + raise ValueError("Image size must contain width and height.") + return np.concatenate([np.tile(size, 2), np.ones(boxes.shape[-1] - 4, dtype=boxes.dtype)]) + + @staticmethod + def normalize_boxes(raw_boxes: np.ndarray, image_source_size: Tuple[int, int]) -> np.ndarray: + """ + Normalize bounding box coordinates from pixel values to [0, 1] range. + + Divides x-coordinates by image width and y-coordinates by image height + to convert from absolute pixel values to normalized coordinates. + + Parameters + ---------- + raw_boxes + Boxes in pixel coordinates of shape (N, 4+). + image_source_size + Image dimensions as (width, height). + + Returns + ------- + normalized_boxes + Normalized boxes with coordinates in [0, 1] range, same shape as input. + """ + raw_boxes = NumpyBoxManager._as_floating(raw_boxes) + image_source_size = np.asarray(image_source_size).reshape(-1) + if image_source_size.size != 2 or np.any(image_source_size == 0): + raise ValueError("Image width and height must be greater than zero for normalization.") + return raw_boxes / NumpyBoxManager._coordinate_scale(raw_boxes, image_source_size) + + @staticmethod + def box_cxcywh_to_xyxy(normalized_boxes: np.ndarray) -> np.ndarray: + """ + Convert boxes from CXCYWH to XYXY format. + + Transforms from (center_x, center_y, width, height) format to + (x_min, y_min, x_max, y_max) format by computing the corner coordinates + from the center point and dimensions. + + Parameters + ---------- + normalized_boxes + Boxes in CXCYWH format of shape (N, 4). + + Returns + ------- + boxes + Boxes in XYXY format of shape (N, 4). + """ + normalized_boxes = NumpyBoxManager._as_floating(normalized_boxes) + coordinates = normalized_boxes[..., :4] + x_c = coordinates[..., 0] + y_c = coordinates[..., 1] + w = coordinates[..., 2] + h = coordinates[..., 3] + b = np.stack([x_c - 0.5 * w, y_c - 0.5 * h, x_c + 0.5 * w, y_c + 0.5 * h], axis=-1) + return np.concatenate([b, normalized_boxes[..., 4:]], axis=-1) + + @staticmethod + def box_xywh_to_xyxy(normalized_boxes: np.ndarray) -> np.ndarray: + """ + Convert boxes from XYWH to XYXY format. + + Transforms from (x_min, y_min, width, height) format to + (x_min, y_min, x_max, y_max) format by computing the bottom-right + corner from the top-left corner and dimensions. + + Parameters + ---------- + normalized_boxes + Boxes in XYWH format of shape (N, 4). + + Returns + ------- + boxes + Boxes in XYXY format of shape (N, 4). + """ + normalized_boxes = NumpyBoxManager._as_floating(normalized_boxes) + coordinates = normalized_boxes[..., :4] + x = coordinates[..., 0] + y = coordinates[..., 1] + w = coordinates[..., 2] + h = coordinates[..., 3] + b = np.stack([x, y, x + w, y + h], axis=-1) + return np.concatenate([b, normalized_boxes[..., 4:]], axis=-1) + + @staticmethod + def box_xyxy_to_cxcywh(xyxy_boxes: np.ndarray) -> np.ndarray: + """ + Convert boxes from XYXY to CXCYWH format. + + Transforms from (x_min, y_min, x_max, y_max) format to + (center_x, center_y, width, height) format. + + Parameters + ---------- + xyxy_boxes + Boxes in XYXY format of shape (N, 4). + + Returns + ------- + boxes + Boxes in CXCYWH format of shape (N, 4). + """ + xyxy_boxes = NumpyBoxManager._as_floating(xyxy_boxes) + coordinates = xyxy_boxes[..., :4] + x_min = coordinates[..., 0] + y_min = coordinates[..., 1] + x_max = coordinates[..., 2] + y_max = coordinates[..., 3] + w = x_max - x_min + h = y_max - y_min + x_c = x_min + 0.5 * w + y_c = y_min + 0.5 * h + b = np.stack([x_c, y_c, w, h], axis=-1) + return np.concatenate([b, xyxy_boxes[..., 4:]], axis=-1) + + @staticmethod + def box_xyxy_to_xywh(xyxy_boxes: np.ndarray) -> np.ndarray: + """ + Convert boxes from XYXY to XYWH format. + + Transforms from (x_min, y_min, x_max, y_max) format to + (x_min, y_min, width, height) format. + + Parameters + ---------- + xyxy_boxes + Boxes in XYXY format of shape (N, 4). + + Returns + ------- + boxes + Boxes in XYWH format of shape (N, 4). + """ + xyxy_boxes = NumpyBoxManager._as_floating(xyxy_boxes) + coordinates = xyxy_boxes[..., :4] + x_min = coordinates[..., 0] + y_min = coordinates[..., 1] + x_max = coordinates[..., 2] + y_max = coordinates[..., 3] + w = x_max - x_min + h = y_max - y_min + b = np.stack([x_min, y_min, w, h], axis=-1) + return np.concatenate([b, xyxy_boxes[..., 4:]], axis=-1) + + @staticmethod + def denormalize_boxes(boxes: np.ndarray, size: Tuple[int, int]) -> np.ndarray: + """ + Convert normalized boxes from [0, 1] range to pixel coordinates. + + Multiplies x-coordinates by image width and y-coordinates by image height + to convert from normalized coordinates to absolute pixel values. + + Parameters + ---------- + boxes + Boxes in normalized coordinates [0, 1] of shape (N, 4+). + size + Image dimensions as (width, height). + + Returns + ------- + denormalized_boxes + Boxes in pixel coordinates, same shape as input. + """ + boxes = NumpyBoxManager._as_floating(boxes) + return boxes * NumpyBoxManager._coordinate_scale(boxes, size) + + @staticmethod + def to_numpy_tuple(*arrays) -> Tuple: + """ + Convert one or more numpy arrays to tuple of numpy arrays. + Always returns a tuple, even if a single array is provided. + + Parameters + ---------- + *arrays + Variable number of arrays to convert + + Returns + ------- + numpy_arrays + Tuple of numpy arrays + """ + return tuple(t for t in arrays) + + @staticmethod + def probas_argmax(proba: np.ndarray) -> int: + """ + Get the class ID from a probability array. + + Parameters + ---------- + proba + Probability array for a single detection + + Returns + ------- + class_id + Class ID as Python int + """ + return int(proba.argmax()) + + +class NumpyBoxCoordinatesTranslator(BaseBoxCoordinatesTranslator): + """ + Translates bounding boxes between different coordinate formats using NumPy. + + Mirrors TorchBoxCoordinatesTranslator and TfBoxCoordinatesTranslator for + use when the MultiBoxTensor is backed by plain NumPy arrays. + """ + + def __init__(self, input_box_type: BoxType, output_box_type: BoxType) -> None: + """ + Initialize the NumPy box coordinates translator. + + Parameters + ---------- + input_box_type + Format specification of input boxes. + output_box_type + Desired format specification for output boxes. + """ + self.input_box_type = input_box_type + self.output_box_type = output_box_type + self._box_manager = NumpyBoxManager() + + def translate(self, box: np.ndarray, image_size: Tuple[int, int] = None) -> np.ndarray: + """ + Translate boxes from input format to output format. + + Performs a multi-step conversion: + 1. Normalize input boxes if needed + 2. Convert to XYXY intermediate format + 3. Convert to output format + 4. Denormalize if needed + + Parameters + ---------- + box + Bounding boxes in input format with shape (N, 4). + image_size + Image dimensions as (width, height). Required if input or output + boxes are not normalized. + + Returns + ------- + translated_boxes + Boxes in output format with shape (N, 4). + + Raises + ------ + ValueError + If image_size is None when required for non-normalized boxes. + """ + box = NumpyBoxManager._as_floating(box) + + # Early return if input and output formats are identical + if ( + self.input_box_type.format == self.output_box_type.format + and self.input_box_type.is_normalized == self.output_box_type.is_normalized + ): + return box + + # normalize the input box if needed + if not self.input_box_type.is_normalized: + if image_size is None: + raise ValueError("Input image size must be provided for non-normalized boxes.") + box = NumpyBoxManager.normalize_boxes(box, image_size) + + # convert the input box to XYXY format if needed + if self.input_box_type.format is BoxFormat.CXCYWH: + box = NumpyBoxManager.box_cxcywh_to_xyxy(box) + elif self.input_box_type.format is BoxFormat.XYWH: + box = NumpyBoxManager.box_xywh_to_xyxy(box) + + # convert to the output format + if self.output_box_type.format is BoxFormat.CXCYWH: + box = NumpyBoxManager.box_xyxy_to_cxcywh(box) + elif self.output_box_type.format is BoxFormat.XYWH: + box = NumpyBoxManager.box_xyxy_to_xywh(box) + + # denormalize if needed + if not self.output_box_type.is_normalized: + if image_size is None: + raise ValueError("Output image size must be provided for non-normalized boxes.") + box = NumpyBoxManager.denormalize_boxes(box, image_size) + + return box + + @property + def box_manager(self) -> NumpyBoxManager: + return self._box_manager diff --git a/xplique/utils_functions/object_detection/base/multi_box_tensor.py b/xplique/utils_functions/object_detection/base/multi_box_tensor.py new file mode 100644 index 00000000..ceef3355 --- /dev/null +++ b/xplique/utils_functions/object_detection/base/multi_box_tensor.py @@ -0,0 +1,64 @@ +""" +Protocol for multi-box tensor representations in object detection. + +This module defines the MultiBoxTensor Protocol, which provides a common interface +for tensors containing multiple detection boxes across different frameworks (TensorFlow, PyTorch). +""" + +from typing import Any, Protocol, runtime_checkable + +from xplique.commons.prediction_types import StructuredPrediction + +# pylint: disable=unnecessary-ellipsis + + +@runtime_checkable +class BaseMultiBoxTensor(StructuredPrediction, Protocol): + """ + Protocol for tensors containing multiple detection boxes. + + Tensor with shape (N, C) where: + - N is the number of boxes + - C is the encoding of a bounding box prediction, C = 4 + 1 + nb_classes + - 4 coordinates (box coordinates) + - 1 score (objectness/detection confidence) + - nb_classes (soft class predictions or one-hot encoded class predictions) + + Example: [100, 85] for 100 boxes with 80 classes (COCO dataset: 4 + 1 + 80 = 85) + + This is a Protocol (structural type) rather than an ABC to avoid metaclass conflicts + with framework-specific tensor types (torch.Tensor, tf.Tensor, etc.) + """ + + def boxes(self) -> Any: + """ + Return box coordinates tensor with shape (N, 4). + + Returns + ------- + tensor + Box coordinates in the format [x1, y1, x2, y2]. + """ + ... + + def scores(self) -> Any: + """ + Return detection scores tensor with shape (N,). + + Returns + ------- + tensor + Objectness or detection confidence scores. + """ + ... + + def probas(self) -> Any: + """ + Return class probabilities tensor with shape (N, nb_classes). + + Returns + ------- + tensor + Class probabilities or one-hot encoded class predictions. + """ + ... diff --git a/xplique/utils_functions/object_detection.py b/xplique/utils_functions/object_detection/object_detection_operators.py similarity index 100% rename from xplique/utils_functions/object_detection.py rename to xplique/utils_functions/object_detection/object_detection_operators.py diff --git a/xplique/utils_functions/object_detection/tf/__init__.py b/xplique/utils_functions/object_detection/tf/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/xplique/utils_functions/object_detection/tf/box_formatter.py b/xplique/utils_functions/object_detection/tf/box_formatter.py new file mode 100644 index 00000000..ccfe4489 --- /dev/null +++ b/xplique/utils_functions/object_detection/tf/box_formatter.py @@ -0,0 +1,107 @@ +""" +TensorFlow implementation of box formatter for object detection predictions. + +This module provides TensorFlow-specific box formatting and coordinate translation +for object detection model outputs, handling conversion between different box formats +and normalization states. +""" + +from abc import ABC, abstractmethod + +import tensorflow as tf + +from xplique.utils_functions.object_detection.base.box_formatter import ( + BaseBoxFormatter, +) +from xplique.utils_functions.object_detection.base.box_manager import ( + BoxFormat, + BoxType, +) +from xplique.utils_functions.object_detection.tf.box_manager import ( + TfBoxCoordinatesTranslator, +) +from xplique.utils_functions.object_detection.tf.multi_box_tensor import TfMultiBoxTensor + + +class TfBaseBoxFormatter(BaseBoxFormatter, ABC): + """ + TensorFlow implementation of the BaseBoxFormatter interface. + + Provides TensorFlow-specific box formatting and coordinate translation + for object detection predictions. Handles conversion between different + box formats and normalization states using TensorFlow operations. + """ + + def __init__( + self, + input_box_type: BoxType, + output_box_type: BoxType = BoxType(BoxFormat.XYXY, is_normalized=True), + ) -> None: + """ + Initialize the TensorFlow box formatter. + + Parameters + ---------- + input_box_type + The format and normalization of input boxes. + output_box_type + The desired format and normalization for output boxes. + """ + super().__init__(input_box_type=input_box_type, output_box_type=output_box_type) + self.box_translator = TfBoxCoordinatesTranslator(self.input_box_type, self.output_box_type) + + @abstractmethod + def forward(self, predictions): + """ + Transform model predictions to MultiBoxTensor format. + + This abstract method must be implemented by subclasses to handle + framework-specific prediction formats. + + Parameters + ---------- + predictions + Model predictions in framework-specific format. + + Raises + ------ + NotImplementedError + This method must be implemented by subclasses. + """ + raise NotImplementedError("This method should be implemented in the subclass") + + def format_predictions(self, predictions, image_size=None) -> TfMultiBoxTensor: + """ + Format detection predictions into unified MultiBoxTensor representation. + + Translates box coordinates to the desired format and concatenates with + scores and class probabilities into a single tensor. + + Parameters + ---------- + predictions + Dictionary with 'boxes', 'scores', and 'probas' keys. + image_size + Optional size of the image for coordinate conversion. + + Returns + ------- + formatted_predictions + MultiBoxTensor containing formatted predictions with shape (N, 4+1+num_classes). + """ + boxes = predictions["boxes"] + boxes = self.box_translator.translate(boxes, image_size=image_size) + probas = predictions["probas"] + scores = predictions["scores"] + if scores.shape.rank == 1: + scores = scores[..., tf.newaxis] + return TfMultiBoxTensor( + tf.concat( + [ + boxes, # boxes coordinates + scores, # detection probability + probas, + ], # class logits predictions for the given box + axis=1, + ) + ) diff --git a/xplique/utils_functions/object_detection/tf/box_manager.py b/xplique/utils_functions/object_detection/tf/box_manager.py new file mode 100644 index 00000000..1763a9cc --- /dev/null +++ b/xplique/utils_functions/object_detection/tf/box_manager.py @@ -0,0 +1,339 @@ +""" +TensorFlow implementation for bounding box management operations. +""" + +from typing import Optional, Tuple + +import tensorflow as tf + +from xplique.utils_functions.object_detection.base.box_manager import ( + BaseBoxCoordinatesTranslator, + BoxFormat, + BoxManager, + BoxType, +) + + +class TfBoxManager(BoxManager): + """ + TensorFlow implementation of box management operations. + + Provides TensorFlow-specific methods for converting between different + bounding box formats, normalizing/denormalizing coordinates, and performing + transformations using TensorFlow operations with gradient support. + """ + + @staticmethod + def _as_floating(boxes: tf.Tensor) -> tf.Tensor: + """Return boxes in a floating dtype suitable for coordinate arithmetic.""" + boxes = tf.convert_to_tensor(boxes) + if not boxes.dtype.is_floating: + boxes = tf.cast(boxes, tf.float32) + return boxes + + @staticmethod + def _coordinate_scale(boxes: tf.Tensor, size: tf.Tensor) -> tf.Tensor: + """Build a scale that leaves prediction fields after coordinates unchanged.""" + size = tf.reshape(tf.cast(tf.convert_to_tensor(size), boxes.dtype), [-1]) + tf.debugging.assert_equal( + tf.size(size), 2, message="Image size must contain width and height." + ) + tf.debugging.assert_greater_equal( + tf.shape(boxes)[-1], 4, message="Boxes must contain at least four coordinate columns." + ) + return tf.concat( + [ + tf.tile(size, [2]), + tf.ones(tf.reshape(tf.shape(boxes)[-1] - 4, [1]), dtype=boxes.dtype), + ], + axis=0, + ) + + @staticmethod + def normalize_boxes(raw_boxes: tf.Tensor, image_source_size: tf.Tensor) -> tf.Tensor: + """ + Normalize bounding box coordinates from pixel values to [0, 1] range. + + Parameters + ---------- + raw_boxes + Boxes in pixel coordinates of shape (N, 4+). + image_source_size + Image dimensions as tensor (width, height). + + Returns + ------- + normalized_boxes + Normalized boxes with coordinates in [0, 1] range, same shape as input. + """ + raw_boxes = TfBoxManager._as_floating(raw_boxes) + size = tf.cast(tf.convert_to_tensor(image_source_size), raw_boxes.dtype) + tf.debugging.assert_positive( + size, message="Image width and height must be greater than zero for normalization." + ) + return raw_boxes / TfBoxManager._coordinate_scale(raw_boxes, size) + + @staticmethod + def box_cxcywh_to_xyxy(normalized_boxes: tf.Tensor) -> tf.Tensor: + """ + Convert boxes from CXCYWH to XYXY format. + + Transforms from (center_x, center_y, width, height) format to + (x_min, y_min, x_max, y_max) format by computing the corner coordinates + from the center point and dimensions. + + Parameters + ---------- + normalized_boxes + Boxes in CXCYWH format of shape (N, 4). + + Returns + ------- + boxes + Boxes in XYXY format of shape (N, 4). + """ + normalized_boxes = TfBoxManager._as_floating(normalized_boxes) + x_c, y_c, w, h = tf.unstack(normalized_boxes[..., :4], axis=-1) + x_min = x_c - 0.5 * w + y_min = y_c - 0.5 * h + x_max = x_c + 0.5 * w + y_max = y_c + 0.5 * h + b = tf.stack([x_min, y_min, x_max, y_max], axis=-1) + return tf.concat([b, normalized_boxes[..., 4:]], axis=-1) + + @staticmethod + def box_xywh_to_xyxy(normalized_boxes: tf.Tensor) -> tf.Tensor: + """ + Convert boxes from XYWH to XYXY format. + + Transforms from (x_min, y_min, width, height) format to + (x_min, y_min, x_max, y_max) format by computing the bottom-right + corner from the top-left corner and dimensions. + + Parameters + ---------- + normalized_boxes + Boxes in XYWH format of shape (N, 4). + + Returns + ------- + boxes + Boxes in XYXY format of shape (N, 4). + """ + normalized_boxes = TfBoxManager._as_floating(normalized_boxes) + x, y, w, h = tf.unstack(normalized_boxes[..., :4], axis=-1) # extract the columns + b = [x, y, x + w, y + h] + return tf.concat([tf.stack(b, axis=-1), normalized_boxes[..., 4:]], axis=-1) + + @staticmethod + def box_xyxy_to_cxcywh(xyxy_boxes: tf.Tensor) -> tf.Tensor: + """ + Convert boxes from XYXY to CXCYWH format. + + Transforms from (x_min, y_min, x_max, y_max) format to + (center_x, center_y, width, height) format by computing the center + point and dimensions from the corner coordinates. + + Parameters + ---------- + xyxy_boxes + Boxes in XYXY format of shape (N, 4). + + Returns + ------- + boxes + Boxes in CXCYWH format of shape (N, 4). + """ + xyxy_boxes = TfBoxManager._as_floating(xyxy_boxes) + x_min, y_min, x_max, y_max = tf.unstack(xyxy_boxes[..., :4], axis=-1) + w = x_max - x_min + h = y_max - y_min + x_c = x_min + 0.5 * w + y_c = y_min + 0.5 * h + b = [x_c, y_c, w, h] + return tf.concat([tf.stack(b, axis=-1), xyxy_boxes[..., 4:]], axis=-1) + + @staticmethod + def box_xyxy_to_xywh(xyxy_boxes: tf.Tensor) -> tf.Tensor: + """ + Convert boxes from XYXY to XYWH format. + + Transforms from (x_min, y_min, x_max, y_max) format to + (x_min, y_min, width, height) format by computing the dimensions + from the corner coordinates. + + Parameters + ---------- + xyxy_boxes + Boxes in XYXY format of shape (N, 4). + + Returns + ------- + boxes + Boxes in XYWH format of shape (N, 4). + """ + xyxy_boxes = TfBoxManager._as_floating(xyxy_boxes) + x_min, y_min, x_max, y_max = tf.unstack(xyxy_boxes[..., :4], axis=-1) + w = x_max - x_min + h = y_max - y_min + b = [x_min, y_min, w, h] + return tf.concat([tf.stack(b, axis=-1), xyxy_boxes[..., 4:]], axis=-1) + + @staticmethod + def denormalize_boxes(boxes: tf.Tensor, size: tf.Tensor) -> tf.Tensor: + """ + Convert normalized boxes from [0, 1] range to pixel coordinates. + + Multiplies x-coordinates by image width and y-coordinates by image height + to convert from normalized coordinates to absolute pixel values. + + Parameters + ---------- + boxes + Boxes in normalized coordinates [0, 1] of shape (N, 4+). + size + Image dimensions as tensor (width, height). + + Returns + ------- + denormalized_boxes + Boxes in pixel coordinates, same shape as input. + """ + boxes = TfBoxManager._as_floating(boxes) + return boxes * TfBoxManager._coordinate_scale(boxes, size) + + @staticmethod + def to_numpy_tuple(*tensors) -> Tuple: + """ + Convert one or more TensorFlow tensors to tuple of NumPy arrays. + + Parameters + ---------- + *tensors + Variable number of tensors or arrays to convert. + + Returns + ------- + arrays + Tuple of NumPy arrays corresponding to input tensors. + """ + return tuple(t.numpy() if isinstance(t, tf.Tensor) else t for t in tensors) + + @staticmethod + def probas_argmax(proba: tf.Tensor) -> int: + """ + Get the class ID with highest probability from a probability tensor. + + Finds the index of the maximum probability value and converts it to a + Python integer for use as a class identifier. + + Parameters + ---------- + proba + Probability tensor for a single detection of shape (num_classes,). + + Returns + ------- + class_id + Class ID as Python int corresponding to highest probability. + """ + return int(tf.argmax(proba).numpy()) + + +class TfBoxCoordinatesTranslator(BaseBoxCoordinatesTranslator): + """ + Translates bounding boxes between different coordinate formats and scales. + + Handles the full pipeline of box coordinate transformations including: + - Normalization/denormalization + - Format conversion (CXCYWH, XYWH, XYXY) + - Image size scaling + + All operations use TensorFlow to maintain gradient flow for attribution methods. + """ + + def __init__(self, input_box_type: BoxType, output_box_type: BoxType) -> None: + """ + Initialize the box coordinates translator. + + Parameters + ---------- + input_box_type + The format and normalization of input boxes. + output_box_type + The desired format and normalization for output boxes. + """ + self.input_box_type = input_box_type + self.output_box_type = output_box_type + self._box_manager = TfBoxManager() + + @property + def box_manager(self) -> TfBoxManager: + return self._box_manager + + def translate( + self, + box: tf.Tensor, + image_size: Optional[tf.TensorShape] = None, + ) -> tf.Tensor: + """ + Translate box coordinates from input format/scale to output format/scale. + + Performs a complete transformation pipeline: + 1. Normalize boxes if input is in pixel coordinates + 2. Convert to XYXY format as intermediate representation + 3. Convert from XYXY to desired output format + 4. Denormalize boxes if output should be in pixel coordinates + + Parameters + ---------- + box + Box tensor of shape (N, 4) to translate. + image_size + Image dimensions as (width, height). Required if input or output + boxes are not normalized. + + Returns + ------- + translated_box + Translated box tensor in the desired format and scale. + + Raises + ------ + ValueError + If image size is required but not provided. + """ + box = TfBoxManager._as_floating(box) + + # Early return if input and output formats are identical + if ( + self.input_box_type.format == self.output_box_type.format + and self.input_box_type.is_normalized == self.output_box_type.is_normalized + ): + return box + + # normalize the input box if needed + if not self.input_box_type.is_normalized: + if image_size is None: + raise ValueError("Input image size must be provided for non-normalized boxes.") + box = TfBoxManager.normalize_boxes(box, image_size) + + # convert the input box to XYXY format if needed + if self.input_box_type.format is BoxFormat.CXCYWH: + box = TfBoxManager.box_cxcywh_to_xyxy(box) + elif self.input_box_type.format is BoxFormat.XYWH: + box = TfBoxManager.box_xywh_to_xyxy(box) + + # now convert to the output format + if self.output_box_type.format is BoxFormat.CXCYWH: + box = TfBoxManager.box_xyxy_to_cxcywh(box) + elif self.output_box_type.format is BoxFormat.XYWH: + box = TfBoxManager.box_xyxy_to_xywh(box) + + # denormalize the box to the output image size if needed + if not self.output_box_type.is_normalized: + if image_size is None: + raise ValueError("Output image size must be provided for non-normalized boxes.") + box = TfBoxManager.denormalize_boxes(box, image_size) + + return box diff --git a/xplique/utils_functions/object_detection/tf/box_model_wrapper.py b/xplique/utils_functions/object_detection/tf/box_model_wrapper.py new file mode 100644 index 00000000..d4e07378 --- /dev/null +++ b/xplique/utils_functions/object_detection/tf/box_model_wrapper.py @@ -0,0 +1,71 @@ +""" +TensorFlow model wrappers for object detection models with box formatting. +""" + +from typing import List, Union + +import tensorflow as tf + +from xplique.utils_functions.object_detection.tf.box_formatter import TfBaseBoxFormatter +from xplique.utils_functions.object_detection.tf.multi_box_tensor import TfMultiBoxTensor +from xplique.utils_functions.output_as_list_mixin import OutputAsListMixin + + +def _pad_and_stack_box_predictions(predictions: List[TfMultiBoxTensor]) -> tf.Tensor: + """Stack variable detection counts, using all-zero rows as padding.""" + predictions = [tf.convert_to_tensor(prediction) for prediction in predictions] + return tf.ragged.stack(predictions).to_tensor() + + +class TfBoxesModelWrapper(OutputAsListMixin, tf.keras.Model): + """ + Wrapper for TensorFlow object detection models with box formatting capabilities. + + This class wraps an object detection model and applies a box formatter to its outputs. + It can return predictions either as a list of formatted boxes (one per image) or as + a single stacked tensor. + """ + + def __init__(self, model: tf.keras.Model, box_formatter: TfBaseBoxFormatter) -> None: + """ + Initialize the TensorFlow box model wrapper. + + Parameters + ---------- + model + TensorFlow object detection model to wrap. + box_formatter + Formatter to process and convert model predictions to Xplique format. + """ + super().__init__() + self.model = model + self.box_formatter = box_formatter + self.output_as_list = True + + def call(self, x, **kwargs) -> Union[tf.Tensor, List[TfMultiBoxTensor]]: + """ + Forward pass through the wrapped model with box formatting. + + Processes input through the object detection model and formats the predictions + using the box formatter. Returns either a list or stacked tensor based on the + output_as_list flag. + + Parameters + ---------- + x + Input tensor of shape (batch_size, height, width, channels). + **kwargs + Additional keyword arguments to pass to the model. + + Returns + ------- + predictions + If output_as_list is True: List of MultiBoxTensor objects, one per image. + If output_as_list is False: Stacked tensor of formatted predictions with shape + (batch_size, max_num_boxes, features), with all-zero padding. + """ + predictions = self.model(x, **kwargs) + list_of_predictions = self.box_formatter(predictions) + if self.output_as_list: + return list_of_predictions + return _pad_and_stack_box_predictions(list_of_predictions) diff --git a/xplique/utils_functions/object_detection/tf/multi_box_tensor.py b/xplique/utils_functions/object_detection/tf/multi_box_tensor.py new file mode 100644 index 00000000..86c5c14c --- /dev/null +++ b/xplique/utils_functions/object_detection/tf/multi_box_tensor.py @@ -0,0 +1,157 @@ +""" +TensorFlow implementation of MultiBoxTensor for object detection predictions. + +This module provides a TensorFlow wrapper for multi-box detection predictions +with a unified format compatible with the MultiBoxTensor protocol. +""" + +from typing import Optional + +import tensorflow as tf + +from xplique.utils_functions.object_detection.base.multi_box_tensor import BaseMultiBoxTensor + + +class TfMultiBoxTensor(BaseMultiBoxTensor): + """ + TensorFlow wrapper for multi-box detection predictions with unified format. + + Encapsulates a tensor with shape (N, C) where: + - N is the number of detected boxes + - C = 4 + 1 + nb_classes encoding: 4 box coordinates, 1 objectness score, + nb_classes class predictions (soft probabilities or one-hot encoding) + + This class provides convenient access to boxes, scores, and class probabilities, + along with filtering capabilities and TensorFlow integration support. + """ + + # Ex: [9, 85] for 9 boxes, 80 classes (COCO dataset) + + def __init__(self, tensor: tf.Tensor) -> None: + """ + Initialize the MultiBoxTensor. + + Parameters + ---------- + tensor + TensorFlow tensor of shape (N, C) containing box predictions. + """ + self.tensor = tensor + + def __len__(self) -> int: + return self.tensor.shape[0] + + @property + def shape(self) -> tf.TensorShape: + """Get the shape of the underlying tensor.""" + return self.tensor.shape + + @property + def dtype(self) -> tf.DType: + """Get the data type of the underlying tensor.""" + return self.tensor.dtype + + def __tensor__(self) -> tf.Tensor: + # Get the underlying TensorFlow tensor. + return self.tensor + + def __tf_tensor__( + self, dtype: Optional[tf.DType] = None, name: Optional[str] = None + ) -> tf.Tensor: + # Critical: This makes tf.stack() work by converting to tensor + return tf.convert_to_tensor(self.tensor, dtype=dtype, name=name) + + def __getitem__(self, key) -> tf.Tensor: + # Support indexing operations + return self.tensor[key] + + def boxes(self) -> tf.Tensor: + """ + Extract bounding box coordinates from predictions. + + Returns + ------- + boxes + Tensor of shape (N, 4) containing box coordinates. + """ + return self[:, :4] + + def scores(self) -> tf.Tensor: + """ + Extract objectness scores from predictions. + + Returns + ------- + scores + Tensor of shape (N,) containing detection confidence scores. + """ + return self[:, 4] + + def probas(self) -> tf.Tensor: + """ + Extract class probability predictions from predictions. + + Returns + ------- + probas + Tensor of shape (N, num_classes) containing class probabilities. + """ + return self[:, 5:] + + def filter( + self, class_id: Optional[int] = None, confidence: Optional[float] = None + ) -> "TfMultiBoxTensor": + """ + Filter boxes by class ID and/or detection confidence threshold. + + Filters the detections based on predicted class and confidence score. + If both filters are provided, boxes must satisfy both conditions. + + Parameters + ---------- + class_id + Optional class ID to filter for specific object class. + confidence + Optional minimum confidence score threshold. + + Returns + ------- + filtered_tensor + New MultiBoxTensor containing only filtered detections. + """ + if class_id is None and confidence is None: + return self + probas = self.probas() + class_ids = tf.argmax(probas, axis=-1) + scores = self.scores() + if class_id is not None and confidence is not None: + keep = (class_ids == class_id) & (scores >= confidence) + elif class_id is None: + keep = scores >= confidence + else: + keep = class_ids == class_id + filtered_tensor = tf.boolean_mask(self.tensor, keep) + return TfMultiBoxTensor(filtered_tensor) + + def to_attribution_target(self, class_id=None): + """ + Return self as the attribution target. + + For object detection, filter() has already selected the relevant boxes; + the filtered box tensor (coordinates + scores) is directly the target + for the attribution method. ``class_id`` is ignored here. + """ + return self + + def to_batched_tensor(self) -> tf.Tensor: + """ + Add batch dimension to MultiBoxTensor. + + Converts (num_boxes, features) -> (1, num_boxes, features) + + Returns + ------- + tensor + tf.Tensor with batch dimension added + """ + return tf.expand_dims(self.tensor, axis=0) diff --git a/xplique/utils_functions/object_detection/torch/__init__.py b/xplique/utils_functions/object_detection/torch/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/xplique/utils_functions/object_detection/torch/box_formatter.py b/xplique/utils_functions/object_detection/torch/box_formatter.py new file mode 100644 index 00000000..fcf1ec7b --- /dev/null +++ b/xplique/utils_functions/object_detection/torch/box_formatter.py @@ -0,0 +1,104 @@ +""" +PyTorch box formatters for converting object detection model outputs to Xplique format. +""" + +from abc import ABC, abstractmethod +from typing import Any, Dict, List, Optional + +import torch + +from xplique.utils_functions.object_detection.base.box_formatter import ( + BaseBoxFormatter, +) +from xplique.utils_functions.object_detection.base.box_manager import ( + BoxFormat, + BoxType, +) +from xplique.utils_functions.object_detection.torch.box_manager import ( + TorchBoxCoordinatesTranslator, +) +from xplique.utils_functions.object_detection.torch.multi_box_tensor import ( + TorchMultiBoxTensor, +) + + +class TorchBaseBoxFormatter(BaseBoxFormatter, ABC): + """ + Abstract base class for PyTorch-based box formatters. + + This class combines BaseBoxFormatter functionality with PyTorch nn.Module + capabilities, enabling gradient computation through box formatting operations. + """ + + def __init__( + self, + input_box_type: BoxType, + output_box_type: BoxType = BoxType(BoxFormat.XYXY, is_normalized=True), + ) -> None: + """ + Initialize the PyTorch box formatter. + + Parameters + ---------- + input_box_type + Format of input bounding boxes (coordinate system and normalization). + output_box_type + Desired output format for bounding boxes. Defaults to normalized XYXY. + """ + super().__init__(input_box_type=input_box_type, output_box_type=output_box_type) + self.box_translator = TorchBoxCoordinatesTranslator( + self.input_box_type, self.output_box_type + ) + + @abstractmethod + def forward(self, predictions: Any) -> List[TorchMultiBoxTensor]: + """ + Transform model predictions into Xplique MultiBoxTensor format. + + Parameters + ---------- + predictions + Raw predictions from the object detection model. + + Returns + ------- + formatted_predictions + List of MultiBoxTensor objects, one per image in the batch. + """ + raise NotImplementedError("This method should be implemented in the subclass") + + # needs boxes, scores, probas + def format_predictions( + self, predictions: Dict[str, torch.Tensor], image_size: Optional[torch.Size] = None + ) -> TorchMultiBoxTensor: + """ + Convert prediction dictionary to MultiBoxTensor format. + + Parameters + ---------- + predictions + Dictionary containing 'boxes', 'scores', and 'probas' keys. + image_size + Optional image size for normalization/denormalization of boxes. + + Returns + ------- + formatted_tensor + MultiBoxTensor with concatenated boxes, scores, and class probabilities. + """ + boxes = predictions["boxes"] + boxes = self.box_translator.translate(boxes, image_size=image_size) + probas = predictions["probas"] + scores = predictions["scores"] + if scores.ndim == 1: + scores = scores.unsqueeze(-1) + return TorchMultiBoxTensor( + torch.cat( + [ + boxes, # boxes coordinates + scores, # detection probability + probas, + ], # class logits predictions for the given box + dim=1, + ) + ) diff --git a/xplique/utils_functions/object_detection/torch/box_manager.py b/xplique/utils_functions/object_detection/torch/box_manager.py new file mode 100644 index 00000000..52d61d23 --- /dev/null +++ b/xplique/utils_functions/object_detection/torch/box_manager.py @@ -0,0 +1,300 @@ +""" +PyTorch-specific box management utilities for coordinate transformations. +""" + +from typing import Optional, Tuple + +import torch + +from xplique.utils_functions.object_detection.base.box_manager import ( + BaseBoxCoordinatesTranslator, + BoxFormat, + BoxManager, + BoxType, +) + + +class TorchBoxManager(BoxManager): + """ + PyTorch implementation of box management for bounding box operations. + + This class provides static methods for converting between different bounding box + coordinate formats and handling normalization/denormalization using PyTorch tensors. + """ + + @staticmethod + def _as_floating(boxes: torch.Tensor) -> torch.Tensor: + """Return boxes in a floating dtype suitable for coordinate arithmetic.""" + boxes = torch.as_tensor(boxes) + if not boxes.is_floating_point(): + boxes = boxes.to(torch.float32) + return boxes + + @staticmethod + def _coordinate_scale(boxes: torch.Tensor, size: torch.Size) -> torch.Tensor: + """Build a scale that leaves prediction fields after coordinates unchanged.""" + if boxes.shape[-1] < 4: + raise ValueError("Boxes must contain at least four coordinate columns.") + size = torch.as_tensor(size, dtype=boxes.dtype, device=boxes.device).reshape(-1) + if size.numel() != 2: + raise ValueError("Image size must contain width and height.") + trailing_scale = torch.ones(boxes.shape[-1] - 4, dtype=boxes.dtype, device=boxes.device) + return torch.cat([size.repeat(2), trailing_scale]) + + @staticmethod + def normalize_boxes(raw_boxes: torch.Tensor, image_source_size: torch.Size) -> torch.Tensor: + """ + Normalize bounding box coordinates to [0,1] range based on image size. + + Parameters + ---------- + raw_boxes + Boxes in pixel coordinates with shape (N, 4). + image_source_size + Image dimensions as (width, height). + + Returns + ------- + normalized_boxes + Normalized boxes with coordinates in [0,1] range. + """ + raw_boxes = TorchBoxManager._as_floating(raw_boxes) + image_source_size = torch.as_tensor(image_source_size, device=raw_boxes.device).reshape(-1) + if image_source_size.numel() != 2 or torch.any(image_source_size == 0): + raise ValueError("Image width and height must be greater than zero for normalization.") + return raw_boxes / TorchBoxManager._coordinate_scale(raw_boxes, image_source_size) + + @staticmethod + def box_cxcywh_to_xyxy(normalized_boxes: torch.Tensor) -> torch.Tensor: + """ + Convert boxes from CXCYWH to corner format XYXY. + + Parameters + ---------- + normalized_boxes + Boxes in CXCYWH format with shape (N, 4). + + Returns + ------- + xyxy_boxes + Boxes in XYXY format with shape (N, 4). + """ + normalized_boxes = TorchBoxManager._as_floating(normalized_boxes) + x_c, y_c, w, h = normalized_boxes[..., :4].unbind(-1) # extract the columns + b = [(x_c - 0.5 * w), (y_c - 0.5 * h), (x_c + 0.5 * w), (y_c + 0.5 * h)] + return torch.cat([torch.stack(b, dim=-1), normalized_boxes[..., 4:]], dim=-1) + + @staticmethod + def box_xyxy_to_cxcywh(xyxy_boxes: torch.Tensor) -> torch.Tensor: + """ + Convert boxes from XYXY to CXCYWH format. + + Parameters + ---------- + xyxy_boxes + Boxes in XYXY format with shape (N, 4). + + Returns + ------- + cxcywh_boxes + Boxes in CXCYWH format with shape (N, 4). + """ + xyxy_boxes = TorchBoxManager._as_floating(xyxy_boxes) + x1, y1, x2, y2 = xyxy_boxes[..., :4].unbind(-1) + w = x2 - x1 + h = y2 - y1 + x_c = x1 + 0.5 * w + y_c = y1 + 0.5 * h + b = [x_c, y_c, w, h] + return torch.cat([torch.stack(b, dim=-1), xyxy_boxes[..., 4:]], dim=-1) + + @staticmethod + def box_xywh_to_xyxy(normalized_boxes: torch.Tensor) -> torch.Tensor: + """ + Convert boxes XYWH to XYXY format. + + Parameters + ---------- + normalized_boxes + Boxes in XYWH format with shape (N, 4). + + Returns + ------- + xyxy_boxes + Boxes in XYXY format with shape (N, 4). + """ + normalized_boxes = TorchBoxManager._as_floating(normalized_boxes) + x, y, w, h = normalized_boxes[..., :4].unbind(-1) + b = [x, y, x + w, y + h] + return torch.cat([torch.stack(b, dim=-1), normalized_boxes[..., 4:]], dim=-1) + + @staticmethod + def box_xyxy_to_xywh(xyxy_boxes: torch.Tensor) -> torch.Tensor: + """ + Convert boxes from XYXY to XYWH format. + + Parameters + ---------- + xyxy_boxes + Boxes in XYXY format with shape (N, 4). + + Returns + ------- + xywh_boxes + Boxes in XYWH format with shape (N, 4). + """ + xyxy_boxes = TorchBoxManager._as_floating(xyxy_boxes) + x_min, y_min, x_max, y_max = xyxy_boxes[..., :4].unbind(-1) + w = x_max - x_min + h = y_max - y_min + b = [x_min, y_min, w, h] + return torch.cat([torch.stack(b, dim=-1), xyxy_boxes[..., 4:]], dim=-1) + + @staticmethod + def denormalize_boxes(boxes: torch.Tensor, size: torch.Size) -> torch.Tensor: + """ + Convert normalized boxes [0,1] to pixel coordinates. + + Parameters + ---------- + boxes + Boxes in normalized coordinates [0,1] + size + Image size (width, height) + + Returns + ------- + denormalized_boxes + Boxes in pixel coordinates + """ + boxes = TorchBoxManager._as_floating(boxes) + return boxes * TorchBoxManager._coordinate_scale(boxes, size) + + @staticmethod + def to_numpy_tuple(*tensors) -> Tuple: + """ + Convert one or more PyTorch tensors to tuple of numpy arrays. + Handles GPU tensors by moving to CPU first. + Always returns a tuple, even if a single tensor is provided. + + Parameters + ---------- + *tensors + Variable number of tensors to convert + + Returns + ------- + numpy_arrays + Tuple of numpy arrays + """ + return tuple( + t.detach().cpu().numpy() if isinstance(t, torch.Tensor) else t for t in tensors + ) + + @staticmethod + def probas_argmax(proba: torch.Tensor) -> int: + """ + Get the class ID from a probability tensor. + + Parameters + ---------- + proba + Probability tensor for a single detection + + Returns + ------- + class_id + Class ID as Python int + """ + return proba.argmax().item() + + +class TorchBoxCoordinatesTranslator(BaseBoxCoordinatesTranslator): + """ + Translates bounding boxes between different coordinate formats for PyTorch tensors. + + This class handles conversions between box formats (XYXY, CXCYWH, XYWH) and + manages normalization/denormalization of coordinates. + """ + + def __init__(self, input_box_type: BoxType, output_box_type: BoxType) -> None: + """ + Initialize the box coordinates translator. + + Parameters + ---------- + input_box_type + Format specification of input boxes. + output_box_type + Desired format specification for output boxes. + """ + self.input_box_type = input_box_type + self.output_box_type = output_box_type + self._box_manager = TorchBoxManager() + + @property + def box_manager(self) -> TorchBoxManager: + return self._box_manager + + def translate(self, box: torch.Tensor, image_size: Optional[torch.Size] = None) -> torch.Tensor: + """ + Translate boxes from input format to output format. + + This method performs a multi-step conversion: + 1. Normalize input boxes if needed + 2. Convert to XYXY intermediate format + 3. Convert to output format + 4. Denormalize if needed + + Parameters + ---------- + box + Bounding boxes in input format with shape (N, 4). + image_size + Image dimensions as (width, height). Required if input or output + boxes are not normalized. + + Returns + ------- + translated_boxes + Boxes in output format with shape (N, 4). + + Raises + ------ + ValueError + If image_size is None when required for non-normalized boxes. + """ + box = TorchBoxManager._as_floating(box) + + # Early return if input and output formats are identical + if ( + self.input_box_type.format == self.output_box_type.format + and self.input_box_type.is_normalized == self.output_box_type.is_normalized + ): + return box + + # normalize the input box if needed + if not self.input_box_type.is_normalized: + if image_size is None: + raise ValueError("Input image size must be provided for non-normalized boxes.") + box = TorchBoxManager.normalize_boxes(box, image_size) + + # convert the input box to XYXY format if needed + if self.input_box_type.format is BoxFormat.CXCYWH: + box = TorchBoxManager.box_cxcywh_to_xyxy(box) + elif self.input_box_type.format is BoxFormat.XYWH: + box = TorchBoxManager.box_xywh_to_xyxy(box) + + # now convert to the output format + if self.output_box_type.format is BoxFormat.CXCYWH: + box = TorchBoxManager.box_xyxy_to_cxcywh(box) + elif self.output_box_type.format is BoxFormat.XYWH: + box = TorchBoxManager.box_xyxy_to_xywh(box) + + # denormalize the box to the output image size if needed + if not self.output_box_type.is_normalized: + if image_size is None: + raise ValueError("Output image size must be provided for non-normalized boxes.") + box = TorchBoxManager.denormalize_boxes(box, image_size) + + return box diff --git a/xplique/utils_functions/object_detection/torch/box_model_wrapper.py b/xplique/utils_functions/object_detection/torch/box_model_wrapper.py new file mode 100644 index 00000000..babf531a --- /dev/null +++ b/xplique/utils_functions/object_detection/torch/box_model_wrapper.py @@ -0,0 +1,72 @@ +"""PyTorch wrapper for object detection models with box formatting capabilities.""" + +from abc import ABC +from typing import List, Union + +import torch + +from xplique.utils_functions.object_detection.torch.box_formatter import ( + TorchBaseBoxFormatter, +) +from xplique.utils_functions.object_detection.torch.multi_box_tensor import TorchMultiBoxTensor +from xplique.utils_functions.output_as_list_mixin import OutputAsListMixin + + +def _pad_and_stack_box_predictions(predictions: List[TorchMultiBoxTensor]) -> torch.Tensor: + """Stack variable detection counts, using all-zero rows as padding.""" + return torch.nn.utils.rnn.pad_sequence(predictions, batch_first=True) + + +class TorchBoxesModelWrapper(OutputAsListMixin, torch.nn.Module, ABC): + """ + Wrapper for PyTorch object detection models with box formatting capabilities. + + This class wraps an object detection model and applies a box formatter to its outputs. + It can return predictions either as a list of formatted boxes (one per image) or as + a single stacked tensor. + """ + + def __init__(self, model: torch.nn.Module, box_formatter: TorchBaseBoxFormatter) -> None: + """ + Initialize the PyTorch box model wrapper. + + Parameters + ---------- + model + PyTorch object detection model to wrap. + box_formatter + Formatter to process and convert model predictions to Xplique format. + """ + super().__init__() + self.model = model + self.box_formatter = box_formatter + self.output_as_list = True + self.eval() # Set training to False by default + + def forward(self, x: torch.Tensor, **kwargs) -> Union[torch.Tensor, List[TorchMultiBoxTensor]]: + """ + Forward pass through the wrapped model with box formatting. + + Processes input through the object detection model and formats the predictions + using the box formatter. Returns either a list or stacked tensor based on the + output_as_list flag. + + Parameters + ---------- + x + Input tensor of shape (batch_size, channels, height, width). + **kwargs + Additional keyword arguments to pass to the model. + + Returns + ------- + predictions + If output_as_list is True: List of MultiBoxTensor objects, one per image. + If output_as_list is False: Stacked tensor of formatted predictions with shape + (batch_size, max_num_boxes, features), with all-zero padding. + """ + predictions = self.model(x, **kwargs) + list_of_predictions = self.box_formatter(predictions) + if self.output_as_list: + return list_of_predictions + return _pad_and_stack_box_predictions(list_of_predictions) diff --git a/xplique/utils_functions/object_detection/torch/multi_box_tensor.py b/xplique/utils_functions/object_detection/torch/multi_box_tensor.py new file mode 100644 index 00000000..20f3d20f --- /dev/null +++ b/xplique/utils_functions/object_detection/torch/multi_box_tensor.py @@ -0,0 +1,152 @@ +""" +PyTorch implementation of MultiBoxTensor for object detection predictions. + +This module provides a PyTorch tensor subclass for multi-box detection predictions +with a unified format. Due to metaclass conflicts with torch.Tensor, this class +cannot explicitly inherit from the StructuredPrediction protocol but implements its +interface via structural typing (duck typing). +""" + +from typing import Optional + +import torch + +from xplique.utils_functions.object_detection.base.multi_box_tensor import BaseMultiBoxTensor + + +class TorchMultiBoxTensor(torch.Tensor): + """ + Tensor representation for multiple bounding box predictions with class probabilities. + + This class extends torch.Tensor to represent object detection predictions with shape + (B, C) where B is the number of boxes and C encodes box coordinates, objectness score, + and class predictions. The encoding is: 4 coordinates + 1 objectness + nb_classes. + For example, (9, 85) represents 9 boxes with 80 classes (COCO dataset). + + Note: This class implements the MultiBoxTensor protocol via structural typing. + The class complies with the protocol by implementing: + - to_batched_tensor(): Adds batch dimension + - filter(class_id, confidence): Filters boxes by class/score + + Additional object detection methods: + - boxes(): Extract box coordinates + - scores(): Extract objectness scores + - probas(): Extract class probabilities + """ + + def __format__(self, format_spec: str) -> str: + """ + Format the tensor as a string. + + For scalar tensors, extracts and formats the single value. Otherwise uses + default tensor formatting. + + Parameters + ---------- + format_spec + Format specification string. + + Returns + ------- + formatted_str + Formatted string representation of the tensor. + """ + if self.numel() == 1: + scalar_value = self.item() + return format(scalar_value, format_spec) + return super().__format__(format_spec) + + def boxes(self) -> torch.Tensor: + """ + Extract box coordinates from the predictions. + + Returns + ------- + boxes + Tensor of shape (B, 4) containing box coordinates for each detection. + """ + return self[:, :4] + + def scores(self) -> torch.Tensor: + """ + Extract objectness scores from the predictions. + + Returns + ------- + scores + Tensor of shape (B,) containing confidence scores for each detection. + """ + return self[:, 4] + + def probas(self) -> torch.Tensor: + """ + Extract class probabilities from the predictions. + + Returns + ------- + probas + Tensor of shape (B, num_classes) containing class probabilities for each box. + """ + return self[:, 5:] + + def filter( + self, class_id: Optional[int] = None, confidence: Optional[float] = None + ) -> "TorchMultiBoxTensor": + """ + Filter detections by class ID and/or confidence threshold. + + Parameters + ---------- + class_id + If provided, keep only detections of this class. + confidence + If provided, keep only detections with score >= this threshold. + + Returns + ------- + filtered_tensor + Filtered MultiBoxTensor containing only detections matching the criteria. + """ + if class_id is None and confidence is None: + return self + probas = self.probas() + class_ids = probas.argmax(dim=-1) + scores = self.scores() + if class_id is not None and confidence is not None: + keep = (class_ids == class_id) & (scores >= confidence) + elif class_id is None: + keep = scores >= confidence + else: + keep = class_ids == class_id + return self[keep, :] + + def to_attribution_target(self, class_id=None): + """ + Return self as the attribution target. + + For object detection, filter() has already selected the relevant boxes; + the filtered box tensor (coordinates + scores) is directly the target + for the attribution method. ``class_id`` is ignored here. + """ + return self + + def to_batched_tensor(self) -> torch.Tensor: + """ + Add batch dimension to MultiBoxTensor. + + Converts (num_boxes, features) -> (1, num_boxes, features) + + Returns + ------- + batched_tensor + torch.Tensor with batch dimension added + """ + return torch.unsqueeze(self, dim=0) + + +# Verify structural compliance with BaseMultiBoxTensor protocol at import time. +# TorchMultiBoxTensor cannot explicitly inherit from BaseMultiBoxTensor due to a +# metaclass conflict between torch.Tensor (_TensorMeta) and Protocol (_ProtocolMeta). +assert issubclass(TorchMultiBoxTensor, BaseMultiBoxTensor), ( + "TorchMultiBoxTensor must structurally satisfy the BaseMultiBoxTensor protocol" +) diff --git a/xplique/utils_functions/output_as_list_mixin.py b/xplique/utils_functions/output_as_list_mixin.py new file mode 100644 index 00000000..6be709d3 --- /dev/null +++ b/xplique/utils_functions/output_as_list_mixin.py @@ -0,0 +1,36 @@ +"""Mixin providing output_as_list and output_as_tensor as inverse boolean properties.""" + + +class OutputAsListMixin: + """ + Mixin that exposes the output format as two inverse boolean properties. + + ``output_as_list`` and ``output_as_tensor`` are backed by a single + ``_output_as_list`` field; setting either one automatically reflects in + the other. + """ + + @staticmethod + def _validate_bool(name: str, value: object) -> None: + if not isinstance(value, bool): + raise TypeError(f"{name} must be a bool, got {type(value).__name__!r}") + + @property + def output_as_list(self) -> bool: + """If True, outputs are returned as a list; if False, as a stacked tensor.""" + return self._output_as_list + + @output_as_list.setter + def output_as_list(self, value: bool) -> None: + self._validate_bool("output_as_list", value) + self._output_as_list = value + + @property + def output_as_tensor(self) -> bool: + """If True, outputs are returned as a stacked tensor; if False, as a list.""" + return not self._output_as_list + + @output_as_tensor.setter + def output_as_tensor(self, value: bool) -> None: + self._validate_bool("output_as_tensor", value) + self._output_as_list = not value diff --git a/xplique/wrappers/pytorch.py b/xplique/wrappers/pytorch.py index 72533186..f8fab78e 100644 --- a/xplique/wrappers/pytorch.py +++ b/xplique/wrappers/pytorch.py @@ -23,6 +23,9 @@ class TorchWrapper(tf.keras.Model): If we are on GPU or CPU is_channel_first A boolean that is true if the torch's model expect a channel dim and if this one come first + requires_grad + Whether gradients are available for the wrapped model. If False, inference + remains available but gradient-based explainers raise a RuntimeError. """ def __init__( @@ -30,6 +33,7 @@ def __init__( torch_model: "nn.Module", # noqa: F821 device: Union["torch.device", str], # noqa: F821 is_channel_first: Optional[bool] = None, + requires_grad: bool = True, ): # pylint: disable=C0415,C0103,W0719 try: super().__init__() @@ -61,6 +65,7 @@ def __init__( self.channel_first = self._has_conv_layers() else: self.channel_first = is_channel_first + self.requires_grad = requires_grad # deactivate all tf.function tf.config.run_functions_eagerly(True) warnings.warn( @@ -90,22 +95,31 @@ def call(self, inputs: np.ndarray) -> Tuple[tf.Tensor, Callable]: The function that allow to compute the gradient of the PyTorch model and broadcast it for Tensorflow """ - # transform your numpy inputs to torch - torch_inputs = self.np_img_to_torch(inputs).to(self.device) - torch_inputs.requires_grad_(True) - - # make predictions - self.model.zero_grad() - outputs = self.model(torch_inputs) + # Transform inputs before enabling gradients so they remain leaf tensors. + torch_inputs = self.np_img_to_torch(inputs).to(self.device).detach() + if self.requires_grad: + with self.torch.enable_grad(): + torch_inputs.requires_grad_(True) + outputs = self.model(torch_inputs) + else: + with self.torch.no_grad(): + outputs = self.model(torch_inputs) output_tensor = tf.constant(outputs.cpu().detach().numpy()) def grad(upstream): - self.torch.autograd.backward( + if not self.requires_grad: + raise RuntimeError( + "TorchWrapper was created with requires_grad=False and cannot be used " + "with gradient-based explainers." + ) + (dx_torch,) = self.torch.autograd.grad( outputs, - grad_tensors=self.from_numpy(upstream.numpy()).to(self.device), - retain_graph=False, + torch_inputs, + grad_outputs=self.from_numpy(upstream.numpy()).to(self.device), + allow_unused=True, ) - dx_torch = torch_inputs.grad + if dx_torch is None: + return None dx_np = dx_torch.cpu().detach().numpy() if self.channel_first: