Initial commit

.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,3 @@
+
+/venv
+__pycache__

LeGrad/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 Walid Bousselham, Angie Boggust, Sofian Chaybouti,Hendrik Strobelt Hilde Kuehne.
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

LeGrad/legrad/__init__.py

@@ -0,0 +1,2 @@
+from .wrapper import LeWrapper, LePreprocess
+from .utils import *

LeGrad/legrad/utils.py

@@ -0,0 +1,722 @@
+from pathlib import Path
+from typing import Optional, List, Tuple
+import matplotlib.pyplot as plt
+import numpy as np
+from PIL import Image
+import cv2 as cv2
+import warnings
+
+import torch
+from torch import Tensor
+from torch.nn import functional as F
+
+import open_clip
+from open_clip import OPENAI_DATASET_MEAN, OPENAI_DATASET_STD
+from open_clip.transformer import _expand_token
+from timm.layers import resample_abs_pos_embed
+
+
+################################################################################
+#                               Hooks utils                                    #
+################################################################################
+
+
+# ------------ Hooked Multi-Head Attention ------------
+# from https://github.com/mlfoundations/open_clip/blob/73fa7f03a33da53653f61841eb6d69aef161e521/src/open_clip/transformer.py#L129
+def hooked_attention_forward(
+    self,
+    x,
+    x_k,
+    x_v,
+    attn_mask: Optional[torch.Tensor] = None,
+    need_weights: bool = False,
+):
+    L, N, C = x.shape
+    q, k, v = F.linear(x, self.in_proj_weight, self.in_proj_bias).chunk(3, dim=-1)
+    q = q.contiguous().view(L, N * self.num_heads, -1).transpose(0, 1)
+    k = k.contiguous().view(L, N * self.num_heads, -1).transpose(0, 1)
+    v = v.contiguous().view(L, N * self.num_heads, -1).transpose(0, 1)
+
+    head_dim = q.shape[-1]
+    scale = float(head_dim) ** -0.5
+    q = q * scale
+    attn = torch.bmm(q, k.transpose(-1, -2))
+
+    if attn_mask is not None:
+        if attn_mask.dtype == torch.bool:
+            new_attn_mask = torch.zeros_like(attn_mask, dtype=q.dtype)
+            new_attn_mask.masked_fill_(attn_mask, float("-inf"))
+            attn_mask = new_attn_mask
+        attn += attn_mask
+
+    attn = attn.softmax(dim=-1)
+    # Hook for attention maps
+    self.attention_map = attn
+
+    x = torch.bmm(attn, v)
+    x = x.transpose(0, 1).reshape(L, N, C)
+    x = self.out_proj(x)
+    return x
+
+
+def hooked_attention_timm_forward(self, x, attn_mask=None):
+    B, N, C = x.shape
+    qkv = (
+        self.qkv(x)
+        .reshape(B, N, 3, self.num_heads, self.head_dim)
+        .permute(2, 0, 3, 1, 4)
+    )
+    q, k, v = qkv.unbind(0)
+    q, k = self.q_norm(q), self.k_norm(k)
+
+    q = q * self.scale
+    attn = q @ k.transpose(-2, -1)
+    attn = attn.softmax(dim=-1)
+    attn = self.attn_drop(attn)
+    x = attn @ v
+
+    # Hook to save attention map for explainability
+    self.attention_map = attn
+
+    x = x.transpose(1, 2).reshape(B, N, C)
+    x = self.proj(x)
+    x = self.proj_drop(x)
+    return x
+
+
+# ------------ Hooked Residual Transformer Block ------------
+# from https://github.com/mlfoundations/open_clip/blob/73fa7f03a33da53653f61841eb6d69aef161e521/src/open_clip/transformer.py#L231
+def hooked_resblock_forward(self, q_x, k_x=None, v_x=None, attn_mask=None):
+    assert k_x is None and v_x is None, "k_x and v_x must be None"
+
+    # Modify this line to include the necessary arguments for hooked_attention_forward
+    x = q_x + self.ls1(
+        self.attn(
+            self.norm1(q_x),
+            k_x=k_x,
+            v_x=v_x,
+            attn_mask=attn_mask,
+        )
+    )
+    # Hook for intermediate features post Attn
+    self.feat_post_attn = x
+    x = x + self.ls2(self.mlp(self.norm2(x)))
+
+    # Hook for intermediate features post MLP
+    self.feat_post_mlp = x
+    return x
+
+
+# ------------ Hooked PyTorch's Multi-Head AttentionResidual ------------
+# modified from PyTorch Library
+# https://github.com/pytorch/pytorch/blob/8c8e4e31f2ddd8e59de18ac733c0c205c23d14ad/torch/nn/functional.py#L5178
+def hooked_torch_multi_head_attention_forward(
+    self, query, key, value, key_padding_mask=None, need_weights=True, attn_mask=None
+):
+    r"""
+    Args:
+        query, key, value: map a query and a set of key-value pairs to an output.
+            See "Attention Is All You Need" for more details.
+        key_padding_mask: if provided, specified padding elements in the key will
+            be ignored by the attention. When given a binary mask and a value is True,
+            the corresponding value on the attention layer will be ignored. When given
+            a byte mask and a value is non-zero, the corresponding value on the attention
+            layer will be ignored
+        need_weights: output attn_output_weights.
+        attn_mask: 2D or 3D mask that prevents attention to certain positions. A 2D mask will be broadcasted for all
+            the batches while a 3D mask allows to specify a different mask for the entries of each batch.
+
+    Shape:
+        - Inputs:
+        - query: :math:`(L, N, E)` where L is the target sequence length, N is the batch size, E is
+          the embedding dimension.
+        - key: :math:`(S, N, E)`, where S is the source sequence length, N is the batch size, E is
+          the embedding dimension.
+        - value: :math:`(S, N, E)` where S is the source sequence length, N is the batch size, E is
+          the embedding dimension.
+        - key_padding_mask: :math:`(N, S)` where N is the batch size, S is the source sequence length.
+          If a ByteTensor is provided, the non-zero positions will be ignored while the position
+          with the zero positions will be unchanged. If a BoolTensor is provided, the positions with the
+          value of ``True`` will be ignored while the position with the value of ``False`` will be unchanged.
+        - attn_mask: 2D mask :math:`(L, S)` where L is the target sequence length, S is the source sequence length.
+          3D mask :math:`(N*num_heads, L, S)` where N is the batch size, L is the target sequence length,
+          S is the source sequence length. attn_mask ensure that position i is allowed to attend the unmasked
+          positions. If a ByteTensor is provided, the non-zero positions are not allowed to attend
+          while the zero positions will be unchanged. If a BoolTensor is provided, positions with ``True``
+          is not allowed to attend while ``False`` values will be unchanged. If a FloatTensor
+          is provided, it will be added to the attention weight.
+
+        - Outputs:
+        - attn_output: :math:`(L, N, E)` where L is the target sequence length, N is the batch size,
+          E is the embedding dimension.
+        - attn_output_weights: :math:`(N, L, S)` where N is the batch size,
+          L is the target sequence length, S is the source sequence length.
+    """
+    if not self._qkv_same_embed_dim:
+        out, _attn_maps = hooked_torch_func_multi_head_attention_forward(
+            query,
+            key,
+            value,
+            self.embed_dim,
+            self.num_heads,
+            self.in_proj_weight,
+            self.in_proj_bias,
+            self.bias_k,
+            self.bias_v,
+            self.add_zero_attn,
+            self.dropout,
+            self.out_proj.weight,
+            self.out_proj.bias,
+            training=self.training,
+            key_padding_mask=key_padding_mask,
+            need_weights=True,
+            attn_mask=attn_mask,
+            use_separate_proj_weight=True,
+            q_proj_weight=self.q_proj_weight,
+            k_proj_weight=self.k_proj_weight,
+            v_proj_weight=self.v_proj_weight,
+        )
+        # Hook for attention maps
+        self.attention_maps = _attn_maps
+        return out, _attn_maps
+    else:
+        out, _attn_maps = hooked_torch_func_multi_head_attention_forward(
+            query,
+            key,
+            value,
+            self.embed_dim,
+            self.num_heads,
+            self.in_proj_weight,
+            self.in_proj_bias,
+            self.bias_k,
+            self.bias_v,
+            self.add_zero_attn,
+            self.dropout,
+            self.out_proj.weight,
+            self.out_proj.bias,
+            training=self.training,
+            key_padding_mask=key_padding_mask,
+            need_weights=True,
+            attn_mask=attn_mask,
+        )
+        # Hook for attention maps
+        self.attention_maps = _attn_maps
+        return out, _attn_maps
+
+
+def hooked_torch_func_multi_head_attention_forward(
+    query: Tensor,
+    key: Tensor,
+    value: Tensor,
+    embed_dim_to_check: int,
+    num_heads: int,
+    in_proj_weight: Tensor,
+    in_proj_bias: Tensor,
+    bias_k: Optional[Tensor],
+    bias_v: Optional[Tensor],
+    add_zero_attn: bool,
+    dropout_p: float,
+    out_proj_weight: Tensor,
+    out_proj_bias: Tensor,
+    training: bool = True,
+    key_padding_mask: Optional[Tensor] = None,
+    need_weights: bool = True,
+    attn_mask: Optional[Tensor] = None,
+    use_separate_proj_weight: bool = False,
+    q_proj_weight: Optional[Tensor] = None,
+    k_proj_weight: Optional[Tensor] = None,
+    v_proj_weight: Optional[Tensor] = None,
+    static_k: Optional[Tensor] = None,
+    static_v: Optional[Tensor] = None,
+) -> Tuple[Tensor, Optional[Tensor]]:
+    if not torch.jit.is_scripting():
+        tens_ops = (
+            query,
+            key,
+            value,
+            in_proj_weight,
+            in_proj_bias,
+            bias_k,
+            bias_v,
+            out_proj_weight,
+            out_proj_bias,
+        )
+        if any([type(t) is not Tensor for t in tens_ops]) and F.has_torch_function(
+            tens_ops
+        ):
+            return F.handle_torch_function(
+                multi_head_attention_forward,
+                tens_ops,
+                query,
+                key,
+                value,
+                embed_dim_to_check,
+                num_heads,
+                in_proj_weight,
+                in_proj_bias,
+                bias_k,
+                bias_v,
+                add_zero_attn,
+                dropout_p,
+                out_proj_weight,
+                out_proj_bias,
+                training=training,
+                key_padding_mask=key_padding_mask,
+                need_weights=need_weights,
+                attn_mask=attn_mask,
+                use_separate_proj_weight=use_separate_proj_weight,
+                q_proj_weight=q_proj_weight,
+                k_proj_weight=k_proj_weight,
+                v_proj_weight=v_proj_weight,
+                static_k=static_k,
+                static_v=static_v,
+            )
+    tgt_len, bsz, embed_dim = query.size()
+    assert embed_dim == embed_dim_to_check
+    # allow MHA to have different sizes for the feature dimension
+    assert key.size(0) == value.size(0) and key.size(1) == value.size(1)
+
+    head_dim = embed_dim // num_heads
+    assert head_dim * num_heads == embed_dim, "embed_dim must be divisible by num_heads"
+    scaling = float(head_dim) ** -0.5
+
+    if not use_separate_proj_weight:
+        if torch.equal(query, key) and torch.equal(key, value):
+            # self-attention
+            q, k, v = F.linear(query, in_proj_weight, in_proj_bias).chunk(3, dim=-1)
+
+        elif torch.equal(key, value):
+            # encoder-decoder attention
+            # This is inline in_proj function with in_proj_weight and in_proj_bias
+            _b = in_proj_bias
+            _start = 0
+            _end = embed_dim
+            _w = in_proj_weight[_start:_end, :]
+            if _b is not None:
+                _b = _b[_start:_end]
+            q = F.linear(query, _w, _b)
+
+            if key is None:
+                assert value is None
+                k = None
+                v = None
+            else:
+
+                # This is inline in_proj function with in_proj_weight and in_proj_bias
+                _b = in_proj_bias
+                _start = embed_dim
+                _end = None
+                _w = in_proj_weight[_start:, :]
+                if _b is not None:
+                    _b = _b[_start:]
+                k, v = F.linear(key, _w, _b).chunk(2, dim=-1)
+
+        else:
+            # This is inline in_proj function with in_proj_weight and in_proj_bias
+            _b = in_proj_bias
+            _start = 0
+            _end = embed_dim
+            _w = in_proj_weight[_start:_end, :]
+            if _b is not None:
+                _b = _b[_start:_end]
+            q = F.linear(query, _w, _b)
+
+            # This is inline in_proj function with in_proj_weight and in_proj_bias
+            _b = in_proj_bias
+            _start = embed_dim
+            _end = embed_dim * 2
+            _w = in_proj_weight[_start:_end, :]
+            if _b is not None:
+                _b = _b[_start:_end]
+            k = F.linear(key, _w, _b)
+
+            # This is inline in_proj function with in_proj_weight and in_proj_bias
+            _b = in_proj_bias
+            _start = embed_dim * 2
+            _end = None
+            _w = in_proj_weight[_start:, :]
+            if _b is not None:
+                _b = _b[_start:]
+            v = F.linear(value, _w, _b)
+    else:
+        q_proj_weight_non_opt = torch.jit._unwrap_optional(q_proj_weight)
+        len1, len2 = q_proj_weight_non_opt.size()
+        assert len1 == embed_dim and len2 == query.size(-1)
+
+        k_proj_weight_non_opt = torch.jit._unwrap_optional(k_proj_weight)
+        len1, len2 = k_proj_weight_non_opt.size()
+        assert len1 == embed_dim and len2 == key.size(-1)
+
+        v_proj_weight_non_opt = torch.jit._unwrap_optional(v_proj_weight)
+        len1, len2 = v_proj_weight_non_opt.size()
+        assert len1 == embed_dim and len2 == value.size(-1)
+
+        if in_proj_bias is not None:
+            q = F.linear(query, q_proj_weight_non_opt, in_proj_bias[0:embed_dim])
+            k = F.linear(
+                key, k_proj_weight_non_opt, in_proj_bias[embed_dim : (embed_dim * 2)]
+            )
+            v = F.linear(value, v_proj_weight_non_opt, in_proj_bias[(embed_dim * 2) :])
+        else:
+            q = F.linear(query, q_proj_weight_non_opt, in_proj_bias)
+            k = F.linear(key, k_proj_weight_non_opt, in_proj_bias)
+            v = F.linear(value, v_proj_weight_non_opt, in_proj_bias)
+    q = q * scaling
+
+    if attn_mask is not None:
+        assert (
+            attn_mask.dtype == torch.float32
+            or attn_mask.dtype == torch.float64
+            or attn_mask.dtype == torch.float16
+            or attn_mask.dtype == torch.uint8
+            or attn_mask.dtype == torch.bool
+        ), "Only float, byte, and bool types are supported for attn_mask, not {}".format(
+            attn_mask.dtype
+        )
+        if attn_mask.dtype == torch.uint8:
+            warnings.warn(
+                "Byte tensor for attn_mask in nn.MultiheadAttention is deprecated. Use bool tensor instead."
+            )
+            attn_mask = attn_mask.to(torch.bool)
+
+        if attn_mask.dim() == 2:
+            attn_mask = attn_mask.unsqueeze(0)
+            if list(attn_mask.size()) != [1, query.size(0), key.size(0)]:
+                raise RuntimeError("The size of the 2D attn_mask is not correct.")
+        elif attn_mask.dim() == 3:
+            if list(attn_mask.size()) != [bsz * num_heads, query.size(0), key.size(0)]:
+                raise RuntimeError("The size of the 3D attn_mask is not correct.")
+        else:
+            raise RuntimeError(
+                "attn_mask's dimension {} is not supported".format(attn_mask.dim())
+            )
+        # attn_mask's dim is 3 now.
+
+    # convert ByteTensor key_padding_mask to bool
+    if key_padding_mask is not None and key_padding_mask.dtype == torch.uint8:
+        warnings.warn(
+            "Byte tensor for key_padding_mask in nn.MultiheadAttention is deprecated. Use bool tensor instead."
+        )
+        key_padding_mask = key_padding_mask.to(torch.bool)
+
+    if bias_k is not None and bias_v is not None:
+        if static_k is None and static_v is None:
+            k = torch.cat([k, bias_k.repeat(1, bsz, 1)])
+            v = torch.cat([v, bias_v.repeat(1, bsz, 1)])
+            if attn_mask is not None:
+                attn_mask = pad(attn_mask, (0, 1))
+            if key_padding_mask is not None:
+                key_padding_mask = pad(key_padding_mask, (0, 1))
+        else:
+            assert static_k is None, "bias cannot be added to static key."
+            assert static_v is None, "bias cannot be added to static value."
+    else:
+        assert bias_k is None
+        assert bias_v is None
+
+    q = q.contiguous().view(tgt_len, bsz * num_heads, head_dim).transpose(0, 1)
+    if k is not None:
+        k = k.contiguous().view(-1, bsz * num_heads, head_dim).transpose(0, 1)
+    if v is not None:
+        v = v.contiguous().view(-1, bsz * num_heads, head_dim).transpose(0, 1)
+
+    if static_k is not None:
+        assert static_k.size(0) == bsz * num_heads
+        assert static_k.size(2) == head_dim
+        k = static_k
+
+    if static_v is not None:
+        assert static_v.size(0) == bsz * num_heads
+        assert static_v.size(2) == head_dim
+        v = static_v
+
+    src_len = k.size(1)
+
+    if key_padding_mask is not None:
+        assert key_padding_mask.size(0) == bsz
+        assert key_padding_mask.size(1) == src_len
+
+    if add_zero_attn:
+        src_len += 1
+        k = torch.cat(
+            [
+                k,
+                torch.zeros(
+                    (k.size(0), 1) + k.size()[2:], dtype=k.dtype, device=k.device
+                ),
+            ],
+            dim=1,
+        )
+        v = torch.cat(
+            [
+                v,
+                torch.zeros(
+                    (v.size(0), 1) + v.size()[2:], dtype=v.dtype, device=v.device
+                ),
+            ],
+            dim=1,
+        )
+        if attn_mask is not None:
+            attn_mask = pad(attn_mask, (0, 1))
+        if key_padding_mask is not None:
+            key_padding_mask = pad(key_padding_mask, (0, 1))
+
+    attn_output_weights = torch.bmm(q, k.transpose(1, 2))
+    assert list(attn_output_weights.size()) == [bsz * num_heads, tgt_len, src_len]
+
+    if attn_mask is not None:
+        if attn_mask.dtype == torch.bool:
+            attn_output_weights.masked_fill_(attn_mask, float("-inf"))
+        else:
+            attn_output_weights += attn_mask
+
+    if key_padding_mask is not None:
+        attn_output_weights = attn_output_weights.view(bsz, num_heads, tgt_len, src_len)
+        attn_output_weights = attn_output_weights.masked_fill(
+            key_padding_mask.unsqueeze(1).unsqueeze(2),
+            float("-inf"),
+        )
+        attn_output_weights = attn_output_weights.view(
+            bsz * num_heads, tgt_len, src_len
+        )
+
+    attn_output_weights = F.softmax(attn_output_weights, dim=-1)
+    attn_output_weights = F.dropout(attn_output_weights, p=dropout_p, training=training)
+
+    # # use hooks for the attention weights if necessary
+    # self.attention_map = attn_output_weights
+    # # if attention_probs_forward_hook is not None and attention_probs_backwards_hook is not None:
+    # if attention_probs_forward_hook is not None:
+    #     attention_probs_forward_hook(attn_output_weights)
+    #     # attn_output_weights.register_hook(attention_probs_backwards_hook)
+
+    attn_output = torch.bmm(attn_output_weights, v)
+    assert list(attn_output.size()) == [bsz * num_heads, tgt_len, head_dim]
+    attn_output = attn_output.transpose(0, 1).contiguous().view(tgt_len, bsz, embed_dim)
+    attn_output = F.linear(attn_output, out_proj_weight, out_proj_bias)
+
+    if need_weights:
+        # --- Fix: removed the unnecessary average over heads, Why?
+        # average attention weights over heads
+        # attn_output_weights = attn_output_weights.view(bsz, num_heads, tgt_len, src_len)
+        # return attn_output, attn_output_weights.sum(dim=1) / num_heads
+        return attn_output, attn_output_weights
+    else:
+        return attn_output, None
+
+
+# ------------ Hooked TimmModel's Residual Transformer Block ------------
+def hooked_resblock_timm_forward(self, x: torch.Tensor) -> torch.Tensor:
+    x = x + self.drop_path1(self.ls1(self.attn(self.norm1(x))))
+    self.feat_post_attn = x
+    x = x + self.drop_path2(self.ls2(self.mlp(self.norm2(x))))
+    self.feat_post_mlp = x
+    return x
+
+
+# ------------ Hooked TimmModel's Attentional Pooler ------------
+def hooked_attentional_pooler_timm_forward(self, x):
+    B, N, C = x.shape
+
+    if self.pos_embed is not None:
+        # FIXME interpolate
+        x = x + self.pos_embed.unsqueeze(0).to(x.dtype)
+
+    q_latent = self.latent.expand(B, -1, -1)
+    q = (
+        self.q(q_latent)
+        .reshape(B, self.latent_len, self.num_heads, self.head_dim)
+        .transpose(1, 2)
+    )
+
+    kv = (
+        self.kv(x)
+        .reshape(B, N, 2, self.num_heads, self.head_dim)
+        .permute(2, 0, 3, 1, 4)
+    )
+    k, v = kv.unbind(0)
+
+    q, k = self.q_norm(q), self.k_norm(k)
+
+    q = q * self.scale
+    attn = q @ k.transpose(-2, -1)
+    attn = attn.softmax(dim=-1)
+    x = attn @ v
+
+    # Hook to save attention map for explainability
+    self.attn_probs = attn
+
+    x = x.transpose(1, 2).reshape(B, self.latent_len, C)
+    x = self.proj(x)
+    x = self.proj_drop(x)
+
+    x = x + self.mlp(self.norm(x))
+
+    # optional pool if latent seq_len > 1 and pooled output is desired
+    if self.pool == "token":
+        x = x[:, 0]
+    elif self.pool == "avg":
+        x = x.mean(1)
+    return x
+
+
+# ------------ OpenCLIP ViT forward with dynamic size ------------
+def vit_dynamic_size_forward(self, x: torch.Tensor):
+    x = self.conv1(x)  # shape = [*, width, grid, grid]
+    grid_h, grid_w = x.shape[2:]
+    x = x.reshape(x.shape[0], x.shape[1], -1)  # shape = [*, width, grid ** 2]
+    x = x.permute(0, 2, 1)  # shape = [*, grid ** 2, width]
+
+    # class embeddings and positional embeddings
+    x = torch.cat(
+        [_expand_token(self.class_embedding, x.shape[0]).to(x.dtype), x], dim=1
+    )
+    # shape = [*, grid ** 2 + 1, width]
+    if x.shape[1] != self.positional_embedding.shape[1]:
+        self.positional_embedding.data = resample_abs_pos_embed(
+            self.positional_embedding.unsqueeze(0),
+            new_size=[grid_h, grid_w],
+            # old_size=list(self.grid_size),
+            num_prefix_tokens=1,
+            interpolation="bicubic",
+            antialias=True,
+        )
+
+    x = x + self.positional_embedding.to(x.dtype)
+
+    x = self.patch_dropout(x)
+    x = self.ln_pre(x)
+
+    x = x.permute(1, 0, 2)  # NLD -> LND
+    x = self.transformer(x)
+    x = x.permute(1, 0, 2)  # LND -> NLD
+
+    if self.attn_pool is not None:
+        if self.attn_pool_contrastive is not None:
+            # This is untested, WIP pooling that should match paper
+            x = self.ln_post(x)  # TBD LN first or separate one after each pool?
+            tokens = self.attn_pool(x)
+            if self.attn_pool_type == "parallel":
+                pooled = self.attn_pool_contrastive(x)
+            else:
+                assert self.attn_pool_type == "cascade"
+                pooled = self.attn_pool_contrastive(tokens)
+        else:
+            # this is the original OpenCLIP CoCa setup, does not match paper
+            x = self.attn_pool(x)
+            x = self.ln_post(x)
+            pooled, tokens = self._global_pool(x)
+    elif self.final_ln_after_pool:
+        pooled, tokens = self._global_pool(x)
+        pooled = self.ln_post(pooled)
+    else:
+        x = self.ln_post(x)
+        pooled, tokens = self._global_pool(x)
+
+    if self.proj is not None:
+        pooled = pooled @ self.proj
+
+    if self.output_tokens:
+        return pooled, tokens
+
+    return pooled
+
+
+################################################################################
+#                               Visualization utils                            #
+################################################################################
+
+
+def min_max(logits):
+    B, num_prompt = logits.shape[:2]
+    logits_min = (
+        logits.reshape(B, num_prompt, -1).min(dim=-1, keepdim=True)[0].unsqueeze(-1)
+    )
+    logits_max = (
+        logits.reshape(B, num_prompt, -1).max(dim=-1, keepdim=True)[0].unsqueeze(-1)
+    )
+    logits = (logits - logits_min) / (logits_max - logits_min)
+    return logits
+
+
+def visualize(image, heatmaps, alpha=0.6, save_path: Path = None):
+    # heatmaps of shape (N, 1, W, H)
+    W, H = heatmaps.shape[-2:]
+    if isinstance(image, Image.Image):
+        image = image.resize((W, H))
+    elif isinstance(image, torch.Tensor):
+        if image.ndim > 3:
+            image = image.squeeze(0)
+        # undo the normalization
+        image_unormed = (
+            image.detach().cpu() * torch.Tensor(OPENAI_DATASET_STD)[:, None, None]
+        ) + torch.Tensor(OPENAI_DATASET_MEAN)[:, None, None]
+        # convert to PIL
+        image = Image.fromarray(
+            (image_unormed.permute(1, 2, 0).numpy() * 255).astype("uint8")
+        )
+    else:
+        raise f"image should be either of type PIL.Image.Image or torch.Tensor but found {type(image)}"
+
+    # plot image
+    plt.imshow(image)
+    plt.axis("off")
+    plt.tight_layout()
+    plt.show()
+
+    if heatmaps.ndim > 3:
+        heatmaps = heatmaps.squeeze(0)
+    heatmaps = heatmaps.detach().cpu().numpy()
+
+    img_cv = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
+    heatmaps = (heatmaps * 255).astype("uint8")
+    heat_maps = [cv2.applyColorMap(logit, cv2.COLORMAP_JET) for logit in heatmaps]
+
+    vizs = [(1 - alpha) * img_cv + alpha * heat_map for heat_map in heat_maps]
+    for i, viz in enumerate(vizs):
+        viz = cv2.cvtColor(viz.astype("uint8"), cv2.COLOR_BGR2RGB)
+        plt.imshow(viz)
+        plt.axis("off")
+        plt.tight_layout()
+        # remove the margin
+        plt.subplots_adjust(left=0, right=1, top=1, bottom=0)
+        plt.show()
+        if save_path is not None:
+            plt.savefig(save_path, bbox_inches="tight", pad_inches=0)
+            print(f"Saved visualization at {save_path}")
+
+
+def list_pretrained():
+    openclip_list_ = open_clip.list_pretrained()
+    filtered_list = [
+        (model_name, pretrained)
+        for (model_name, pretrained) in openclip_list_
+        if model_name
+    ]
+    unsupported_models = [
+        "RN",
+        "convnext",
+    ]  # legrad doesn't support CNN-based VLMs (for the moment)
+    _str = (
+        ": ".join(["model_name" + " " * (25 - len("model_name")), "pretrained"]) + "\n"
+    )  # for nice display
+    for model_name, pretrained in openclip_list_:
+        for unsup_model in unsupported_models:
+            if unsup_model in model_name:
+                skip = True
+                break
+            else:
+                skip = False
+        if not skip:
+            filtered_list.append((model_name, pretrained))
+            _str += (
+                ": ".join([model_name + " " * (25 - len(model_name)), pretrained])
+                + "\n"
+            )  # for nice display
+
+    print(_str)
+    return filtered_list
+
+
+if __name__ == "__main__":
+    list_pretrained()

LeGrad/legrad/wrapper.py

@@ -0,0 +1,447 @@
+import math
+import types
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torchvision.transforms import Compose, Resize, InterpolationMode
+import open_clip
+from open_clip.transformer import VisionTransformer
+from open_clip.timm_model import TimmModel
+from einops import rearrange
+
+from .utils import (
+    hooked_attention_timm_forward,
+    hooked_resblock_forward,
+    hooked_attention_forward,
+    hooked_resblock_timm_forward,
+    hooked_attentional_pooler_timm_forward,
+    vit_dynamic_size_forward,
+    min_max,
+    hooked_torch_multi_head_attention_forward,
+)
+
+
+class LeWrapper(nn.Module):
+    """
+    Wrapper around OpenCLIP to add LeGrad to OpenCLIP's model while keep all the functionalities of the original model.
+    """
+
+    def __init__(self, model, layer_index=-2):
+        super(LeWrapper, self).__init__()
+        # ------------ copy of model's attributes and methods ------------
+        for attr in dir(model):
+            if not attr.startswith("__"):
+                setattr(self, attr, getattr(model, attr))
+
+        # ------------ activate hooks & gradient ------------
+        self._activate_hooks(layer_index=layer_index)
+
+    def _activate_hooks(self, layer_index):
+        # ------------ identify model's type ------------
+        print("Activating necessary hooks and gradients ....")
+        if isinstance(self.visual, VisionTransformer):
+            # --- Activate dynamic image size ---
+            self.visual.forward = types.MethodType(
+                vit_dynamic_size_forward, self.visual
+            )
+            # Get patch size
+            self.patch_size = self.visual.patch_size[0]
+            # Get starting depth (in case of negative layer_index)
+            self.starting_depth = (
+                layer_index
+                if layer_index >= 0
+                else len(self.visual.transformer.resblocks) + layer_index
+            )
+
+            if self.visual.attn_pool is None:
+                self.model_type = "clip"
+                self._activate_self_attention_hooks()
+            else:
+                self.model_type = "coca"
+                self._activate_att_pool_hooks(layer_index=layer_index)
+
+        elif isinstance(self.visual, TimmModel):
+            # --- Activate dynamic image size ---
+            self.visual.trunk.dynamic_img_size = True
+            self.visual.trunk.patch_embed.dynamic_img_size = True
+            self.visual.trunk.patch_embed.strict_img_size = False
+            self.visual.trunk.patch_embed.flatten = False
+            self.visual.trunk.patch_embed.output_fmt = "NHWC"
+            self.model_type = "timm_siglip"
+            # --- Get patch size ---
+            self.patch_size = self.visual.trunk.patch_embed.patch_size[0]
+            # --- Get starting depth (in case of negative layer_index) ---
+            self.starting_depth = (
+                layer_index
+                if layer_index >= 0
+                else len(self.visual.trunk.blocks) + layer_index
+            )
+            if (
+                hasattr(self.visual.trunk, "attn_pool")
+                and self.visual.trunk.attn_pool is not None
+            ):
+                self._activate_timm_attn_pool_hooks(layer_index=layer_index)
+            else:
+                self._activate_timm_self_attention_hooks()
+        else:
+            raise ValueError(
+                "Model currently not supported, see legrad.list_pretrained() for a list of available models"
+            )
+        print("Hooks and gradients activated!")
+
+    def _activate_self_attention_hooks(self):
+        # Adjusting to use the correct structure
+        if isinstance(self.visual, VisionTransformer):
+            blocks = self.visual.transformer.resblocks
+        elif isinstance(self.visual, TimmModel):
+            blocks = self.visual.trunk.blocks
+        else:
+            raise ValueError("Unsupported model type for self-attention hooks")
+
+        # ---------- Apply Hooks + Activate/Deactivate gradients ----------
+        # Necessary steps to get intermediate representations
+        for name, param in self.named_parameters():
+            param.requires_grad = False
+            if name.startswith("visual.trunk.blocks"):
+                depth = int(name.split("visual.trunk.blocks.")[-1].split(".")[0])
+                if depth >= self.starting_depth:
+                    param.requires_grad = True
+
+        # --- Activate the hooks for the specific layers ---
+        for layer in range(self.starting_depth, len(blocks)):
+            blocks[layer].attn.forward = types.MethodType(
+                hooked_attention_forward, blocks[layer].attn
+            )
+            blocks[layer].forward = types.MethodType(
+                hooked_resblock_forward, blocks[layer]
+            )
+
+    def _activate_timm_self_attention_hooks(self):
+        # Adjusting to use the correct structure
+        blocks = self.visual.trunk.blocks
+
+        # ---------- Apply Hooks + Activate/Deactivate gradients ----------
+        # Necessary steps to get intermediate representations
+        for name, param in self.named_parameters():
+            param.requires_grad = False
+            if name.startswith("visual.trunk.blocks"):
+                depth = int(name.split("visual.trunk.blocks.")[-1].split(".")[0])
+                if depth >= self.starting_depth:
+                    param.requires_grad = True
+
+        # --- Activate the hooks for the specific layers ---
+        for layer in range(self.starting_depth, len(blocks)):
+            blocks[layer].attn.forward = types.MethodType(
+                hooked_attention_timm_forward, blocks[layer].attn
+            )
+            blocks[layer].forward = types.MethodType(
+                hooked_resblock_timm_forward, blocks[layer]
+            )
+
+    def _activate_att_pool_hooks(self, layer_index):
+        # ---------- Apply Hooks + Activate/Deactivate gradients ----------
+        # Necessary steps to get intermediate representations
+        for name, param in self.named_parameters():
+            param.requires_grad = False
+            if name.startswith("visual.transformer.resblocks"):
+                # get the depth
+                depth = int(
+                    name.split("visual.transformer.resblocks.")[-1].split(".")[0]
+                )
+                if depth >= self.starting_depth:
+                    param.requires_grad = True
+
+        # --- Activate the hooks for the specific layers ---
+        for layer in range(self.starting_depth, len(self.visual.transformer.resblocks)):
+            self.visual.transformer.resblocks[layer].forward = types.MethodType(
+                hooked_resblock_forward, self.visual.transformer.resblocks[layer]
+            )
+        # --- Apply hook on the attentional pooler ---
+        self.visual.attn_pool.attn.forward = types.MethodType(
+            hooked_torch_multi_head_attention_forward, self.visual.attn_pool.attn
+        )
+
+    def _activate_timm_attn_pool_hooks(self, layer_index):
+        # Ensure all components are present before attaching hooks
+        if (
+            not hasattr(self.visual.trunk, "attn_pool")
+            or self.visual.trunk.attn_pool is None
+        ):
+            raise ValueError("Attentional pooling not found in TimmModel")
+
+        self.visual.trunk.attn_pool.forward = types.MethodType(
+            hooked_attentional_pooler_timm_forward, self.visual.trunk.attn_pool
+        )
+        for block in self.visual.trunk.blocks:
+            if hasattr(block, "attn"):
+                block.attn.forward = types.MethodType(
+                    hooked_attention_forward, block.attn
+                )
+
+        # --- Deactivate gradient for module that don't need it ---
+        for name, param in self.named_parameters():
+            param.requires_grad = False
+            if name.startswith("visual.trunk.attn_pool"):
+                param.requires_grad = True
+            if name.startswith("visual.trunk.blocks"):
+                # get the depth
+                depth = int(name.split("visual.trunk.blocks.")[-1].split(".")[0])
+                if depth >= self.starting_depth:
+                    param.requires_grad = True
+
+        # --- Activate the hooks for the specific layers by modifying the block's forward ---
+        for layer in range(self.starting_depth, len(self.visual.trunk.blocks)):
+            self.visual.trunk.blocks[layer].forward = types.MethodType(
+                hooked_resblock_timm_forward, self.visual.trunk.blocks[layer]
+            )
+
+        self.visual.trunk.attn_pool.forward = types.MethodType(
+            hooked_attentional_pooler_timm_forward, self.visual.trunk.attn_pool
+        )
+
+    def compute_legrad(self, text_embedding, image=None, apply_correction=True):
+        if "clip" in self.model_type:
+            return self.compute_legrad_clip(text_embedding, image)
+        elif "siglip" in self.model_type:
+            return self.compute_legrad_siglip(
+                text_embedding, image, apply_correction=apply_correction
+            )
+        elif "coca" in self.model_type:
+            return self.compute_legrad_coca(text_embedding, image)
+
+    def compute_legrad_clip(self, text_embedding, image=None):
+        num_prompts = text_embedding.shape[0]
+        if image is not None:
+            # Ensure the image is passed through the model to get the intermediate features
+            _ = self.encode_image(image)
+
+        blocks_list = list(dict(self.visual.trunk.blocks.named_children()).values())
+
+        image_features_list = []
+
+        for layer in range(self.starting_depth, len(self.visual.trunk.blocks)):
+            # [num_patch, batch, dim]
+            intermediate_feat = blocks_list[layer].feat_post_mlp
+            # Mean over the patch tokens
+            intermediate_feat = intermediate_feat.mean(dim=1)
+            intermediate_feat = self.visual.head(
+                self.visual.trunk.norm(intermediate_feat)
+            )
+            intermediate_feat = F.normalize(intermediate_feat, dim=-1)
+            image_features_list.append(intermediate_feat)
+
+        num_tokens = blocks_list[-1].feat_post_mlp.shape[1] - 1
+        w = h = int(math.sqrt(num_tokens))
+
+        # ----- Get explainability map
+        accum_expl_map = 0
+        for layer, (blk, img_feat) in enumerate(
+            zip(blocks_list[self.starting_depth :], image_features_list)
+        ):
+            self.visual.zero_grad()
+            sim = text_embedding @ img_feat.transpose(-1, -2)  # [1, 1]
+            one_hot = (
+                F.one_hot(torch.arange(0, num_prompts))
+                .float()
+                .requires_grad_(True)
+                .to(text_embedding.device)
+            )
+            one_hot = torch.sum(one_hot * sim)
+
+            # [b, num_heads, N, N]
+            attn_map = blocks_list[self.starting_depth + layer].attn.attention_map
+
+            # -------- Get explainability map --------
+            # [batch_size * num_heads, N, N]
+            grad = torch.autograd.grad(
+                one_hot, [attn_map], retain_graph=True, create_graph=True
+            )[0]
+            # grad = rearrange(grad, '(b h) n m -> b h n m', b=num_prompts)  # separate batch and attn heads
+            grad = torch.clamp(grad, min=0.0)
+
+            # average attn over [CLS] + patch tokens
+            image_relevance = grad.mean(dim=1).mean(dim=1)[:, 1:]
+            expl_map = rearrange(image_relevance, "b (w h) -> 1 b w h", w=w, h=h)
+            # [B, 1, H, W]
+            expl_map = F.interpolate(
+                expl_map, scale_factor=self.patch_size, mode="bilinear"
+            )
+            accum_expl_map += expl_map
+
+        # Min-Max Norm
+        accum_expl_map = min_max(accum_expl_map)
+        return accum_expl_map
+
+    def compute_legrad_coca(self, text_embedding, image=None):
+        if image is not None:
+            _ = self.encode_image(image)
+
+        blocks_list = list(
+            dict(self.visual.transformer.resblocks.named_children()).values()
+        )
+
+        image_features_list = []
+
+        for layer in range(self.starting_depth, len(self.visual.transformer.resblocks)):
+            intermediate_feat = self.visual.transformer.resblocks[
+                layer
+            ].feat_post_mlp  # [num_patch, batch, dim]
+            intermediate_feat = intermediate_feat.permute(
+                1, 0, 2
+            )  # [batch, num_patch, dim]
+            image_features_list.append(intermediate_feat)
+
+        num_tokens = blocks_list[-1].feat_post_mlp.shape[0] - 1
+        w = h = int(math.sqrt(num_tokens))
+
+        # ----- Get explainability map
+        accum_expl_map = 0
+        for layer, (blk, img_feat) in enumerate(
+            zip(blocks_list[self.starting_depth :], image_features_list)
+        ):
+            self.visual.zero_grad()
+            # --- Apply attn_pool ---
+            image_embedding = self.visual.attn_pool(img_feat)[
+                :, 0
+            ]  # we keep only the first pooled token as it is only this one trained with the contrastive loss
+            image_embedding = image_embedding @ self.visual.proj
+
+            sim = text_embedding @ image_embedding.transpose(-1, -2)  # [1, 1]
+            one_hot = torch.sum(sim)
+
+            attn_map = (
+                self.visual.attn_pool.attn.attention_maps
+            )  # [num_heads, num_latent, num_patch]
+
+            # -------- Get explainability map --------
+            grad = torch.autograd.grad(
+                one_hot, [attn_map], retain_graph=True, create_graph=True
+            )[
+                0
+            ]  # [num_heads, num_latent, num_patch]
+            grad = torch.clamp(grad, min=0.0)
+
+            image_relevance = grad.mean(dim=0)[
+                0, 1:
+            ]  # average attn over heads + select first latent
+            expl_map = rearrange(image_relevance, "(w h) -> 1 1 w h", w=w, h=h)
+            expl_map = F.interpolate(
+                expl_map, scale_factor=self.patch_size, mode="bilinear"
+            )  # [B, 1, H, W]
+            accum_expl_map += expl_map
+
+        # Min-Max Norm
+        accum_expl_map = (accum_expl_map - accum_expl_map.min()) / (
+            accum_expl_map.max() - accum_expl_map.min()
+        )
+        return accum_expl_map
+
+    def _init_empty_embedding(self):
+        if not hasattr(self, "empty_embedding"):
+            # For the moment only SigLIP is supported & they all have the same tokenizer
+            _tok = open_clip.get_tokenizer(model_name="ViT-B-16-SigLIP")
+            empty_text = _tok(["a photo of a"]).to(self.logit_scale.data.device)
+            empty_embedding = self.encode_text(empty_text)
+            empty_embedding = F.normalize(empty_embedding, dim=-1)
+            self.empty_embedding = empty_embedding.t()
+
+    def compute_legrad_siglip(
+        self,
+        text_embedding,
+        image=None,
+        apply_correction=True,
+        correction_threshold=0.8,
+    ):
+        # --- Forward CLIP ---
+        blocks_list = list(dict(self.visual.trunk.blocks.named_children()).values())
+        if image is not None:
+            _ = self.encode_image(image)  # [bs, num_patch, dim] bs=num_masks
+
+        image_features_list = []
+        for blk in blocks_list[self.starting_depth :]:
+            intermediate_feat = blk.feat_post_mlp
+            image_features_list.append(intermediate_feat)
+
+        num_tokens = blocks_list[-1].feat_post_mlp.shape[1]
+        w = h = int(math.sqrt(num_tokens))
+
+        if apply_correction:
+            self._init_empty_embedding()
+            accum_expl_map_empty = 0
+
+        accum_expl_map = 0
+        for layer, (blk, img_feat) in enumerate(
+            zip(blocks_list[self.starting_depth :], image_features_list)
+        ):
+            self.zero_grad()
+            pooled_feat = self.visual.trunk.attn_pool(img_feat)
+            pooled_feat = F.normalize(pooled_feat, dim=-1)
+            # -------- Get explainability map --------
+            sim = text_embedding @ pooled_feat.transpose(-1, -2)  # [num_mask, num_mask]
+            one_hot = torch.sum(sim)
+            grad = torch.autograd.grad(
+                one_hot,
+                [self.visual.trunk.attn_pool.attn_probs],
+                retain_graph=True,
+                create_graph=True,
+            )[0]
+            grad = torch.clamp(grad, min=0.0)
+
+            image_relevance = grad.mean(dim=1)[
+                :, 0
+            ]  # average attn over [CLS] + patch tokens
+            expl_map = rearrange(image_relevance, "b (w h) -> b 1 w h", w=w, h=h)
+            accum_expl_map += expl_map
+
+            if apply_correction:
+                # -------- Get empty explainability map --------
+                sim_empty = pooled_feat @ self.empty_embedding
+                one_hot_empty = torch.sum(sim_empty)
+                grad_empty = torch.autograd.grad(
+                    one_hot_empty,
+                    [self.visual.trunk.attn_pool.attn_probs],
+                    retain_graph=True,
+                    create_graph=True,
+                )[0]
+                grad_empty = torch.clamp(grad_empty, min=0.0)
+
+                image_relevance_empty = grad_empty.mean(dim=1)[
+                    :, 0
+                ]  # average attn over heads + select query's row
+                expl_map_empty = rearrange(
+                    image_relevance_empty, "b (w h) -> b 1 w h", w=w, h=h
+                )
+                accum_expl_map_empty += expl_map_empty
+
+        if apply_correction:
+            heatmap_empty = min_max(accum_expl_map_empty)
+            accum_expl_map[heatmap_empty > correction_threshold] = 0
+
+        Res = min_max(accum_expl_map)
+        Res = F.interpolate(
+            Res, scale_factor=self.patch_size, mode="bilinear"
+        )  # [B, 1, H, W]
+
+        return Res
+
+
+class LePreprocess(nn.Module):
+    """
+    Modify OpenCLIP preprocessing to accept arbitrary image size.
+    """
+
+    def __init__(self, preprocess, image_size):
+        super(LePreprocess, self).__init__()
+        self.transform = Compose(
+            [
+                Resize(
+                    (image_size, image_size), interpolation=InterpolationMode.BICUBIC
+                ),
+                preprocess.transforms[-3],
+                preprocess.transforms[-2],
+                preprocess.transforms[-1],
+            ]
+        )
+
+    def forward(self, image):
+        return self.transform(image)

README.md

@@ -0,0 +1,12 @@
+---
+title: MedicalVisualPromptEngineering
+emoji: 🐠
+colorFrom: green
+colorTo: pink
+sdk: gradio
+sdk_version: 4.38.0
+app_file: app.py
+pinned: false
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,192 @@
+import gradio as gr
+import torch
+from PIL import Image
+import open_clip
+import numpy as np
+from LeGrad.legrad import LeWrapper, LePreprocess
+
+import cv2
+import numpy as np
+from PIL import Image
+
+# Load BiomedCLIP model
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model_name = "hf-hub:microsoft/BiomedCLIP-PubMedBERT_256-vit_base_patch16_224"
+model, preprocess = open_clip.create_model_from_pretrained(
+    model_name=model_name, device=device
+)
+tokenizer = open_clip.get_tokenizer(model_name=model_name)
+model = LeWrapper(model)  # Equip the model with LeGrad
+preprocess = LePreprocess(
+    preprocess=preprocess, image_size=448
+)  # Optional higher-res preprocessing
+
+
+def classify_image_with_biomedclip(editor_value, prompts):
+    # editor_value is a dict with keys: 'background', 'layers', 'composite'
+    # The 'composite' key contains the final annotated image
+
+    if editor_value is None:
+        return None, None
+
+    # Get the composite image (background + annotations)
+    image = editor_value["composite"]
+
+    # Ensure image is in PIL format
+    if not isinstance(image, Image.Image):
+        image = Image.fromarray(image)
+
+    # Preprocess and encode the image
+    image_input = preprocess(image).unsqueeze(0).to(device)
+    text_inputs = tokenizer(prompts).to(device)
+
+    # Encode text and image
+
+    text_embeddings = model.encode_text(text_inputs, normalize=True)
+    image_embeddings = model.encode_image(image_input, normalize=True)
+
+    # Generate probabilities (optional - not required for LeGrad explanations but included for completeness)
+    similarity = (
+        model.logit_scale.exp() * image_embeddings @ text_embeddings.T
+    ).softmax(dim=-1)
+    probabilities = similarity[0].detach().cpu().numpy()
+    explanation_maps = model.compute_legrad_clip(
+        image=image_input, text_embedding=text_embeddings[probabilities.argmax()]
+    )
+
+    # Convert explanation maps to heatmap
+    explanation_maps = explanation_maps.squeeze(0).detach().cpu().numpy()
+    explanation_map = (explanation_maps * 255).astype(np.uint8)  # Rescale to [0, 255]
+
+    return probabilities, explanation_map
+
+def update_output(editor_value, prompts_input):
+    prompts_list = [p.strip() for p in prompts_input.split(",") if p.strip()]
+    if not prompts_list:
+        return None, "Please enter at least one prompt."
+
+    probabilities, explanation_map = classify_image_with_biomedclip(
+        editor_value, prompts_list
+    )
+
+    if probabilities is None:
+        return None, "Please upload and annotate an image."
+
+    # Create probability display
+    prob_text = "\n".join(
+        [
+            f"{prompt}: {prob*100:.2f}%"
+            for prompt, prob in zip(prompts_list, probabilities)
+        ]
+    )
+
+    # Prepare the explanation map overlay
+    image = editor_value["composite"]
+    if not isinstance(image, Image.Image):
+        image = Image.fromarray(image)
+
+    explanation_image = explanation_map[0]
+    if isinstance(explanation_image, torch.Tensor):
+        explanation_image = explanation_image.cpu().numpy()
+
+    # Resize the explanation map to match the size of the original image
+    explanation_image_resized = cv2.resize(
+        explanation_image, (image.width, image.height)
+    )
+
+    # Normalize the explanation map for proper colormap application
+    explanation_image_resized = cv2.normalize(
+        explanation_image_resized, None, 0, 255, cv2.NORM_MINMAX
+    )
+
+    # Apply the colormap (e.g., COLORMAP_JET)
+    explanation_colormap = cv2.applyColorMap(
+        explanation_image_resized.astype(np.uint8), cv2.COLORMAP_JET
+    )
+
+    # Convert the original image to a format that OpenCV understands (RGB to BGR)
+    image_cv = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
+
+    # Blend the original image and the colormap
+    alpha = 0.5  # Transparency factor
+    blended_image = cv2.addWeighted(image_cv, 1 - alpha, explanation_colormap, alpha, 0)
+
+    # Convert back to RGB for displaying with PIL or matplotlib
+    blended_image_rgb = cv2.cvtColor(blended_image, cv2.COLOR_BGR2RGB)
+    output_image = Image.fromarray(blended_image_rgb)
+
+    return output_image, prob_text
+
+
+def clear_inputs():
+    return None, ""
+
+
+with gr.Blocks() as demo:
+    gr.Markdown(
+        "# ✨ Visual Prompt Engineering for Medical Vision Language Models in Radiology ✨",
+        elem_id="main-header",
+    )
+
+    gr.Markdown(
+        "This tool applies **visual prompt engineering to improve the classification of medical images using the BiomedCLIP**[3], the current state of the art in zero-shot biomedical image classification. By uploading biomedical images (e.g., chest X-rays), you can manually annotate areas of interest directly on the image. These annotations serve as visual prompts, which guide the model's attention on the region of interest. This technique improves the model's ability to focus on subtle yet important details.\n\n"
+        "After annotating and inputting text prompts (e.g., 'A chest X-ray with a benign/malignant lung nodule indicated by a red circle'), the tool returns classification results. These results are accompanied by **explainability maps** generated by **LeGrad** [3], which show where the model focused its attention, conditioned on the highest scoring text prompt. This helps to better interpret the model's decision-making process.\n\n"
+        "In our paper **[Visual Prompt Engineering for Medical Vision Language Models in Radiology](https://arxiv.org/pdf/2408.15802)**, we show, that visual prompts such as arrows, circles, and contours improve the zero-shot classification of biomedical vision language models in radiology."
+    )
+
+    gr.Markdown("---")  # Horizontal rule for separation
+
+    gr.Markdown(
+        "## 📝 **How It Works**:\n"
+        "1. **Upload** a biomedical image.\n"
+        "2. **Annotate** the image using the built-in editor to highlight regions of interest.\n"
+        "3. **Enter text prompts** separated by comma (e.g., 'A chest X-ray with a (benign/malignant) lung nodule indicated by a red circle').\n"
+        "4. **Submit** to get class probabilities and an explainability map conditioned on the highest scoring text prompt."
+    )
+
+    gr.Markdown("---")  # Horizontal rule for separation
+
+    with gr.Row():
+        with gr.Column():
+            image_editor = gr.ImageEditor(
+                label="Upload and Annotate Image",
+                type="pil",
+                interactive=True,
+                mirror_webcam=False,
+                layers=False,
+                # placeholder="Upload an image",
+                scale=2,
+            )
+            prompts_input = gr.Textbox(
+                placeholder="Enter prompts, comma-separated", label="Text Prompts"
+            )
+            submit_button = gr.Button("Submit", variant="primary")
+        with gr.Column():
+            output_image = gr.Image(
+                type="pil",
+                label="Output Image with Explanation Map",
+            )
+            prob_text = gr.Textbox(
+                label="Class Probabilities", interactive=False, lines=10
+            )
+
+    # Manually trigger the computation with the submit button
+    inputs = [image_editor, prompts_input]
+    outputs = [output_image, prob_text]
+    submit_button.click(fn=update_output, inputs=inputs, outputs=outputs)
+
+    gr.Markdown("---")  # Horizontal rule for separation
+
+    gr.Markdown("### 📝 **References**:\n")
+    gr.Markdown(
+        "[1] Denner, S., Bujotzek, M., Bounias, D., Zimmerer, D., Stock, R., Jäger, P.F. and Maier-Hein, K., 2024. **Visual Prompt Engineering for Medical Vision Language Models in Radiology**. arXiv preprint arXiv:2408.15802."
+    )
+    gr.Markdown(
+        "[2] Zhang, S., Xu, Y., Usuyama, N., Bagga, J., Tinn, R., Preston, S., Rao, R., Wei, M., Valluri, N., Wong, C. and Lungren, M.P., 2023. **Large-scale domain-specific pretraining for biomedical vision-language processing**. arXiv preprint arXiv:2303.00915, 2(3), p.6.\n"
+    )
+    gr.Markdown(
+        "[3] Bousselham, W., Boggust, A., Chaybouti, S., Strobelt, H. and Kuehne, H., 2024. **LeGrad: An Explainability Method for Vision Transformers via Feature Formation Sensitivity**. arXiv preprint arXiv:2404.03214."
+    )
+
+if __name__ == "__main__":
+    demo.launch(share=True)

requirements.txt

@@ -0,0 +1,5 @@
+gradio==4.42.0
+torch==2.4.1
+open_clip_torch==2.26.1
+legrad_torch
+transformers==4.44.2