Initial Commit

app.py

@@ -0,0 +1,98 @@
+import torch
+from matplotlib import pyplot as plt
+from PIL import Image
+import numpy as np
+from numpy import matlib as mb
+import torchvision.transforms as transforms
+import torchvision.models as models
+import skimage.transform
+import gradio as gr
+
+
+def compute_spatial_similarity(conv1, conv2):
+  """
+  Takes in the last convolutional layer from two images, computes the pooled output
+  feature, and then generates the spatial similarity map for both images.
+  """
+  conv1 = conv1.reshape(-1, 7*7).T
+  conv2 = conv2.reshape(-1, 7*7).T
+
+  pool1 = np.mean(conv1, axis=0)
+  pool2 = np.mean(conv2, axis=0)
+  out_sz = (int(np.sqrt(conv1.shape[0])),int(np.sqrt(conv1.shape[0])))
+  conv1_normed = conv1 / np.linalg.norm(pool1) / conv1.shape[0]
+  conv2_normed = conv2 / np.linalg.norm(pool2) / conv2.shape[0]
+  im_similarity = np.zeros((conv1_normed.shape[0], conv1_normed.shape[0]))
+
+  for zz in range(conv1_normed.shape[0]):
+    repPx = mb.repmat(conv1_normed[zz,:],conv1_normed.shape[0],1)
+    im_similarity[zz,:] = np.multiply(repPx,conv2_normed).sum(axis=1)
+  similarity1 = np.reshape(np.sum(im_similarity,axis=1),out_sz)
+  similarity2 = np.reshape(np.sum(im_similarity,axis=0),out_sz)
+  return similarity1, similarity2
+
+# Get Layer 4
+
+display_transform = transforms.Compose([transforms.Resize(256), transforms.CenterCrop((224, 224))])
+
+imagenet_transform = transforms.Compose(
+  [transforms.Resize(256),
+   transforms.CenterCrop((224, 224)), 
+   transforms.ToTensor(), 
+   transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])
+
+
+class Wrapper(torch.nn.Module):
+  def __init__(self, model):
+    super(Wrapper, self).__init__()
+    self.model = model
+    self.layer4_ouputs = None
+    def fw_hook(module, input, output):
+      self.layer4_ouputs = output
+    self.model.layer4.register_forward_hook(fw_hook)
+
+  def forward(self, input):
+    _ = self.model(input)
+    return self.layer4_ouputs
+    
+  def __repr__(self):
+    return "Wrapper"
+  
+def get_layer4(input_image):
+  l4_model = models.resnet50(pretrained=True)
+  l4_model = l4_model.cuda()
+  l4_model.eval();
+  wrapped_model = Wrapper(l4_model)
+
+  with torch.no_grad():
+    data = imagenet_transform(input_image).unsqueeze(0)
+    data = data.cuda()
+    reference_layer4 = wrapped_model(data)
+
+  return reference_layer4.data.to('cpu').numpy()
+
+# Visualization
+def visualize_similarities(image1, image2):
+  a = get_layer4(image1).squeeze()
+  b = get_layer4(image2).squeeze()
+  sim1, sim2 = compute_spatial_similarity(a, b)
+
+  fig, axes = plt.subplots(1, 2, figsize=(12, 5))
+  axes[0].imshow(display_transform(image1))
+  im1=axes[0].imshow(skimage.transform.resize(sim1, (224, 224)), alpha=0.6, cmap='jet')
+  # axes[0].colorbar()
+
+  axes[1].imshow(display_transform(image2))
+  im2=axes[1].imshow(skimage.transform.resize(sim2, (224, 224)), alpha=0.6, cmap='jet')
+  # axes[1].colorbar()
+
+  fig.colorbar(im1, ax=axes[0])
+  fig.colorbar(im2, ax=axes[1])
+  plt.tight_layout()
+  return fig
+
+# GRADIO APP
+iface = gr.Interface(fn=visualize_similarities, 
+                     inputs=[gr.inputs.Image(shape=(300, 300), type='pil'), 
+                             gr.inputs.Image(shape=(300, 300), type='pil')], outputs="plot")
+iface.launch()

requirements.txt

@@ -0,0 +1,8 @@
+gradio==2.4.5
+matplotlib==3.4.3
+numpy==1.21.2
+Pillow==8.4.0
+scikit_image==0.18.3
+skimage==0.0
+torch==1.10.0
+torchvision==0.11.1