|
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") |
|
|
|
|
|
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() |
|
|
