import os import tempfile import warnings from PIL import Image import cv2 import gradio as gr import numpy as np import torch import torch.nn.functional as F from gradio_imageslider import ImageSlider from utils import interpret import clip_explain.CLIP.clip as clip device = "cuda" if torch.cuda.is_available() else "cpu" vit_model, vit_preprocess = clip.load("ViT-B/32", device=device, jit=False) cnn_model, cnn_preprocess = clip.load("RN50", device=device, jit=False) gradients_and_activations = {} def save_gradients(module, grad_in, grad_out): gradients_and_activations["gradients"] = grad_out[0].detach() def save_activations(module, input, output): gradients_and_activations["activations"] = output.detach() def get_cnn_gradcam(model, image, text, device): gradients_and_activations.clear() target_layer = model.visual.layer4[-1] handle_forward = target_layer.register_forward_hook(save_activations) handle_backward = target_layer.register_full_backward_hook(save_gradients) model.zero_grad() image_features = model.encode_image(image) text_features = model.encode_text(text) image_features = image_features / image_features.norm(dim=-1, keepdim=True) text_features = text_features / text_features.norm(dim=-1, keepdim=True) logit_scale = model.logit_scale.exp() logits = logit_scale * image_features @ text_features.t() score = logits[0, 0] score.backward(retain_graph=True) handle_forward.remove() handle_backward.remove() gradients = gradients_and_activations.get("gradients") activations = gradients_and_activations.get("activations") if gradients is None or activations is None: print("Warning: Gradients or activations are None. Using fallback layer.") gradients_and_activations.clear() model.zero_grad() target_layer = model.visual.layer3[-1] handle_forward = target_layer.register_forward_hook(save_activations) handle_backward = target_layer.register_full_backward_hook(save_gradients) image_features = model.encode_image(image) text_features = model.encode_text(text) image_features = image_features / image_features.norm(dim=-1, keepdim=True) text_features = text_features / text_features.norm(dim=-1, keepdim=True) logits = model.logit_scale.exp() * image_features @ text_features.t() score = logits[0, 0] score.backward(retain_graph=True) handle_forward.remove() handle_backward.remove() gradients = gradients_and_activations.get("gradients") activations = gradients_and_activations.get("activations") print(f"Gradients shape: {gradients.shape}") print(f"Activations shape: {activations.shape}") if len(activations.shape) == 4: weights = torch.mean(gradients, dim=(2, 3), keepdim=True) cam = torch.sum(weights * activations, dim=1, keepdim=True) cam = F.relu(cam) cam = F.interpolate(cam, size=(224, 224), mode='bilinear', align_corners=False) cam = cam.squeeze().detach().cpu().numpy() cam = (cam - cam.min()) / (cam.max() - cam.min() + 1e-8) else: pooled_gradients = torch.mean(gradients, dim=0) cam = torch.outer(activations.squeeze(0), pooled_gradients) cam = torch.sum(cam, dim=0).detach().cpu().numpy() size = int(np.sqrt(cam.shape[0])) cam = cam[:size*size].reshape(size, size) cam = cv2.resize(cam, (224, 224)) cam = np.maximum(cam, 0) cam = cam / (cam.max() + 1e-8) return cam def process_image(image, text, model_type="Transformer (ViT)"): if image is None or text.strip() == "": return (image, image), None if isinstance(image, np.ndarray): original_img = Image.fromarray(image) else: original_img = image if original_img.mode == 'I;16' or original_img.mode == 'I': original_img = np.array(original_img) original_img = ((original_img / 65535.0) * 255).astype(np.uint8) original_img = Image.fromarray(original_img) elif original_img.mode not in ['RGB', 'L']: original_img = original_img.convert('RGB') if model_type == "Transformer (ViT)": model = vit_model preprocess = vit_preprocess use_transformer = True else: model = cnn_model preprocess = cnn_preprocess use_transformer = False try: processed_img = preprocess(original_img).unsqueeze(0).to(device) except Exception as e: warnings.warn(f"Error preprocessing image: {e}") return (original_img, original_img), None texts = [text] tokenized_text = clip.tokenize(texts).to(device) if use_transformer: _, image_relevance = interpret( model=model, image=processed_img, texts=tokenized_text, device=device ) dim = int(image_relevance[0].numel() ** 0.5) rel = image_relevance[0].reshape(1, 1, dim, dim) rel = torch.nn.functional.interpolate(rel, size=224, mode='bilinear') rel = rel.reshape(224, 224).cpu().data.numpy() else: heatmap = get_cnn_gradcam(model, processed_img, tokenized_text, device) rel = cv2.resize(heatmap, (224, 224)) rel = (rel - rel.min()) / (rel.max() - rel.min()) resized_img = original_img.resize((224, 224)) img_array = np.array(resized_img) if len(img_array.shape) == 2: img_array = np.stack([img_array, img_array, img_array], axis=2) img_array = img_array / 255.0 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 np.uint8(255 * cam) vis = show_cam_on_image(img_array, rel) saliency_img = Image.fromarray(vis) with tempfile.NamedTemporaryFile(delete=False, suffix='.npy') as tmp: np_path = tmp.name np.save(np_path, rel) return (resized_img, saliency_img), np_path title = "CLIP Saliency Visualization" description = "Use this tool to visualize attention maps and salient regions of the CLIP model. Based on https://github.com/hila-chefer/Transformer-MM-Explainability." demo = gr.Interface( fn=process_image, inputs=[ gr.Image( type="pil", label="Upload Image", sources=["upload", "clipboard"] ), gr.Textbox( label="Enter Text Prompt", placeholder="Enter text to compare with..." ), gr.Radio( ["Transformer (ViT)", "CNN (ResNet)"], label="Model Architecture", value="Transformer (ViT)" ) ], outputs=[ ImageSlider(label="Original vs Saliency Map", type="pil", slider_color="#2BB8AB"), gr.File(label="Raw Depth (NumPy File)") ], title=title, description=description, allow_flagging="never" ) if __name__ == "__main__": demo.launch()