| |
| |
| |
| |
| |
|
|
| import io |
| import json |
| import os |
| import math |
| import warnings |
| import base64 |
| import collections.abc |
| from itertools import repeat |
| from functools import partial |
| from typing import Dict, List, Optional, Tuple, Any |
|
|
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| import numpy as np |
| from PIL import Image |
| import matplotlib |
| matplotlib.use('Agg') |
| import matplotlib.pyplot as plt |
| import matplotlib.cm as cm |
| import torchvision.transforms as transforms |
|
|
| |
| try: |
| from einops import rearrange |
| except ImportError: |
| def rearrange(tensor, pattern, **kwargs): |
| if pattern == 'b n (qkv h d) -> qkv b h n d': |
| h = kwargs['h'] |
| qkv = kwargs['qkv'] |
| b, n, total_dim = tensor.shape |
| d = total_dim // (qkv * h) |
| return tensor.reshape(b, n, qkv, h, d).permute(2, 0, 3, 1, 4) |
| elif pattern == 'b h n d -> b n (h d)': |
| b, h, n, d = tensor.shape |
| return tensor.permute(0, 2, 1, 3).reshape(b, n, h * d) |
| raise NotImplementedError(f"Pattern {pattern} not implemented in fallback") |
|
|
|
|
| |
| |
| |
| def _ntuple(n): |
| def parse(x): |
| if isinstance(x, collections.abc.Iterable): |
| return x |
| return tuple(repeat(x, n)) |
| return parse |
|
|
| to_2tuple = _ntuple(2) |
|
|
|
|
| |
| |
| |
| def _no_grad_trunc_normal_(tensor, mean, std, a, b): |
| def norm_cdf(x): |
| return (1. + math.erf(x / math.sqrt(2.))) / 2. |
|
|
| if (mean < a - 2 * std) or (mean > b + 2 * std): |
| warnings.warn("mean is more than 2 std from [a, b] in nn.init.trunc_normal_. " |
| "The distribution of values may be incorrect.", |
| stacklevel=2) |
|
|
| with torch.no_grad(): |
| l = norm_cdf((a - mean) / std) |
| u = norm_cdf((b - mean) / std) |
| tensor.uniform_(2 * l - 1, 2 * u - 1) |
| tensor.erfinv_() |
| tensor.mul_(std * math.sqrt(2.)) |
| tensor.add_(mean) |
| tensor.clamp_(min=a, max=b) |
| return tensor |
|
|
|
|
| def trunc_normal_(tensor, mean=0., std=1., a=-2., b=2.): |
| return _no_grad_trunc_normal_(tensor, mean, std, a, b) |
|
|
|
|
| |
| |
| |
| def _conv_filter(state_dict, patch_size=16): |
| """ convert patch embedding weight from manual patchify + linear proj to conv""" |
| out_dict = {} |
| for k, v in state_dict.items(): |
| if 'patch_embed.proj.weight' in k: |
| if len(v.shape) == 2: |
| v = v.reshape((v.shape[0], 3, patch_size, patch_size)) |
| out_dict[k] = v |
| return out_dict |
|
|
|
|
| |
| |
| |
| class PatchEmbed(nn.Module): |
| """ Image to Patch Embedding |
| """ |
| def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768): |
| super().__init__() |
| img_size = to_2tuple(img_size) |
| patch_size = to_2tuple(patch_size) |
| num_patches = (img_size[1] // patch_size[1]) * (img_size[0] // patch_size[0]) |
| self.img_size = img_size |
| self.patch_size = patch_size |
| self.num_patches = num_patches |
|
|
| self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size) |
|
|
| def forward(self, x): |
| B, C, H, W = x.shape |
| assert H == self.img_size[0] and W == self.img_size[1], \ |
| f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})." |
| x = self.proj(x).flatten(2).transpose(1, 2) |
| return x |
|
|
|
|
| |
| |
| |
| class Mlp(nn.Module): |
| def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.): |
| super().__init__() |
| out_features = out_features or in_features |
| hidden_features = hidden_features or in_features |
| self.fc1 = nn.Linear(in_features, hidden_features) |
| self.act = act_layer() |
| self.fc2 = nn.Linear(hidden_features, out_features) |
| self.drop = nn.Dropout(drop) |
|
|
| def save_fc1_gradient(self, fc1_grad): |
| self.fc1_gradient = fc1_grad |
|
|
| def get_fc1_gradient(self): |
| return self.fc1_gradient |
|
|
| def save_fc2_gradient(self, fc2_grad): |
| self.fc2_gradient = fc2_grad |
|
|
| def get_fc2_gradient(self): |
| return self.fc2_gradient |
|
|
| def save_act_gradient(self, act_grad): |
| self.act_gradient = act_grad |
|
|
| def get_act_gradient(self): |
| return self.act_gradient |
|
|
| def forward(self, x): |
| x = self.fc1(x) |
| x = self.act(x) |
| x = self.drop(x) |
| x = self.fc2(x) |
| x = self.drop(x) |
| return x |
|
|
|
|
| |
| |
| |
| class Attention(nn.Module): |
| def __init__(self, dim, num_heads=8, qkv_bias=False, attn_drop=0., proj_drop=0., verbose=False): |
| super().__init__() |
| self.num_heads = num_heads |
| head_dim = dim // num_heads |
| self.scale = head_dim ** -0.5 |
|
|
| self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) |
| self.attn_drop = nn.Dropout(attn_drop) |
| self.proj = nn.Linear(dim, dim) |
| self.proj_drop = nn.Dropout(proj_drop) |
|
|
| self.attn_gradients = None |
| self.attention_map = None |
| self.vproj = None |
| self.input = None |
|
|
| def save_input(self, z): |
| self.input = z |
|
|
| def get_input(self): |
| return self.input |
|
|
| def save_attn_gradients(self, attn_gradients): |
| self.attn_gradients = attn_gradients |
|
|
| def get_attn_gradients(self): |
| return self.attn_gradients |
|
|
| def save_attention_map(self, attention_map): |
| self.attention_map = attention_map |
|
|
| def get_attention_map(self): |
| return self.attention_map |
|
|
| def forward(self, x, register_hook=False): |
| b, n, _ = x.shape |
| h = self.num_heads |
| self.save_input(x) |
|
|
| qkv = self.qkv(x) |
| q, k, v = rearrange(qkv, 'b n (qkv h d) -> qkv b h n d', qkv=3, h=h) |
|
|
| W_qkv = self.qkv.weight |
| dim = x.shape[-1] |
|
|
| W_v = W_qkv[2 * dim : 3 * dim, :] |
| self.W_v = W_v |
| self.W_proj = self.proj.weight |
|
|
| dots = torch.einsum('bhid,bhjd->bhij', q, k) * self.scale |
| attn = dots.softmax(dim=-1) |
| attn = self.attn_drop(attn) |
| attention_score = torch.einsum('bhij,bhjd->bhid', attn, v) |
|
|
| self.save_attention_map(attn) |
| if register_hook: |
| attn.register_hook(self.save_attn_gradients) |
|
|
| attention_score_rearranged = rearrange(attention_score, 'b h n d -> b n (h d)') |
| out = self.proj(attention_score_rearranged) |
| out = self.proj_drop(out) |
|
|
| self.vproj = torch.matmul(rearrange(v, 'b h n d -> b n (h d)'), self.proj.weight.t()) |
| return out |
|
|
|
|
| |
| |
| |
| class Block(nn.Module): |
| def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, drop=0., attn_drop=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm, block_index=None, verbose=False): |
| super().__init__() |
| self.norm1 = norm_layer(dim) |
| self.attn = Attention(dim, num_heads=num_heads, qkv_bias=qkv_bias, attn_drop=attn_drop, proj_drop=drop, verbose=verbose) |
| self.norm2 = norm_layer(dim) |
|
|
| mlp_hidden_dim = int(dim * mlp_ratio) |
| self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop) |
| self.input = None |
| self.output = None |
| self.tilde = None |
| self.block_index = block_index |
| self.verbose = verbose |
|
|
| def save_input(self, z): |
| self.input = z |
|
|
| def get_input(self): |
| return self.input |
|
|
| def save_tilde(self, z): |
| self.tilde = z |
|
|
| def get_tilde(self): |
| return self.tilde |
|
|
| def save_output(self, z): |
| self.output = z |
|
|
| def get_output(self): |
| return self.output |
|
|
| def forward(self, x, register_hook=False): |
| self.save_input(x) |
| attn_out = self.attn(self.norm1(x), register_hook=register_hook) |
| out1 = x + attn_out |
| self.save_tilde(out1) |
|
|
| mlp_out = self.mlp(self.norm2(out1)) |
| out = out1 + mlp_out |
| self.save_output(out) |
| return out |
|
|
|
|
| |
| |
| |
| class VisionTransformer(nn.Module): |
| """ Vision Transformer |
| """ |
| def __init__(self, img_size=224, patch_size=16, in_chans=3, num_classes=1000, embed_dim=768, depth=12, |
| num_heads=12, mlp_ratio=4., qkv_bias=False, drop_rate=0., attn_drop_rate=0., norm_layer=nn.LayerNorm, |
| global_pooling=False, dino=False, verbose=False): |
| super().__init__() |
|
|
| self.global_pooling = global_pooling |
| self.num_classes = num_classes |
| self.num_features = self.embed_dim = embed_dim |
| self.patch_embed = PatchEmbed( |
| img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim) |
| num_patches = self.patch_embed.num_patches |
|
|
| self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim)) |
| self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim)) |
| self.pos_drop = nn.Dropout(p=drop_rate) |
|
|
| self.blocks = nn.ModuleList([ |
| Block( |
| dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, |
| drop=drop_rate, attn_drop=attn_drop_rate, norm_layer=norm_layer, block_index=i, verbose=verbose) |
| for i in range(depth)]) |
| if global_pooling: |
| self.fc_norm = norm_layer(embed_dim) |
| else: |
| self.norm = norm_layer(embed_dim) |
|
|
| if dino: |
| self.head = nn.Linear(embed_dim * 2, num_classes) if num_classes > 0 else nn.Identity() |
| else: |
| self.head = nn.Linear(embed_dim, num_classes) if num_classes > 0 else nn.Identity() |
|
|
| self.cls_gradients = None |
| self.input_grad = None |
| self.dino = dino |
|
|
| trunc_normal_(self.pos_embed, std=.02) |
| trunc_normal_(self.cls_token, std=.02) |
| self.apply(self._init_weights) |
|
|
| def _init_weights(self, m): |
| if isinstance(m, nn.Linear): |
| trunc_normal_(m.weight, std=.02) |
| if isinstance(m, nn.Linear) and m.bias is not None: |
| nn.init.constant_(m.bias, 0) |
| elif isinstance(m, nn.LayerNorm): |
| nn.init.constant_(m.bias, 0) |
| nn.init.constant_(m.weight, 1.0) |
|
|
| def forward(self, x, register_hook=False): |
| B = x.shape[0] |
| x = self.patch_embed(x) |
| cls_tokens = self.cls_token.expand(B, -1, -1) |
| x = torch.cat((cls_tokens, x), dim=1) |
| x = x + self.pos_embed |
| x = self.pos_drop(x) |
|
|
| for blk in self.blocks[:-1]: |
| x = blk(x, register_hook=register_hook) |
|
|
| x = self.blocks[-1](x, register_hook=register_hook) |
|
|
| if self.global_pooling: |
| x = x[:, 1:, :].mean(axis=1) |
| x = self.fc_norm(x) |
| elif self.dino: |
| x = self.norm(x) |
| x = torch.cat((x[:, 0].unsqueeze(-1), x[:, 1:, :].mean(axis=1).unsqueeze(-1)), dim=-1) |
| x = x.reshape(x.shape[0], -1) |
| else: |
| x = self.norm(x) |
| x = x[:, 0] |
|
|
| x = self.head(x) |
| return x |
|
|
|
|
| |
| |
| |
| def compute_rollout_attention(all_layer_matrices, start_layer=0): |
| num_tokens = all_layer_matrices[0].shape[1] |
| batch_size = all_layer_matrices[0].shape[0] |
| eye = torch.eye(num_tokens).expand(batch_size, num_tokens, num_tokens).to(all_layer_matrices[0].device) |
| all_layer_matrices = [all_layer_matrices[i] + eye for i in range(len(all_layer_matrices))] |
| matrices_aug = [all_layer_matrices[i] / all_layer_matrices[i].sum(dim=-1, keepdim=True) |
| for i in range(len(all_layer_matrices))] |
| joint_attention = matrices_aug[start_layer] |
| for i in range(start_layer+1, len(matrices_aug)): |
| joint_attention = matrices_aug[i].bmm(joint_attention) |
| return joint_attention |
|
|
|
|
| def generate_rollout(model, input_tensor, device=None, start_layer=0, mae=False): |
| input_tensor = input_tensor.to(device) |
| out = model(input_tensor) |
|
|
| all_layer_attentions = [] |
| for blk in model.blocks: |
| attn = blk.attn.get_attention_map() |
| attn_avg = attn.mean(dim=1).detach() |
| all_layer_attentions.append(attn_avg) |
|
|
| rollout = compute_rollout_attention(all_layer_attentions, start_layer=start_layer) |
| if mae: |
| return rollout[:, 1:, 1:].mean(1) |
| else: |
| return rollout[:, 0, 1:] |
|
|
|
|
| |
| |
| |
| def beyond_intuition_tokenwise(model, x, device, onehot=None, ffd_consideration=False, index=None, steps=20, start_layer=1, samples=20, noise=0.1, mae=False, ssl=False, dino=False, taken_head_idx = None): |
| b = x.shape[0] |
| batch_size = b |
| x = x.to(device) |
|
|
| output = model(x) |
| logit_score = output |
|
|
| if index is None: |
| index = np.argmax(output.cpu().data.numpy(), axis=-1) |
| elif isinstance(index, int): |
| index = np.array([index]) |
|
|
| out_sum = torch.sum(output) |
|
|
| one_hot = np.zeros((batch_size, output.size()[-1]), dtype=np.float32) |
| one_hot[np.arange(batch_size), index] = 1 |
| one_hot = torch.from_numpy(one_hot).requires_grad_(True).to(device) |
| one_hot = torch.sum(one_hot) |
|
|
| model.zero_grad() |
| one_hot.backward(retain_graph=True) |
|
|
| _, num_head, num_tokens, _ = model.blocks[-1].attn.get_attention_map().shape |
| R = torch.eye(num_tokens, num_tokens).expand(b, num_tokens, num_tokens).to(device) |
|
|
| for blk_idx, blk in enumerate(model.blocks): |
| if blk_idx < start_layer - 1: |
| continue |
|
|
| |
| fc2_w = blk.mlp.fc2.weight |
| h_dim = fc2_w.shape[1] // num_head |
| impw = torch.abs(fc2_w.reshape((num_head, h_dim, -1))).mean(2).mean(1).unsqueeze(1).unsqueeze(2) |
|
|
| z = blk.get_input() |
| vproj = blk.attn.vproj |
| order = torch.linalg.norm(vproj, dim=-1).squeeze() / (torch.linalg.norm(z, dim=-1).squeeze() + 1e-8) |
| m = torch.diag_embed(order) |
|
|
| cam = blk.attn.get_attention_map() |
| cam = cam.reshape(-1, cam.shape[-2], cam.shape[-1]) |
|
|
| if taken_head_idx is None: |
| if ffd_consideration: |
| cam = cam * impw |
| cam = cam.mean(0) |
| else: |
| cam = cam[taken_head_idx] |
|
|
| O_t = torch.matmul(cam.to(device), m.to(device)) |
| R = R + torch.matmul(O_t, R.to(device)) |
|
|
| if ssl: |
| if mae: |
| return R[:, 1:, 1:].abs().mean(axis=1), logit_score |
| elif dino: |
| return R[:, 1:, 1:].abs().mean(axis=1) + R[:, 0, 1:].abs(), logit_score |
| else: |
| return R[:, 0, 1:].abs(), logit_score |
|
|
| total_gradients = torch.zeros(b, num_head, num_tokens, num_tokens, device=device) |
| for alpha in np.linspace(0, 1, steps): |
| data_scaled = x * alpha |
| output = model(data_scaled, register_hook=True) |
| one_hot = np.zeros((b, output.size()[-1]), dtype=np.float32) |
| one_hot[np.arange(b), index] = 1 |
| one_hot = torch.from_numpy(one_hot).requires_grad_(True).to(device) |
| one_hot = torch.sum(one_hot * output) |
|
|
| model.zero_grad() |
| one_hot.backward(retain_graph=True) |
|
|
| gradients = model.blocks[-1].attn.get_attn_gradients() |
| if gradients is not None: |
| total_gradients += gradients |
|
|
| W_state = (total_gradients / steps).clamp(min=0).mean(1).reshape(b, num_tokens, num_tokens) |
| R = W_state * R.abs() |
|
|
| if mae: |
| R_MAE = R[:, 1:, 1:].mean(axis=1) |
| return (R_MAE, logit_score) |
| elif dino: |
| R_DINO = (R[:, 1:, 1:].mean(axis=1) + R[:, 0, 1:]) |
| return (R_DINO, logit_score) |
| else: |
| R_general = R[:, 0, 1:] |
| return (R_general, logit_score) |
|
|
|
|
| |
| |
| |
| def integrated_gradient(model, x, device, index=None, steps=20, mae=False, dino=False, ssl=False): |
| b = x.shape[0] |
| batch_size = b |
| x = x.to(device) |
|
|
| output = model(x) |
| if index is None: |
| index = np.argmax(output.cpu().data.numpy(), axis=-1) |
| elif isinstance(index, int): |
| index = np.array([index]) |
|
|
| one_hot = np.zeros((batch_size, output.size()[-1]), dtype=np.float32) |
| one_hot[np.arange(batch_size), index] = 1 |
| one_hot = torch.from_numpy(one_hot).requires_grad_(True).to(device) |
| one_hot = torch.sum(one_hot) |
|
|
| model.zero_grad() |
| one_hot.backward(retain_graph=True) |
|
|
| _, num_head, num_tokens, _ = model.blocks[-1].attn.get_attention_map().shape |
| total_gradients = torch.zeros(b, num_head, num_tokens, num_tokens, device=device) |
|
|
| for alpha in np.linspace(0, 1, steps): |
| data_scaled = x * alpha |
| output = model(data_scaled, register_hook=True) |
| one_hot = np.zeros((b, output.size()[-1]), dtype=np.float32) |
| one_hot[np.arange(b), index] = 1 |
| one_hot = torch.from_numpy(one_hot).requires_grad_(True).to(device) |
| one_hot = torch.sum(one_hot * output) |
|
|
| model.zero_grad() |
| one_hot.backward(retain_graph=True) |
|
|
| gradients = model.blocks[-1].attn.get_attn_gradients() |
| if gradients is not None: |
| total_gradients += gradients |
|
|
| W_state = (total_gradients / steps).clamp(min=0).mean(1).reshape(b, num_tokens, num_tokens) |
| R = W_state |
|
|
| if mae: |
| R_MAE = R[:, 1:, 1:].mean(axis=1) |
| return R_MAE |
| elif dino: |
| R_DINO = (R[:, 1:, 1:].mean(axis=1) + R[:, 0, 1:]) |
| return R_DINO |
| else: |
| R_general = R[:, 0, 1:] |
| return R_general |
|
|
|
|
| |
| |
| |
| def gradcam(model, x, device, target_cls_idx=None, mae=False): |
| x = x.to(device) |
| output = model(x, register_hook=True) |
|
|
| if target_cls_idx is None: |
| target_cls_idx = np.argmax(output.cpu().data.numpy()) |
|
|
| one_hot = np.zeros((1, output.size()[-1]), dtype=np.float32) |
| one_hot[0][target_cls_idx] = 1 |
| one_hot = torch.from_numpy(one_hot).requires_grad_(True).to(device) |
| one_hot = torch.sum(one_hot * output) |
|
|
| model.zero_grad() |
| one_hot.backward(retain_graph=True) |
|
|
| grad = model.blocks[-1].attn.get_attn_gradients().to(device) |
| cam = model.blocks[-1].attn.get_attention_map().to(device) |
|
|
| if mae: |
| cam = cam[0, :, 1:, 1:] |
| grad = grad[0, :, 1:, 1:] |
| grad = grad.mean(dim=[0, 2], keepdim=True) |
| cam = (cam * grad).mean(0).mean(1).clamp(min=0) |
| cam = (cam - cam.min()) / (cam.max() - cam.min() + 1e-8) |
| else: |
| cam = cam[0, :, 0, 1:].reshape(-1, 14, 14) |
| grad = grad[0, :, 0, 1:].reshape(-1, 14, 14) |
| grad = grad.mean(dim=[1, 2], keepdim=True) |
| cam = (cam * grad).mean(0).clamp(min=0) |
| cam = (cam - cam.min()) / (cam.max() - cam.min() + 1e-8) |
|
|
| return cam |
|
|
|
|
| def normalize_class_name(s: str) -> str: |
| import re |
| |
| s = re.sub(r'^n\d{8}[_\s]+', '', str(s).strip().lower()) |
| |
| return re.sub(r'[\s_\-]+', ' ', s).strip() |
|
|
|
|
| def compact_class_name(s: str) -> str: |
| import re |
| return re.sub(r'[^a-z0-9]', '', str(s).lower()) |
|
|
|
|
| |
| |
| |
| class ViTExplanationService: |
| def __init__(self, base_dir: str = "."): |
| self.base_dir = base_dir |
| |
| if os.environ.get("SPACES_ZERO_GPU") or os.environ.get("FORCE_CPU"): |
| self.device = torch.device("cpu") |
| else: |
| self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
| self.models: Dict[str, VisionTransformer] = {} |
| self.id2label: Dict[int, str] = {} |
| self.label2id: Dict[str, int] = {} |
| self.norm_label2id: Dict[str, int] = {} |
| self.compact_label2id: Dict[str, int] = {} |
| self._load_classes() |
|
|
| |
| self.transform = transforms.Compose([ |
| transforms.Resize((224, 224)), |
| transforms.ToTensor(), |
| transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]), |
| ]) |
|
|
| def _index_label(self, idx: int, label_str: str): |
| self.id2label[idx] = label_str |
| candidates = [label_str] + [p.strip() for p in label_str.split(",") if p.strip()] |
| for cand in candidates: |
| raw_l = cand.lower() |
| self.label2id[raw_l] = idx |
| self.label2id[cand.replace(" ", "_").lower()] = idx |
| norm = normalize_class_name(cand) |
| if norm: |
| self.norm_label2id[norm] = idx |
| comp = compact_class_name(cand) |
| if comp: |
| self.compact_label2id[comp] = idx |
|
|
| def _load_classes(self): |
| possible_paths = [ |
| os.path.join(self.base_dir, "draft", "vit_config.json"), |
| os.path.join(self.base_dir, "vit_config.json"), |
| os.path.join(self.base_dir, "static", "classes.json") |
| ] |
| loaded = False |
| for p in possible_paths: |
| if os.path.exists(p): |
| try: |
| with open(p, "r") as f: |
| data = json.load(f) |
| if "id2label" in data: |
| for k, v in data["id2label"].items(): |
| self._index_label(int(k), v) |
| loaded = True |
| break |
| elif isinstance(data, list): |
| for item in data: |
| self._index_label(int(item["index"]), item["name"]) |
| loaded = True |
| break |
| except Exception as e: |
| print(f"Notice: error reading {p}: {e}") |
|
|
| if not loaded: |
| for i in range(1000): |
| self._index_label(i, f"Class {i}") |
|
|
| def get_classes_list(self) -> List[Dict[str, Any]]: |
| return [{"index": i, "name": self.id2label.get(i, f"Class {i}")} for i in range(1000)] |
|
|
| def resolve_class_index(self, query: str, query_type: str = "auto") -> Tuple[int, str]: |
| query_clean = str(query).strip() |
| if not query_clean: |
| return 0, self.id2label.get(0, "tench") |
|
|
| if query_type == "number" or (query_type == "auto" and query_clean.isdigit()): |
| idx = int(query_clean) |
| idx = max(0, min(999, idx)) |
| return idx, self.id2label.get(idx, f"Class {idx}") |
|
|
| q_lower = query_clean.lower() |
| q_norm = normalize_class_name(query_clean) |
| q_comp = compact_class_name(query_clean) |
|
|
| |
| if q_lower in self.label2id: |
| idx = self.label2id[q_lower] |
| return idx, self.id2label[idx] |
| if q_norm in self.norm_label2id: |
| idx = self.norm_label2id[q_norm] |
| return idx, self.id2label[idx] |
| if q_comp in self.compact_label2id: |
| idx = self.compact_label2id[q_comp] |
| return idx, self.id2label[idx] |
|
|
| |
| for name, idx in self.norm_label2id.items(): |
| if q_norm in name or name in q_norm: |
| return idx, self.id2label[idx] |
|
|
| |
| for comp, idx in self.compact_label2id.items(): |
| if q_comp in comp or comp in q_comp: |
| return idx, self.id2label[idx] |
|
|
| return 0, self.id2label.get(0, "tench") |
|
|
| def get_model(self, model_key: str = "vit_base_patch16_224") -> VisionTransformer: |
| if model_key in self.models: |
| return self.models[model_key] |
|
|
| models_dir = os.path.join(self.base_dir, "models") |
| if model_key == "vit_small_patch16_224": |
| model = VisionTransformer( |
| patch_size=16, embed_dim=768, depth=8, num_classes=1000, num_heads=6, |
| mlp_ratio=3, qkv_bias=False, norm_layer=partial(nn.LayerNorm, eps=1e-6) |
| ) |
| ckpt_path = os.path.join(models_dir, "vit_small_p16_224-15ec54c9.pth") |
| else: |
| model = VisionTransformer( |
| patch_size=16, embed_dim=768, depth=12, num_classes=1000, num_heads=12, |
| mlp_ratio=4, qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6) |
| ) |
| ckpt_path = os.path.join(models_dir, "jx_vit_base_p16_224-80ecf9dd.pth") |
|
|
| if os.path.exists(ckpt_path): |
| state = torch.load(ckpt_path, map_location="cpu") |
| state = _conv_filter(state) |
| model.load_state_dict(state, strict=False) |
|
|
| model.to(self.device) |
| model.eval() |
| self.models[model_key] = model |
| return model |
|
|
| def _render_heatmaps(self, raw_map: torch.Tensor, original_pil: Image.Image) -> Tuple[str, str]: |
| arr = raw_map.detach().cpu().numpy() |
| if arr.size == 196: |
| arr = arr.reshape(14, 14) |
| elif arr.shape[-2:] != (14, 14): |
| arr = arr.reshape(14, 14) |
|
|
| |
| |
| fig, ax = plt.subplots(figsize=(3, 3)) |
| ax.imshow(arr, cmap="Reds") |
| ax.axis("off") |
| fig.subplots_adjust(left=0, right=1, bottom=0, top=1) |
| raw_buf = io.BytesIO() |
| fig.savefig(raw_buf, format="png", bbox_inches="tight", pad_inches=0) |
| plt.close(fig) |
| saliency_base64 = "data:image/png;base64," + base64.b64encode(raw_buf.getvalue()).decode("utf-8") |
|
|
| |
| h_min, h_max = float(arr.min()), float(arr.max()) |
| if h_max > h_min: |
| norm_map = (arr - h_min) / (h_max - h_min) |
| else: |
| norm_map = np.zeros_like(arr) |
|
|
| norm_tensor = torch.from_numpy(norm_map).unsqueeze(0).unsqueeze(0).float() |
| upsampled = F.interpolate(norm_tensor, size=(224, 224), mode='bilinear', align_corners=False).squeeze().numpy() |
| seismic_rgba = cm.seismic(upsampled) |
| seismic_rgb = (seismic_rgba[:, :, :3] * 255).astype(np.uint8) |
| orig_resized = original_pil.resize((224, 224)).convert("RGB") |
| orig_arr = np.array(orig_resized) |
| blended = (0.55 * orig_arr + 0.45 * seismic_rgb).astype(np.uint8) |
| overlay_buf = io.BytesIO() |
| Image.fromarray(blended).save(overlay_buf, format="PNG") |
| overlay_base64 = "data:image/png;base64," + base64.b64encode(overlay_buf.getvalue()).decode("utf-8") |
|
|
| return saliency_base64, overlay_base64 |
|
|
| def run_simulation( |
| self, |
| pil_image: Image.Image, |
| class_query: str, |
| query_type: str = "auto", |
| model_name: str = "vit_base_patch16_224", |
| selected_methods: Optional[List[str]] = None |
| ) -> Dict[str, Any]: |
| if selected_methods is None or len(selected_methods) == 0: |
| selected_methods = ["f2hf", "beyond_intuition", "integrated_gradients", "gradcam", "rollout"] |
|
|
| model = self.get_model(model_name) |
| input_tensor = self.transform(pil_image).unsqueeze(0).to(self.device) |
|
|
| with torch.no_grad(): |
| logits = model(input_tensor) |
| probs = F.softmax(logits, dim=-1).squeeze().cpu().numpy() |
|
|
| top5_indices = np.argsort(probs)[::-1][:5] |
| top5 = [ |
| {"index": int(i), "name": self.id2label.get(int(i), f"Class {i}"), "score": float(probs[i])} |
| for i in top5_indices |
| ] |
| pred_top1 = top5[0] |
|
|
| target_idx, target_name = self.resolve_class_index(class_query, query_type) |
| heatmaps: Dict[str, Any] = {} |
|
|
| |
| if "f2hf" in selected_methods: |
| h_f2hf, _ = beyond_intuition_tokenwise(model, input_tensor, self.device, index=target_idx, ffd_consideration=True) |
| sal_b64, over_b64 = self._render_heatmaps(h_f2hf, pil_image) |
| heatmaps["f2hf"] = { |
| "id": "f2hf", |
| "name": "F2HF", |
| "badge": "Proposed Method", |
| "description": "Feed-Forward Network Filtered Head Attention", |
| "image": sal_b64, |
| "saliency_image": sal_b64, |
| "overlay_image": over_b64 |
| } |
|
|
| |
| if "beyond_intuition" in selected_methods: |
| h_bi, _ = beyond_intuition_tokenwise(model, input_tensor, self.device, index=target_idx, ffd_consideration=False, start_layer=1) |
| sal_b64, over_b64 = self._render_heatmaps(h_bi, pil_image) |
| heatmaps["beyond_intuition"] = { |
| "id": "beyond_intuition", |
| "name": "BeyondIntuition", |
| "badge": "Baseline", |
| "description": "Tokenwise Self-Attention Attribution", |
| "image": sal_b64, |
| "saliency_image": sal_b64, |
| "overlay_image": over_b64 |
| } |
|
|
| |
| if "integrated_gradients" in selected_methods: |
| h_ig = integrated_gradient(model, input_tensor, self.device, index=target_idx, steps=10) |
| sal_b64, over_b64 = self._render_heatmaps(h_ig, pil_image) |
| heatmaps["integrated_gradients"] = { |
| "id": "integrated_gradients", |
| "name": "Integrated Gradients", |
| "badge": "Baseline", |
| "description": "Path-Integrated Head Attentions", |
| "image": sal_b64, |
| "saliency_image": sal_b64, |
| "overlay_image": over_b64 |
| } |
|
|
| |
| if "gradcam" in selected_methods: |
| h_gcam = gradcam(model, input_tensor, self.device, target_cls_idx=target_idx) |
| sal_b64, over_b64 = self._render_heatmaps(h_gcam, pil_image) |
| heatmaps["gradcam"] = { |
| "id": "gradcam", |
| "name": "Grad-CAM", |
| "badge": "Baseline", |
| "description": "Attention Map Weighted by Target Gradients", |
| "image": sal_b64, |
| "saliency_image": sal_b64, |
| "overlay_image": over_b64 |
| } |
|
|
| |
| if "rollout" in selected_methods: |
| h_roll = generate_rollout(model, input_tensor, device=self.device, start_layer=1) |
| sal_b64, over_b64 = self._render_heatmaps(h_roll, pil_image) |
| heatmaps["rollout"] = { |
| "id": "rollout", |
| "name": "Attention Rollout", |
| "badge": "Baseline", |
| "description": "Joint Inter-Layer Attention Flow", |
| "image": sal_b64, |
| "saliency_image": sal_b64, |
| "overlay_image": over_b64 |
| } |
|
|
| thumb = pil_image.copy() |
| thumb.thumbnail((300, 300)) |
| orig_buf = io.BytesIO() |
| thumb.save(orig_buf, format="PNG") |
| orig_base64 = "data:image/png;base64," + base64.b64encode(orig_buf.getvalue()).decode("utf-8") |
|
|
| return { |
| "prediction": pred_top1, |
| "top5": top5, |
| "target": {"index": target_idx, "name": target_name}, |
| "model_name": model_name, |
| "original_image": orig_base64, |
| "heatmaps": heatmaps |
| } |
|
|
