app.py

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()