Spaces:
Running
Running
| import gradio as gr | |
| import numpy as np | |
| import skimage.transform | |
| import torch | |
| import torchvision.models as models | |
| import torchvision.transforms as transforms | |
| from matplotlib import pyplot as plt | |
| from numpy import matlib as mb | |
| from PIL import Image | |
| import csv | |
| import sys | |
| csv.field_size_limit(sys.maxsize) | |
| 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.5, cmap="jet" | |
| ) | |
| # axes[0].colorbar() | |
| axes[1].imshow(display_transform(image2)) | |
| im2 = axes[1].imshow( | |
| skimage.transform.resize(sim2, (224, 224)), alpha=0.5, 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.Image(type="pil"), | |
| gr.Image(type="pil"), | |
| ], | |
| allow_flagging="never", | |
| outputs=[gr.Plot(type="matplotlib")], | |
| cache_examples=True, | |
| enable_queue=False, | |
| examples=[ | |
| [ | |
| "./examples/Red_Winged_Blackbird_0012_6015.jpg", | |
| "./examples/Red_Winged_Blackbird_0025_5342.jpg", | |
| ] | |
| ], | |
| ) | |
| iface.launch() | |