added heatmap visualizer

Author: jerome-revaud

imgs/peppers.png

@@ -0,0 +1,122 @@
+import pdb
+import os
+import sys
+import tqdm
+
+import numpy as np
+import torch
+
+from PIL import Image
+from matplotlib import pyplot as pl; pl.ion()
+from scipy.ndimage import uniform_filter
+smooth = lambda arr: uniform_filter(arr, 3)
+
+def transparent(img, alpha, cmap, **kw):
+    from matplotlib.colors import Normalize
+    colored_img = cmap(Normalize(clip=True,**kw)(img))
+    colored_img[:,:,-1] = alpha
+    return colored_img
+
+from tools import common
+from tools.dataloader import norm_RGB
+from nets.patchnet import *
+from extract import NonMaxSuppression
+
+
+if __name__ == '__main__':
+    import argparse
+    parser = argparse.ArgumentParser("Visualize the patch detector and descriptor")
+    
+    parser.add_argument("--img", type=str, default="imgs/brooklyn.png")
+    parser.add_argument("--resize", type=int, default=512)
+    parser.add_argument("--out", type=str, default="viz.png")
+
+    parser.add_argument("--checkpoint", type=str, required=True, help='network path')
+    parser.add_argument("--net", type=str, default="", help='network command')
+
+    parser.add_argument("--max-kpts", type=int, default=200)
+    parser.add_argument("--reliability-thr", type=float, default=0.8)
+    parser.add_argument("--repeatability-thr", type=float, default=0.7)
+    parser.add_argument("--border", type=int, default=20,help='rm keypoints close to border')
+
+    parser.add_argument("--gpu", type=int, nargs='+', required=True, help='-1 for CPU')
+    parser.add_argument("--dbg", type=str, nargs='+', default=(), help='debug options')
+    
+    args = parser.parse_args()
+    args.dbg = set(args.dbg)
+    
+    iscuda = common.torch_set_gpu(args.gpu)
+    device = torch.device('cuda' if iscuda else 'cpu')
+
+    # create network
+    checkpoint = torch.load(args.checkpoint, lambda a,b:a)
+    args.net = args.net or checkpoint['net']
+    print("\n>> Creating net = " + args.net) 
+    net = eval(args.net)
+    net.load_state_dict({k.replace('module.',''):v for k,v in checkpoint['state_dict'].items()})
+    if iscuda: net = net.cuda()
+    print(f" ( Model size: {common.model_size(net)/1000:.0f}K parameters )")
+
+    img = Image.open(args.img).convert('RGB')
+    if args.resize: img.thumbnail((args.resize,args.resize))
+    img = np.asarray(img)
+        
+    detector = NonMaxSuppression(
+        rel_thr = args.reliability_thr, 
+        rep_thr = args.repeatability_thr)
+
+    with torch.no_grad():
+        print(">> computing features...")
+        res = net(imgs=[norm_RGB(img).unsqueeze(0).to(device)])
+        rela = res.get('reliability')
+        repe = res.get('repeatability')
+        kpts = detector(**res).T[:,[1,0]]
+        kpts = kpts[repe[0][0,0][kpts[:,1],kpts[:,0]].argsort()[-args.max_kpts:]]
+
+    fig = pl.figure("viz")
+    kw = dict(cmap=pl.cm.RdYlGn, vmax=1)
+    crop = (slice(args.border,-args.border or 1),)*2
+    
+    if 'reliability' in args.dbg:
+    
+        ax1 = pl.subplot(131)
+        pl.imshow(img[crop], cmap=pl.cm.gray)
+        pl.xticks(()); pl.yticks(())
+
+        pl.subplot(132)
+        pl.imshow(img[crop], cmap=pl.cm.gray, alpha=0)
+        pl.xticks(()); pl.yticks(())
+
+        x,y = kpts[:,0:2].cpu().numpy().T - args.border
+        pl.plot(x,y,'+',c=(0,1,0),ms=10, scalex=0, scaley=0)
+
+        ax1 = pl.subplot(133)
+        rela = rela[0][0,0].cpu().numpy()
+        pl.imshow(rela[crop], cmap=pl.cm.RdYlGn, vmax=1, vmin=0.9)
+        pl.xticks(()); pl.yticks(())
+
+    else:
+        ax1 = pl.subplot(131)
+        pl.imshow(img[crop], cmap=pl.cm.gray)
+        pl.xticks(()); pl.yticks(())
+
+        x,y = kpts[:,0:2].cpu().numpy().T - args.border
+        pl.plot(x,y,'+',c=(0,1,0),ms=10, scalex=0, scaley=0)
+
+        pl.subplot(132)
+        pl.imshow(img[crop], cmap=pl.cm.gray)
+        pl.xticks(()); pl.yticks(())
+        c = repe[0][0,0].cpu().numpy()
+        pl.imshow(transparent(smooth(c)[crop], 0.5, vmin=0, **kw))
+
+        ax1 = pl.subplot(133)
+        pl.imshow(img[crop], cmap=pl.cm.gray)
+        pl.xticks(()); pl.yticks(())
+        rela = rela[0][0,0].cpu().numpy()
+        pl.imshow(transparent(rela[crop], 0.5, vmin=0.9, **kw))
+
+    pl.gcf().set_size_inches(9, 2.73)
+    pl.subplots_adjust(0.01,0.01,0.99,0.99,hspace=0.1)
+    pl.savefig(args.out)
+    pdb.set_trace()
+