Spaces:
Runtime error
Runtime error
| import requests | |
| import numpy as np | |
| import cv2 as cv2 | |
| from PIL import Image | |
| import torch | |
| import torch.nn.functional as F | |
| import open_clip | |
| import gradio as gr | |
| import spaces | |
| from legrad import LeWrapper, LePreprocess | |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| layer_index = -2 # will run on cpu | |
| image_size = 448 | |
| # ---------- Init CLIP Model ---------- | |
| model_name = 'ViT-B-16' | |
| pretrained = 'laion2b_s34b_b88k' | |
| patch_size = 16 | |
| model, _, preprocess = open_clip.create_model_and_transforms(model_name, pretrained=pretrained, device=device) | |
| tokenizer = open_clip.get_tokenizer(model_name) | |
| # ---------- Apply LeGrad's wrappers ---------- | |
| model = LeWrapper(model) | |
| preprocess = LePreprocess(preprocess=preprocess, image_size=image_size) | |
| # ---------- Function to load image from URL ---------- | |
| def change_to_url(url): | |
| img_pil = Image.open(requests.get(url, stream=True).raw).convert('RGB') | |
| return img_pil | |
| def _get_text_embedding(model, tokenizer, classes: list, device): | |
| prompts = [f'a photo of a {cls}.' for cls in classes] | |
| tokenized_prompts = tokenizer(prompts).to(device) | |
| text_embedding = model.encode_text(tokenized_prompts) | |
| text_embedding = F.normalize(text_embedding, dim=-1) | |
| return text_embedding.unsqueeze(0) | |
| # ---------- Function to convert logits to heatmaps ---------- | |
| def logits_to_heatmaps(logits, image_cv): | |
| logits = logits[0, 0].detach().cpu().numpy() | |
| logits = (logits * 255).astype('uint8') | |
| heat_map = cv2.applyColorMap(logits, cv2.COLORMAP_JET) | |
| viz = 0.4 * image_cv + 0.6 * heat_map | |
| viz = cv2.cvtColor(viz.astype('uint8'), cv2.COLOR_BGR2RGB) | |
| return viz | |
| # ---------- Main visualization function ---------- | |
| # @spaces.GPU | |
| def viz_func(url, image, text_query): | |
| image_torch = preprocess(image).unsqueeze(0).to(device) | |
| text_emb = _get_text_embedding(model, tokenizer, classes=[text_query], device=device) | |
| # ------- Get LeGrad output ------- | |
| logits_legrad = model.compute_legrad(image=image_torch, text_embedding=text_emb) | |
| # ------- Get Heatmpas ------- | |
| image_cv = cv2.cvtColor(np.array(image.resize((image_size, image_size))), cv2.COLOR_RGB2BGR) | |
| viz_legrad = logits_to_heatmaps(logits=logits_legrad, image_cv=image_cv) | |
| return viz_legrad | |
| inputs = [ | |
| gr.Textbox(label="Paste the url to the selected image"), | |
| gr.Image(type="pil", interactive=True, label='Select An Image'), | |
| gr.Textbox(label="Text query"), | |
| ] | |
| with gr.Blocks(css="#gradio-app-title { text-align: center; }") as demo: | |
| gr.Markdown( | |
| """ | |
| # **LeGrad: An Explainability Method for Vision Transformers via Feature Formation Sensitivity** | |
| ### This demo that showcases LeGrad method to visualize the important regions in an image that correspond to a given text query. | |
| The model used is OpenCLIP-ViT-B-16 (weights: `laion2b_s34b_b88k`) | |
| """ | |
| ) | |
| with gr.Row(): | |
| with gr.Column(): | |
| gr.Markdown('# Select An Image') | |
| selected_image = gr.Image(type="pil", interactive=True, label='') | |
| gr.Markdown('## Paste the url to the selected image') | |
| url_query = gr.Textbox(label="") | |
| gr.Markdown('# Create your Own query') | |
| text_query = gr.Textbox(label='') | |
| run_button = gr.Button(icon='https://cdn-icons-png.flaticon.com/512/3348/3348036.png') | |
| inputs[0].change(fn=change_to_url, outputs=inputs[1], inputs=inputs[0]) | |
| gr.Markdown('## LeGrad Explanation') | |
| le_grad_output = gr.Image(label='LeGrad') | |
| run_button.click(fn=viz_func, | |
| inputs=[url_query, selected_image, text_query], | |
| outputs=[le_grad_output]) | |
| with gr.Column(): | |
| gr.Markdown('# Select a Premade Example') | |
| gr.Examples( | |
| examples=[ | |
| ["assets/cats_remote_control.jpeg", "cat"], | |
| ["assets/cats_remote_control.jpeg", "remote control"], | |
| ["assets/la_baguette.webp", "la baguette"], | |
| ["assets/la_baguette.webp", "beret"], | |
| ["assets/pokemons.jpeg", "Pikachu"], | |
| ["assets/pokemons.jpeg", "Bulbasaur"], | |
| ["assets/pokemons.jpeg", "Charmander"], | |
| ["assets/pokemons.jpeg", "Pokemons"], | |
| ], | |
| inputs=[selected_image, text_query], | |
| outputs=[le_grad_output], | |
| fn=viz_func, | |
| label='' | |
| ) | |
| demo.queue() | |
| demo.launch() |