Spaces:
Sleeping
Sleeping
| import torch | |
| import numpy as np | |
| import torch.nn as nn | |
| import torchvision.transforms as transforms | |
| import matplotlib | |
| import matplotlib.pyplot as plt | |
| from PIL import Image | |
| import cv2 | |
| import gradio as gr | |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| from data_transforms import normal_transforms, no_shift_transforms, ig_transforms, modify_transforms | |
| from utils import overlay_heatmap, viz_map, show_image, deprocess, get_ssl_model, fig2img | |
| from methods import occlusion, pairwise_occlusion | |
| from methods import create_mixed_images, averaged_transforms, sailency, smooth_grad | |
| from methods import get_gradcam, get_interactioncam | |
| matplotlib.use('Agg') | |
| def load_model(model_name): | |
| global network, ssl_model, denorm | |
| if model_name == "simclrv2 (1X)": | |
| variant = '1x' | |
| network = 'simclrv2' | |
| denorm = False | |
| elif model_name == "simclrv2 (2X)": | |
| variant = '2x' | |
| network = 'simclrv2' | |
| denorm = False | |
| elif model_name == "Barlow Twins": | |
| network = 'barlow_twins' | |
| variant = None | |
| denorm = True | |
| ssl_model = get_ssl_model(network, variant) | |
| if network != 'simclrv2': | |
| global normal_transforms, no_shift_transforms, ig_transforms | |
| normal_transforms, no_shift_transforms, ig_transforms = modify_transforms(normal_transforms, no_shift_transforms, ig_transforms) | |
| return "Loaded Model Successfully" | |
| def load_or_augment_images(img1_input, img2_input, use_aug): | |
| global img_main, img1, img2 | |
| img_main = img1_input.convert('RGB') | |
| if use_aug: | |
| img1 = normal_transforms['pure'](img_main).unsqueeze(0).to(device) | |
| img2 = normal_transforms['aug'](img_main).unsqueeze(0).to(device) | |
| else: | |
| img1 = normal_transforms['pure'](img_main).unsqueeze(0).to(device) | |
| img2 = img2_input.convert('RGB') | |
| img2 = normal_transforms['pure'](img2).unsqueeze(0).to(device) | |
| similarity = "Similarity: {:.3f}".format(nn.CosineSimilarity(dim=-1)(ssl_model(img1), ssl_model(img2)).item()) | |
| fig, axs = plt.subplots(1, 2, figsize=(10,10)) | |
| np.vectorize(lambda ax:ax.axis('off'))(axs) | |
| axs[0].imshow(show_image(img1, denormalize = denorm)) | |
| axs[1].imshow(show_image(img2, denormalize = denorm)) | |
| plt.subplots_adjust(wspace=0.1, hspace = 0) | |
| pil_output = fig2img(fig) | |
| return pil_output, similarity | |
| def run_occlusion(w_size, stride): | |
| heatmap1, heatmap2 = occlusion(img1, img2, ssl_model, w_size = 64, stride = 8, batch_size = 32) | |
| heatmap1_po, heatmap2_po = pairwise_occlusion(img1, img2, ssl_model, batch_size = 32, erase_scale = (0.1, 0.3), erase_ratio = (1, 1.5), num_erases = 100) | |
| added_image1 = overlay_heatmap(img1, heatmap1, denormalize = denorm) | |
| added_image2 = overlay_heatmap(img2, heatmap2, denormalize = denorm) | |
| fig, axs = plt.subplots(2, 3, figsize=(20,10)) | |
| np.vectorize(lambda ax:ax.axis('off'))(axs) | |
| axs[0, 0].imshow(show_image(img1, denormalize = denorm)) | |
| axs[0, 1].imshow(added_image1) | |
| axs[0, 1].set_title("Conditional Occlusion") | |
| axs[0, 2].imshow((deprocess(img1, denormalize = denorm) * heatmap1_po[:,:,None]).astype('uint8')) | |
| axs[0, 2].set_title("Pairwise Occlusion") | |
| axs[1, 0].imshow(show_image(img2, denormalize = denorm)) | |
| axs[1, 1].imshow(added_image2) | |
| axs[1, 2].imshow((deprocess(img2, denormalize = denorm) * heatmap2_po[:,:,None]).astype('uint8')) | |
| plt.subplots_adjust(wspace=0, hspace = 0.01) | |
| pil_output = fig2img(fig) | |
| return pil_output | |
| def get_avg_trasforms(transform_type, add_noise, blur_output, guided): | |
| mixed_images = create_mixed_images(transform_type = transform_type, | |
| ig_transforms = ig_transforms, | |
| step = 0.1, | |
| img_path = img_main, | |
| add_noise = add_noise) | |
| # vanilla gradients (for comparison purposes) | |
| sailency1_van, sailency2_van = sailency(guided = guided, ssl_model = ssl_model, | |
| img1 = mixed_images[0], img2 = mixed_images[-1], | |
| blur_output = blur_output) | |
| # smooth gradients (for comparison purposes) | |
| sailency1_s, sailency2_s = smooth_grad(guided = guided, ssl_model = ssl_model, | |
| img1 = mixed_images[0], img2 = mixed_images[-1], | |
| blur_output = blur_output, steps = 50) | |
| # integrated transform | |
| sailency1, sailency2 = averaged_transforms(guided = guided, ssl_model = ssl_model, | |
| mixed_images = mixed_images, | |
| blur_output = blur_output) | |
| fig, axs = plt.subplots(2, 4, figsize=(20,10)) | |
| np.vectorize(lambda ax:ax.axis('off'))(axs) | |
| axs[0,0].imshow(show_image(mixed_images[0], denormalize = denorm)) | |
| axs[0,1].imshow(show_image(sailency1_van.detach(), squeeze = False).squeeze(), cmap = plt.cm.jet) | |
| axs[0,1].imshow(show_image(mixed_images[0], denormalize = denorm), alpha=0.5) | |
| axs[0,1].set_title("Vanilla Gradients") | |
| axs[0,2].imshow(show_image(sailency1_s.detach(), squeeze = False).squeeze(), cmap = plt.cm.jet) | |
| axs[0,2].imshow(show_image(mixed_images[0], denormalize = denorm), alpha=0.5) | |
| axs[0,2].set_title("Smooth Gradients") | |
| axs[0,3].imshow(show_image(sailency1.detach(), squeeze = False).squeeze(), cmap = plt.cm.jet) | |
| axs[0,3].imshow(show_image(mixed_images[0], denormalize = denorm), alpha=0.5) | |
| axs[0,3].set_title("Integrated Transform") | |
| axs[1,0].imshow(show_image(mixed_images[-1], denormalize = denorm)) | |
| axs[1,1].imshow(show_image(sailency2_van.detach(), squeeze = False).squeeze(), cmap = plt.cm.jet) | |
| axs[1,1].imshow(show_image(mixed_images[-1], denormalize = denorm), alpha=0.5) | |
| axs[1,2].imshow(show_image(sailency2_s.detach(), squeeze = False).squeeze(), cmap = plt.cm.jet) | |
| axs[1,2].imshow(show_image(mixed_images[-1], denormalize = denorm), alpha=0.5) | |
| axs[1,3].imshow(show_image(sailency2.detach(), squeeze = False).squeeze(), cmap = plt.cm.jet) | |
| axs[1,3].imshow(show_image(mixed_images[-1], denormalize = denorm), alpha=0.5) | |
| plt.subplots_adjust(wspace=0.02, hspace = 0.02) | |
| pil_output = fig2img(fig) | |
| return pil_output | |
| def get_cams(): | |
| gradcam1, gradcam2 = get_gradcam(ssl_model, img1, img2) | |
| intcam1_mean, intcam2_mean = get_interactioncam(ssl_model, img1, img2, reduction = 'mean') | |
| fig, axs = plt.subplots(2, 3, figsize=(20,8)) | |
| np.vectorize(lambda ax:ax.axis('off'))(axs) | |
| axs[0,0].imshow(show_image(img1[0], squeeze = False, denormalize = denorm)) | |
| axs[0,1].imshow(overlay_heatmap(img1, gradcam1, denormalize = denorm)) | |
| axs[0,1].set_title("Grad-CAM") | |
| axs[0,2].imshow(overlay_heatmap(img1, intcam1_mean, denormalize = denorm)) | |
| axs[0,2].set_title("IntCAM") | |
| axs[1,0].imshow(show_image(img2[0], squeeze = False, denormalize = denorm)) | |
| axs[1,1].imshow(overlay_heatmap(img2, gradcam2, denormalize = denorm)) | |
| axs[1,2].imshow(overlay_heatmap(img2, intcam2_mean, denormalize = denorm)) | |
| plt.subplots_adjust(wspace=0.01, hspace = 0.01) | |
| pil_output = fig2img(fig) | |
| return pil_output | |
| xai = gr.Blocks() | |
| with xai: | |
| gr.Markdown("<h1>Visualizing and Understanding Contrastive Learning, TIP Submission</h1>") | |
| gr.Markdown("The interface is simplified as much as possible with only necessary options to select for each method") | |
| gr.Markdown("<b>Due to the latency in Hugging Face machines (this demo is using the free CPU Basic plan with 2 CPUs), the methods are very slow. We advice to use a local machine or our Google Colab demo (link in the GitHub)</b>") | |
| with gr.Row(): | |
| model_name = gr.Dropdown(["simclrv2 (1X)", "simclrv2 (2X)", "Barlow Twins"], label="Choose Model and press \"Load Model\"") | |
| load_model_button = gr.Button("Load Model") | |
| status_or_similarity = gr.inputs.Textbox(label = "Status") | |
| with gr.Row(): | |
| gr.Markdown("You can either load two images or load a single image and augment it to get the second image (in that case please check the \"Use Augmentations\" checkbox). After that, please press on \"Show Images\". The similarity will be shown in the \"Status\" bar.") | |
| img1 = gr.Image(type='pil', label = "First Image") | |
| img2 = gr.Image(type='pil', label = "Second Image") | |
| with gr.Row(): | |
| use_aug = gr.Checkbox(value = False, label = "Use Augmentations") | |
| load_images_button = gr.Button("Show Images") | |
| gr.Markdown("Choose a method from the different tabs. You may leave the default options as they are and press on \"Run\" ") | |
| with gr.Row(): | |
| with gr.Column(): | |
| with gr.Tabs(): | |
| with gr.TabItem("Interaction-CAM"): | |
| cams_button = gr.Button("Get Heatmaps") | |
| with gr.TabItem("Perturbation Methods"): | |
| w_size = gr.Number(value = 64, label = "Occlusion Window Size", precision = 0) | |
| stride = gr.Number(value = 8, label = "Occlusion Stride", precision = 0) | |
| occlusion_button = gr.Button("Get Heatmap") | |
| with gr.TabItem("Averaged Transforms"): | |
| transform_type = gr.inputs.Radio(label="Data Augment", choices=['color_jitter', 'blur', 'grayscale', 'solarize', 'combine'], default="combine") | |
| add_noise = gr.Checkbox(value = True, label = "Add Noise") | |
| blur_output = gr.Checkbox(value = True, label = "Blur Output") | |
| guided = gr.Checkbox(value = True, label = "Guided Backprop") | |
| avgtransform_button = gr.Button("Get Saliency") | |
| with gr.Column(): | |
| output_image = gr.Image(type='pil', show_label = False) | |
| load_model_button.click(load_model, inputs = model_name, outputs = status_or_similarity) | |
| load_images_button.click(load_or_augment_images, inputs = [img1, img2, use_aug], outputs = [output_image, status_or_similarity]) | |
| occlusion_button.click(run_occlusion, inputs=[w_size,stride], outputs=output_image) | |
| avgtransform_button.click(get_avg_trasforms, inputs = [transform_type, add_noise, blur_output, guided], outputs = output_image) | |
| cams_button.click(get_cams, inputs = [], outputs = output_image) | |
| xai.launch() | |