|
import sys |
|
|
|
import cv2 |
|
import numpy as np |
|
import torch |
|
|
|
from imagenet_class_indices import CLS2IDX |
|
|
|
sys.path.append("Transformer-Explainability") |
|
|
|
|
|
from baselines.ViT.ViT_explanation_generator import LRP, Baselines |
|
from baselines.ViT.ViT_LRP import vit_base_patch16_224 as vit_LRP |
|
from baselines.ViT.ViT_new import vit_base_patch16_224 as vit |
|
|
|
|
|
|
|
def show_cam_on_image(img, mask): |
|
heatmap = cv2.applyColorMap(np.uint8(255 * mask), cv2.COLORMAP_JET) |
|
heatmap = np.float32(heatmap) / 255 |
|
cam = heatmap + np.float32(img) |
|
cam = cam / np.max(cam) |
|
return cam |
|
|
|
|
|
|
|
model = vit_LRP(pretrained=True) |
|
model.eval() |
|
attribution_generator = LRP(model) |
|
model_baseline = vit(pretrained=True) |
|
model_baseline.eval() |
|
baselines_generator = Baselines(model_baseline) |
|
|
|
|
|
def generate_visualization( |
|
original_image, class_index=None, method="transformer_attribution", LRP=True |
|
): |
|
if LRP: |
|
transformer_attribution = attribution_generator.generate_LRP( |
|
original_image.unsqueeze(0), method=method, index=class_index |
|
).detach() |
|
else: |
|
if method == "gradcam": |
|
transformer_attribution = baselines_generator.generate_cam_attn( |
|
original_image.unsqueeze(0), index=class_index |
|
).detach() |
|
else: |
|
transformer_attribution = baselines_generator.generate_rollout( |
|
original_image.unsqueeze(0) |
|
).detach() |
|
if method != "full": |
|
transformer_attribution = transformer_attribution.reshape(1, 1, 14, 14) |
|
transformer_attribution = torch.nn.functional.interpolate( |
|
transformer_attribution, scale_factor=16, mode="bilinear" |
|
) |
|
else: |
|
transformer_attribution = transformer_attribution.reshape(1, 1, 224, 224) |
|
transformer_attribution = ( |
|
transformer_attribution.reshape(224, 224).data.cpu().numpy() |
|
) |
|
transformer_attribution = ( |
|
transformer_attribution - transformer_attribution.min() |
|
) / (transformer_attribution.max() - transformer_attribution.min()) |
|
|
|
image_transformer_attribution = original_image.permute(1, 2, 0).data.cpu().numpy() |
|
image_transformer_attribution = ( |
|
image_transformer_attribution - image_transformer_attribution.min() |
|
) / (image_transformer_attribution.max() - image_transformer_attribution.min()) |
|
vis = show_cam_on_image(image_transformer_attribution, transformer_attribution) |
|
vis = np.uint8(255 * vis) |
|
vis = cv2.cvtColor(np.array(vis), cv2.COLOR_RGB2BGR) |
|
return vis |
|
|
|
|
|
def print_top_classes(predictions, **kwargs): |
|
|
|
prob = torch.softmax(predictions, dim=1) |
|
class_indices = predictions.data.topk(5, dim=1)[1][0].tolist() |
|
max_str_len = 0 |
|
class_names = [] |
|
for cls_idx in class_indices: |
|
class_names.append(CLS2IDX[cls_idx]) |
|
if len(CLS2IDX[cls_idx]) > max_str_len: |
|
max_str_len = len(CLS2IDX[cls_idx]) |
|
|
|
print("Top 5 classes:") |
|
for cls_idx in class_indices: |
|
output_string = "\t{} : {}".format(cls_idx, CLS2IDX[cls_idx]) |
|
output_string += " " * (max_str_len - len(CLS2IDX[cls_idx])) + "\t\t" |
|
output_string += "value = {:.3f}\t prob = {:.1f}%".format( |
|
predictions[0, cls_idx], 100 * prob[0, cls_idx] |
|
) |
|
print(output_string) |
|
|
