app.py

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