naver/SuperFeatures
4
1import gradio as gr2 3import cv24 5import torch6import torch.utils.data as data7from torchvision import transforms8from torch import nn9import torch.nn.functional as F10 11import matplotlib.pyplot as plt12from matplotlib import cm13from matplotlib import colors14from mpl_toolkits.axes_grid1 import ImageGrid15 16import fire_network17 18import numpy as np19 20from PIL import Image21 22# Possible Scales for multiscale inference23scales = [2.0, 1.414, 1.0, 0.707, 0.5, 0.353, 0.25] 24 25device = 'cpu'26 27# Load nets28state = torch.load('fire.pth', map_location='cpu')29state['net_params']['pretrained'] = None # no need for imagenet pretrained model30net_sfm = fire_network.init_network(**state['net_params']).to(device)31net_sfm.load_state_dict(state['state_dict'])32dim_red_params_dict = {}33for name, param in net_sfm.named_parameters():34 if 'dim_reduction' in name:35 dim_red_params_dict[name] = param36 37 38state2 = torch.load('fire_imagenet.pth', map_location='cpu')39state2['net_params'] = state['net_params']40state2['state_dict'] = dict(state2['state_dict'], **dim_red_params_dict);41net_imagenet = fire_network.init_network(**state['net_params']).to(device)42net_imagenet.load_state_dict(state2['state_dict'], strict=False)43 44transform = transforms.Compose([45 transforms.Resize(1024),46 transforms.ToTensor(), 47 transforms.Normalize(**dict(zip(["mean", "std"], net_sfm.runtime['mean_std'])))48 ])49 50 51def match(query_feat, pos_feat, LoweRatioTh=0.9):52 # first perform reciprocal nn53 dist = torch.cdist(query_feat, pos_feat)54 # print('dist.size',dist.size())55 best1 = torch.argmin(dist, dim=1)56 best2 = torch.argmin(dist, dim=0)57 # print('best2.size',best2.size())58 arange = torch.arange(best2.size(0))59 reciprocal = best1[best2]==arange60 # check Lowe ratio test61 dist2 = dist.clone()62 dist2[best2,arange] = float('Inf')63 dist2_second2 = torch.argmin(dist2, dim=0)64 ratio1to2 = dist[best2,arange] / dist2_second265 valid = torch.logical_and(reciprocal, ratio1to2<=LoweRatioTh)66 pindices = torch.where(valid)[0]67 qindices = best2[pindices]68 # keep only the ones with same indices 69 valid = pindices==qindices70 return pindices[valid]71 72 73def clear_figures():74 plt.figure().clear()75 plt.close()76 plt.cla()77 plt.clf()78 79 80 81 82def generate_matching_superfeatures(83 im1, im2, 84 Imagenet_model=False, 85 scale_id=6, threshold=50, 86 random_mode=False, sf_ids=''): #, only_matching=True):87 # print('im1:', im1.size)88 # print('im2:', im2.size)89 90 clear_figures()91 col = plt.get_cmap('tab10')92 93 net = net_sfm94 if Imagenet_model:95 net = net_imagenet96 97 im1_tensor = transform(im1).unsqueeze(0)98 im2_tensor = transform(im2).unsqueeze(0)99 100 im1_cv = np.array(im1)[:, :, ::-1].copy() 101 im2_cv = np.array(im2)[:, :, ::-1].copy() 102 103 # extract features104 with torch.no_grad():105 output1 = net.get_superfeatures(im1_tensor.to(device), scales=[scales[scale_id]])106 feats1 = output1[0][0]107 attns1 = output1[1][0]108 strenghts1 = output1[2][0]109 110 output2 = net.get_superfeatures(im2_tensor.to(device), scales=[scales[scale_id]])111 feats2 = output2[0][0]112 attns2 = output2[1][0]113 strenghts2 = output2[2][0]114 115 feats1n = F.normalize(torch.t(torch.squeeze(feats1)), dim=1)116 feats2n = F.normalize(torch.t(torch.squeeze(feats2)), dim=1)117 ind_match = match(feats1n, feats2n)118 119 # which sf 120 sf_idx_ = []121 n_sf_ids = 10122 if random_mode or sf_ids == '':123 sf_idx_ = np.random.randint(256, size=n_sf_ids)124 else:125 sf_idx_ = map(int, sf_ids.strip().split(','))126 127 # only_matching:128 if random_mode:129 sf_idx_ = [int(jj) for jj in ind_match[np.random.randint(len(list(ind_match)), size=n_sf_ids)].numpy()]130 sf_idx_ = list( dict.fromkeys(sf_idx_) )131 else:132 sf_idx_ = [i for i in sf_idx_ if i in list(ind_match)]133 134 n_sf_ids = len(sf_idx_)135 136 # Store all binary SF att maps to show them all at once in the end137 all_att_bin1 = []138 all_att_bin2 = []139 for n, i in enumerate(sf_idx_):140 att_heat = np.array(attns1[0,i,:,:].numpy(), dtype=np.float32)141 att_heat = np.uint8(att_heat / np.max(att_heat[:]) * 255.0)142 att_heat_bin = np.where(att_heat>threshold, 255, 0)143 all_att_bin1.append(att_heat_bin)144 145 att_heat = np.array(attns2[0,i,:,:].numpy(), dtype=np.float32)146 att_heat = np.uint8(att_heat / np.max(att_heat[:]) * 255.0)147 att_heat_bin = np.where(att_heat>threshold, 255, 0)148 all_att_bin2.append(att_heat_bin)149 150 151 fin_img = []152 img1rsz = np.copy(im1_cv)153 for j, att in enumerate(all_att_bin1):154 att = cv2.resize(att, im1.size, interpolation=cv2.INTER_NEAREST)155 mask2d = zip(*np.where(att==255))156 for m,n in mask2d:157 col_ = col.colors[j]158 col_ = 255*np.array(colors.to_rgba(col_))[:3]159 img1rsz[m,n, :] = col_[::-1] 160 161 img2rsz = np.copy(im2_cv)162 for j, att in enumerate(all_att_bin2):163 att = cv2.resize(att, im2.size, interpolation=cv2.INTER_NEAREST)164 mask2d = zip(*np.where(att==255))165 for m,n in mask2d:166 col_ = col.colors[j]167 col_ = 255*np.array(colors.to_rgba(col_))[:3]168 img2rsz[m,n, :] = col_[::-1] 169 170 fig1 = plt.figure(1)171 plt.imshow(cv2.cvtColor(img1rsz, cv2.COLOR_BGR2RGB))172 ax1 = plt.gca()173 ax1.axis('off')174 plt.tight_layout() 175 176 fig2 = plt.figure(2)177 plt.imshow(cv2.cvtColor(img2rsz, cv2.COLOR_BGR2RGB))178 ax2 = plt.gca()179 ax2.axis('off')180 plt.tight_layout() 181 182 f = lambda m,c: plt.plot([],[],marker=m, color=c, ls="none")[0]183 handles = [f("s", col.colors[i]) for i in range(n_sf_ids)]184 fig_leg = plt.figure(3)185 legend = plt.legend(handles, sf_idx_, framealpha=1, frameon=False, facecolor='w',fontsize=25, loc="center")186 ax3 = plt.gca()187 ax3.axis('off')188 plt.tight_layout() 189 190 im1 = None191 im2 = None192 return fig1, fig2, fig_leg193 194 195# GRADIO APP196title = "Visualizing Super-features"197description = "This is a visualization demo for the ICLR 2022 paper <b><a href='https://github.com/naver/fire' target='_blank'>Learning Super-Features for Image Retrieval</a></p></b>" 198article = "<p style='text-align: center'><a href='https://github.com/naver/fire' target='_blank'>Original Github Repo</a></p>"199 200iface = gr.Interface(201 fn=generate_matching_superfeatures,202 inputs=[203 gr.inputs.Image(shape=(1024, 1024), type="pil", label="First Image"),204 gr.inputs.Image(shape=(1024, 1024), type="pil", label="Second Image"),205 gr.inputs.Checkbox(default=False, label="ImageNet Model (Default: SfM-120k)"),206 gr.inputs.Slider(minimum=0, maximum=6, step=1, default=4, label="Scale"),207 gr.inputs.Slider(minimum=0, maximum=255, step=25, default=150, label="Binarization Threshold"),208 gr.inputs.Checkbox(default=True, label="Show random (matching) SFs"),209 gr.inputs.Textbox(lines=1, default="", label="...or show specific SF IDs:", optional=True),210 ],211 outputs=[212 gr.outputs.Image(type="plot", label="First Image SFs"),213 gr.outputs.Image(type="plot", label="Second Image SFs"),214 gr.outputs.Image(type="plot", label="SF legend")],215 title=title,216 theme='peach',217 layout="horizontal",218 description=description,219 article=article,220 examples=[221 ["chateau_1.png", "chateau_2.png", False, 3, 150, False, '170,15,25,63,193,125,92,214,107'],222 ["areopoli1.jpeg", "areopoli2.jpeg", False, 4, 150, False, '205,2,163,130'],223 ["jaipur1.jpeg", "jaipur2.jpeg", False, 4, 50, False, '51,206,216,49,27'],224 ["basil1.jpeg", "basil2.jpeg", True, 4, 100, False, '75,152,19,36,156'],225 ["mill1.jpeg", "mill2.jpeg", False, 4, 100, False, '177,88,170,190,151,155'],226 ]227)228iface.launch(enable_queue=True)