Upload app.py

app.py

@@ -0,0 +1,283 @@
+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
+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, occlusion_context_agnositc, pairwise_occlusion
+from methods import get_difference
+from methods import create_mixed_images, averaged_transforms, sailency, smooth_grad 
+from methods import get_sample_dataset, pixel_invariance, 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_ca, heatmap2_ca = occlusion_context_agnositc(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)
+  added_image1_ca = overlay_heatmap(img1, heatmap1_ca, denormalize = denorm)
+  added_image2_ca = overlay_heatmap(img2, heatmap2_ca, denormalize = denorm)
+
+  fig, axs = plt.subplots(2, 4, 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(added_image1_ca)
+  axs[0, 2].set_title("CA Cond. Occlusion")
+  axs[0, 3].imshow((deprocess(img1, denormalize = denorm) * heatmap1_po[:,:,None]).astype('uint8'))
+  axs[0, 3].set_title("Pairwise Occlusion")
+  axs[1, 0].imshow(show_image(img2, denormalize = denorm))
+  axs[1, 1].imshow(added_image2)
+  axs[1, 2].imshow(added_image2_ca)
+  axs[1, 3].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_model_difference(later):
+
+  imagenet_images, ssl_images = get_difference(ssl_model = ssl_model, baseline = 'imagenet', image = img2, lr = 1e4, 
+                                              l2_weight = 0.1, alpha_weight = 1e-7, alpha_power = 6, tv_weight = 1e-8, 
+                                              init_scale = 0.1, network = network)
+
+  fig, axs = plt.subplots(3, 3, figsize=(10,10))
+  np.vectorize(lambda ax:ax.axis('off'))(axs)
+
+  for aa, (in_img, ssl_img) in enumerate(zip(imagenet_images, ssl_images)):
+      axs[aa,0].imshow(deprocess(img2, denormalize = denorm))
+      axs[aa,1].imshow(deprocess(in_img))
+      axs[aa,2].imshow(deprocess(ssl_img))
+      
+  axs[0,0].set_title("Original Image")
+  axs[0,1].set_title("Synthesized (cls)")
+  axs[0,2].set_title("Synthesized (contastive)")
+
+  plt.subplots_adjust(wspace=0.01, 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')
+  intcam1_maxmax, intcam2_maxmax = get_interactioncam(ssl_model, img1, img2, reduction = 'max', grad_interact = True)
+  intcam1_attnmax, intcam2_attnmax = get_interactioncam(ssl_model, img1, img2, reduction = 'attn', grad_interact = True)
+
+  fig, axs = plt.subplots(2, 5, 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 Mean")
+  axs[0,3].imshow(overlay_heatmap(img1, intcam1_maxmax, denormalize = denorm))
+  axs[0,3].set_title("IntCAM Max + IntGradMax")
+  axs[0,4].imshow(overlay_heatmap(img1, intcam1_attnmax, denormalize = denorm))
+  axs[0,4].set_title("IntCAM Attn + IntGradMax")
+
+  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))
+  axs[1,3].imshow(overlay_heatmap(img2, intcam2_maxmax, denormalize = denorm))
+  axs[1,4].imshow(overlay_heatmap(img2, intcam2_attnmax, denormalize = denorm))
+
+  plt.subplots_adjust(wspace=0.01, hspace = 0.01)
+  pil_output = fig2img(fig)
+  return pil_output
+
+def get_pixel_invariance():
+
+  data_samples1, data_samples2, data_labels, labels_invariance = get_sample_dataset(img_path = img_main, 
+                                                                                    num_augments = 1000, 
+                                                                                    batch_size =  32, 
+                                                                                    no_shift_transforms = no_shift_transforms, 
+                                                                                    ssl_model = ssl_model, 
+                                                                                    n_components = 10)
+
+  inv_heatmap = pixel_invariance(data_samples1 = data_samples1, data_samples2 = data_samples2, data_labels = data_labels,
+                                labels_invariance = labels_invariance, resize_transform = transforms.Resize, size = 64, 
+                                epochs = 1000, learning_rate = 0.1, l1_weight = 0.2, zero_small_values = True, 
+                                blur_output = True, nmf_weight = 0)
+
+  inv_heatmap_nmf = pixel_invariance(data_samples1 = data_samples1, data_samples2 = data_samples2, data_labels = data_labels,
+                                    labels_invariance = labels_invariance, resize_transform = transforms.Resize, size = 64, 
+                                    epochs = 100, learning_rate = 0.1, l1_weight = 0.2, zero_small_values = True, 
+                                    blur_output = True, nmf_weight = 1)
+
+  fig, axs = plt.subplots(1, 2, figsize=(10,5))
+  np.vectorize(lambda ax:ax.axis('off'))(axs)
+
+  axs[0].imshow(viz_map(img_main, inv_heatmap))
+  axs[0].set_title("Heatmap w/o NMF")
+  axs[1].imshow(viz_map(img_main, inv_heatmap_nmf))
+  axs[1].set_title("Heatmap w/ NMF")
+  plt.subplots_adjust(wspace=0.01, hspace = 0.01)
+
+  pil_output = fig2img(fig)
+  return pil_output
+
+xai = gr.Blocks()
+
+with xai:
+  gr.Markdown("<h1>Methods for Explaining Contrastive Learning, CVPR 2023 Submission</h1>")
+  gr.Markdown("The interface is simplified as much as possible with only necessary options to select for each method. Please use our Google Colab demo for more flexibility.")
+  
+  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\" button). After that, please press on \"Show Images\"")
+    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.TabItem("Pixel Invariance"):
+          gr.Markdown("Note: Invariance map will be obtained for the first image")
+          pixel_invariance_button = gr.Button("Get Invariance Map")
+        with gr.TabItem("Image Synthesization"):
+          baseline = gr.inputs.Radio(label="Compare With", choices=["Supervised Classification"], default="Supervised Classification")
+          modeldiff_button = gr.Button("Compare")
+
+    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)
+  modeldiff_button.click(get_model_difference, inputs = baseline, 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) 
+  pixel_invariance_button.click(get_pixel_invariance, inputs = [], outputs = output_image)
+
+xai.launch()
+