Craft object detection

[fredericboisnard]

CRAFT for Object Detection

This PR brings support of CRAFT for Object Detection.

Object detection models often have very different architectures and output their results in various formats, which frequently do not allow gradients to flow through. To 'normalize' these outputs and make them compatible with Xplique, several formatters are introduced in this PR. These formatters can reorganize or reformat the outputs and be integrated into different parts of the model pipeline to expose gradients.

So, this PR is composed of several patches that can be organized into the following parts:

• BoxFormatters -> building blocks for representing the models results in a unified way
• BoxModelWrappers -> wrap models with a formatter, allow diplaying results and use of Xplique attribution methods
• CRAFT itself, with LatentData and LatentExtractor -> use all of the above + other concepts to bring support of CRAFT for object detection models

Box Formatters

Those formatters are used to:

• standardize predictions: some models output results as list of predictions in their own data format, some others as a single dictionary containing the predictions ; the BoxFormattes now always return a list of type MultiBoxTensor
• translate coordinate systems
• preserve gradients (pytorch & tf implementations)

MultiBoxTensor

This class gathers the results of the model for 1 image. It offers a convenient way to filter those results by class_id or accuracy.

flowchart LR
  Images --> Model -->BoxFormatter --> |MultiBoxTensor list| Predictions

Loading

Supported Box Formats

• XYXY: (x_min, y_min, x_max, y_max) - corner coordinates format
• CXCYWH: (center_x, center_y, width, height) - center-based format
• XYWH: (x_min, y_min, width, height) - top-left corner format

Supported Models and Formatters

Several formatters are delivered, they can be used for most common models.

PyTorch Models:

• DETR (DetrBoxFormatter): Processes transformer outputs with CXCYWH normalized format
• FCOS, RetinaNet, SSD (TorchvisionBoxFormatter): Works with XYXY normalized format from torchvision models
• YOLO (YoloResultBoxFormatter): Handles modern YOLO v8+ models that return Results objects containing boxes, scores, and classes in XYXY format
• YOLO (YoloRawBoxFormatter): Handles CXCYWH normalized format with differentiable head modifications for gradient flow

TensorFlow Models:

• RetinaNet: Dedicated TensorFlow support with corresponding formatters

BoxesModelWrappers

Those wrappers are shortcuts to wrap a model with its dedicated formatter.
Depending on where the formatter is plugged, gradients can or can not be propagated.

• Pytorch: Detr, Retinanet, Fcos, SSD, YOLO
• Tf: Retinanet

flowchart LR
  subgraph BoxesModelWrapper["BoxesModelWrapper"]
    direction LR
    Model --> BoxFormatter
  end
  Images --> BoxesModelWrapper -->|MultiBoxTensor list| Predictions

Loading

It is also possible to set the BoxesModelWrapper as "tensor mode output", to make them compatible with Xplique attribution framework (in this case the process is done image by image).

Visualization helpers

Some methods have been added to display the outputs in a common way for all the models (display_image_with_boxes)
They can also display boxes with heatmaps:

CRAFT

To use CRAFT, we must split the model 2 in parts. Problem with the object detection models is they are often not 'splittable' easily because they are not composed of a succession of layers.
A new class called LatentData is introduced to handle the activations on 1 side, and the remaining data that are not activations but need to be transmitted to the next part of the model (we need this to reencode a concept to compare their relative importance / attributions).
The LatentExtractor decomposes object detection models into two functions: g (input -> latent data) and h (latent data -> predictions). This splitting enables concept extraction from intermediate representations.

flowchart LR
inputs --> g(["g \(input to latent\)"]) --> LatentData --> h(["h \(latent to logit\)"]) --> Formatter([Formatter]) --> MultiBoxTensor

Loading

The LatentData acts as a buffer between the 2 parts of the split model, allowing to store the full context of execution in the middle of the model, therefore each model has a different LatentData implementation:

• Detr: LatentDataDetr stores positional encoding & mask information for the transformer
• YOLO: LatentDataYolo stores skip connection outputs from earlier layers
• RetinaNet: LatentDataRetinaNet stores image shape for bounding box coordinate conversion
• etc...

Several builders are available to produce such extractors for the most common models:

• Pytorch: Detr, YOLO (needs a YOLO specially trained with ReLU in order to be able to fit the NMF), Retinaner, SSD
• Tf: Retinanet

HolisticCraftObjectDetection applies CRAFT to object detection models by extracting concepts from the latent space (LatentData) produced by the LatentExtractor.

Visualization Methods:

• display_images_per_concept(): Shows spatial heatmaps of concept activations on images
• display_top_images_per_concept(): Displays images with highest activation for each concept

CRAFT Fit

The CRAFT fitting operation gathers the activations resulting of the 1st part of the model applied on the input dataset, thanks to the classes described above (LatentExtractor & LatenData).

flowchart LR
    User([User])
    User -->|1 fit| UC1Input

    subgraph UC1["1 CRAFT fit"]
        direction LR
        UC1Input[Training Images] --> UC1Fit([fit])
        UC1Fit --> UC1Bank[W]
    end
    
    style UC1 fill:#e6f3ff

Loading

flowchart LR
    User([User])

    User -->|2 Images to Concepts| UC2Input
    subgraph UC2["2 CRAFT transform"]
        direction LR
        UC2Input[New Images] --> UC2Transform([transform])
        UC2Transform --> UC2Encode([encode])
        UC2Encode --> UC2LatentExtractor([LatentExtractor.input_to_latent])
        UC2LatentExtractor --> UC2LatentData[LatentData]
        UC2LatentData --> UC2TransformLatent([transform_latent])
        UC2TransformLatent --> UC2EncodedData["EncodedData<br> \(LatentData + U\)"]
        UC2EncodedData --> UC2Coeffs[U]
    end

    style UC2 fill:#e6ffe6

Loading

Importance evaluation: the decode path

• GradientInput and Sobol attributions are supported to evaluate the importance of the concepts (GradientInput needs the model to be differenciable).
• More generally all attributions methods provided by Xplique should work.

flowchart LR
    User([User])

    User -->|3 Concepts ranking| UC3Input
    subgraph UC3["3 CRAFT.estimate_importance"]
        direction LR
        UC3Explainer[Explainer: GradientInput, Sobol] --> compute_explanation_per_concept
        UC3Operator[operator] --> compute_explanation_per_concept
        UC3Input[Images] --> compute_explanation_per_concept
        compute_explanation_per_concept --> UC3LatentExtractor
        subgraph UC3LatentExtractor["LatentExtractor"]
            direction LR
            UC3Encode([encode<br>images]) --> EncodedData -->|LatentData, U| UC3Decode([decode]) -->|MultiBoxTensor| filtered_targets["filter targets<br>class_id, accuracy"] --> explain
            EncodedData -->|LatentData| UC3MakeDec(["make_concept_decoder<br>\(TorchWrapper\)"]) -->|concept_decoder| Exp([explainer.model = concept_decoder]) --> explain["explainer.explain(U, targets)"]
            explain --> explanation
        end
    end

    style UC3 fill:#fff4e6

Loading

Feature Visualization: the encode differentiable path

To do this, use the encode() method with argument differentiable=True, it will use a specific path on torch/tf framework (using torch.linalg.lstsq() for example on torch side) instead of the regular NMF operation proficed by skleanrn package.

Input (requires_grad) -> g (encoder) -> latent_activations -> lstsq -> concept_coeffs

It can be used with an external toolbox such as horama to maximize a concept activation. See the notebooks for CRAFT application for Detr and YOLO.

flowchart LR
    User([User])
    
    User -->|4 Feature Visualization| UC4Input
    
    subgraph UC4["4 Visualize Concepts"]
        direction LR
        UC4Input(["make objective"]) --> objective
        subgraph objective
            direction LR
            images --> encode(["encode"]) -->|EncodeData| U --> mean["mean on concept id"]
        end
        mean --> maco([maco])
    end
    style UC4 fill:#ffe6f3

Loading

Support of classifiers

Classifiers are also supported by this Holistic CRAFT.

Support of external NMF (Overcomplete)

External factorizers are also supported, through a protocol "ConceptFactorizer". In particular Overcomplete NMF, SemiNMF and ConvexNMF are compatible.

Misc

In the last patch, I added several helper scripts in a new examples folder: we might in the end only choose 1 or 2 to keep.

Finally, maybe I should move each model specific code (in particular each LatentData) outside of the code of Xplique, for example in dedicated notebook for each model.

[Agustin-Picard]

Thank you very much for this PR, there is a ton of work that has been put into creating a lot of new features! The architecture that was implemented allows for a nice generalization across different object detection architectures and frameworks, with plenty of unit tests to ensure non-regression and that everything is working well, and the new functionality seems easy to use by new users. Documentation, the new plotting functions, and backward compatibility with the original Craft method are really nice to have, too.

However, I spotted some potential issues that might need fixing before merging. Note: do not hesitate to argue if I'm wrong here. This PR is quite large, and I might have missed stuff:

• The most important one is the encode method in the Torch and TF NMF Factorizers. These factorizers apply a non-linear transformation that can be implemented through specific solvers (like they do in scikit-learn and overcomplete). However, linalg.lstsq is being used, which solves a linear least-squares problem. I think we should consider whether we should try to find another library that already provides the NMF transform method, or implement it ourselves. In practice, we're computing a SemiNMF but with the non-negative constraint on the concept bank (which is the opposite of the SemiNMF in overcomplete, for example).

• On the TorchBoxedModelWrapper, you override the __call__ method of a nn.Module. I'm worried this might affect torch hooks, since that is where they are implemented.

• I think we should also add some tests on the new Factorizer classes, to make sure that they are working as intended (ie, NMF gives non-negative banks and coefficients, reconstruction loss is low, etc.).

• Finally, for the holistic Craft tests, you download images from the internet. I'm concerned that network issues might cause unstable unit tests.

I've left plenty of comments and questions. I hope I didn't miss anything!

[fredericboisnard]

I pushed a new PR to take into account Agu's reviews, except the new FSITA solver (I am currently merging it)

[Agustin-Picard]

Awesome work, looks good to me!