gim-online / third_party /r2d2 /viz_heatmaps.py
Vincentqyw
fix: roma
8b973ee
raw
history blame
4.3 kB
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()