195 lines
8.5 KiB
Python
195 lines
8.5 KiB
Python
import torch
|
|
from torch import nn
|
|
import torch.nn.functional as F
|
|
|
|
|
|
# coordinates system
|
|
# ------------------------------> [ x: range=-1.0~1.0; w: range=0~W ]
|
|
# | -----------------------------
|
|
# | | |
|
|
# | | |
|
|
# | | |
|
|
# | | image |
|
|
# | | |
|
|
# | | |
|
|
# | | |
|
|
# | |---------------------------|
|
|
# v
|
|
# [ y: range=-1.0~1.0; h: range=0~H ]
|
|
|
|
def simple_nms(scores, nms_radius: int):
|
|
""" Fast Non-maximum suppression to remove nearby points """
|
|
assert (nms_radius >= 0)
|
|
|
|
def max_pool(x):
|
|
return torch.nn.functional.max_pool2d(
|
|
x, kernel_size=nms_radius * 2 + 1, stride=1, padding=nms_radius)
|
|
|
|
zeros = torch.zeros_like(scores)
|
|
max_mask = scores == max_pool(scores)
|
|
|
|
for _ in range(2):
|
|
supp_mask = max_pool(max_mask.float()) > 0
|
|
supp_scores = torch.where(supp_mask, zeros, scores)
|
|
new_max_mask = supp_scores == max_pool(supp_scores)
|
|
max_mask = max_mask | (new_max_mask & (~supp_mask))
|
|
return torch.where(max_mask, scores, zeros)
|
|
|
|
|
|
def sample_descriptor(descriptor_map, kpts, bilinear_interp=False):
|
|
"""
|
|
:param descriptor_map: BxCxHxW
|
|
:param kpts: list, len=B, each is Nx2 (keypoints) [h,w]
|
|
:param bilinear_interp: bool, whether to use bilinear interpolation
|
|
:return: descriptors: list, len=B, each is NxD
|
|
"""
|
|
batch_size, channel, height, width = descriptor_map.shape
|
|
|
|
descriptors = []
|
|
for index in range(batch_size):
|
|
kptsi = kpts[index] # Nx2,(x,y)
|
|
|
|
if bilinear_interp:
|
|
descriptors_ = torch.nn.functional.grid_sample(descriptor_map[index].unsqueeze(0), kptsi.view(1, 1, -1, 2),
|
|
mode='bilinear', align_corners=True)[0, :, 0, :] # CxN
|
|
else:
|
|
kptsi = (kptsi + 1) / 2 * kptsi.new_tensor([[width - 1, height - 1]])
|
|
kptsi = kptsi.long()
|
|
descriptors_ = descriptor_map[index, :, kptsi[:, 1], kptsi[:, 0]] # CxN
|
|
|
|
descriptors_ = torch.nn.functional.normalize(descriptors_, p=2, dim=0)
|
|
descriptors.append(descriptors_.t())
|
|
|
|
return descriptors
|
|
|
|
|
|
class DKD(nn.Module):
|
|
def __init__(self, radius=2, top_k=0, scores_th=0.2, n_limit=20000):
|
|
"""
|
|
Args:
|
|
radius: soft detection radius, kernel size is (2 * radius + 1)
|
|
top_k: top_k > 0: return top k keypoints
|
|
scores_th: top_k <= 0 threshold mode: scores_th > 0: return keypoints with scores>scores_th
|
|
else: return keypoints with scores > scores.mean()
|
|
n_limit: max number of keypoint in threshold mode
|
|
"""
|
|
super().__init__()
|
|
self.radius = radius
|
|
self.top_k = top_k
|
|
self.scores_th = scores_th
|
|
self.n_limit = n_limit
|
|
self.kernel_size = 2 * self.radius + 1
|
|
self.temperature = 0.1 # tuned temperature
|
|
self.unfold = nn.Unfold(kernel_size=self.kernel_size, padding=self.radius)
|
|
|
|
# local xy grid
|
|
x = torch.linspace(-self.radius, self.radius, self.kernel_size)
|
|
# (kernel_size*kernel_size) x 2 : (w,h)
|
|
self.hw_grid = torch.stack(torch.meshgrid([x, x])).view(2, -1).t()[:, [1, 0]]
|
|
|
|
def detect_keypoints(self, scores_map, sub_pixel=True):
|
|
b, c, h, w = scores_map.shape
|
|
scores_nograd = scores_map.detach()
|
|
# nms_scores = simple_nms(scores_nograd, self.radius)
|
|
nms_scores = simple_nms(scores_nograd, 2)
|
|
|
|
# remove border
|
|
nms_scores[:, :, :self.radius + 1, :] = 0
|
|
nms_scores[:, :, :, :self.radius + 1] = 0
|
|
nms_scores[:, :, h - self.radius:, :] = 0
|
|
nms_scores[:, :, :, w - self.radius:] = 0
|
|
|
|
# detect keypoints without grad
|
|
if self.top_k > 0:
|
|
topk = torch.topk(nms_scores.view(b, -1), self.top_k)
|
|
indices_keypoints = topk.indices # B x top_k
|
|
else:
|
|
if self.scores_th > 0:
|
|
masks = nms_scores > self.scores_th
|
|
if masks.sum() == 0:
|
|
th = scores_nograd.reshape(b, -1).mean(dim=1) # th = self.scores_th
|
|
masks = nms_scores > th.reshape(b, 1, 1, 1)
|
|
else:
|
|
th = scores_nograd.reshape(b, -1).mean(dim=1) # th = self.scores_th
|
|
masks = nms_scores > th.reshape(b, 1, 1, 1)
|
|
masks = masks.reshape(b, -1)
|
|
|
|
indices_keypoints = [] # list, B x (any size)
|
|
scores_view = scores_nograd.reshape(b, -1)
|
|
for mask, scores in zip(masks, scores_view):
|
|
indices = mask.nonzero(as_tuple=False)[:, 0]
|
|
if len(indices) > self.n_limit:
|
|
kpts_sc = scores[indices]
|
|
sort_idx = kpts_sc.sort(descending=True)[1]
|
|
sel_idx = sort_idx[:self.n_limit]
|
|
indices = indices[sel_idx]
|
|
indices_keypoints.append(indices)
|
|
|
|
keypoints = []
|
|
scoredispersitys = []
|
|
kptscores = []
|
|
if sub_pixel:
|
|
# detect soft keypoints with grad backpropagation
|
|
patches = self.unfold(scores_map) # B x (kernel**2) x (H*W)
|
|
self.hw_grid = self.hw_grid.to(patches) # to device
|
|
for b_idx in range(b):
|
|
patch = patches[b_idx].t() # (H*W) x (kernel**2)
|
|
indices_kpt = indices_keypoints[b_idx] # one dimension vector, say its size is M
|
|
patch_scores = patch[indices_kpt] # M x (kernel**2)
|
|
|
|
# max is detached to prevent undesired backprop loops in the graph
|
|
max_v = patch_scores.max(dim=1).values.detach()[:, None]
|
|
x_exp = ((patch_scores - max_v) / self.temperature).exp() # M * (kernel**2), in [0, 1]
|
|
|
|
# \frac{ \sum{(i,j) \times \exp(x/T)} }{ \sum{\exp(x/T)} }
|
|
xy_residual = x_exp @ self.hw_grid / x_exp.sum(dim=1)[:, None] # Soft-argmax, Mx2
|
|
|
|
hw_grid_dist2 = torch.norm((self.hw_grid[None, :, :] - xy_residual[:, None, :]) / self.radius,
|
|
dim=-1) ** 2
|
|
scoredispersity = (x_exp * hw_grid_dist2).sum(dim=1) / x_exp.sum(dim=1)
|
|
|
|
# compute result keypoints
|
|
keypoints_xy_nms = torch.stack([indices_kpt % w, indices_kpt // w], dim=1) # Mx2
|
|
keypoints_xy = keypoints_xy_nms + xy_residual
|
|
keypoints_xy = keypoints_xy / keypoints_xy.new_tensor(
|
|
[w - 1, h - 1]) * 2 - 1 # (w,h) -> (-1~1,-1~1)
|
|
|
|
kptscore = torch.nn.functional.grid_sample(scores_map[b_idx].unsqueeze(0),
|
|
keypoints_xy.view(1, 1, -1, 2),
|
|
mode='bilinear', align_corners=True)[0, 0, 0, :] # CxN
|
|
|
|
keypoints.append(keypoints_xy)
|
|
scoredispersitys.append(scoredispersity)
|
|
kptscores.append(kptscore)
|
|
else:
|
|
for b_idx in range(b):
|
|
indices_kpt = indices_keypoints[b_idx] # one dimension vector, say its size is M
|
|
keypoints_xy_nms = torch.stack([indices_kpt % w, indices_kpt // w], dim=1) # Mx2
|
|
keypoints_xy = keypoints_xy_nms / keypoints_xy_nms.new_tensor(
|
|
[w - 1, h - 1]) * 2 - 1 # (w,h) -> (-1~1,-1~1)
|
|
kptscore = torch.nn.functional.grid_sample(scores_map[b_idx].unsqueeze(0),
|
|
keypoints_xy.view(1, 1, -1, 2),
|
|
mode='bilinear', align_corners=True)[0, 0, 0, :] # CxN
|
|
keypoints.append(keypoints_xy)
|
|
scoredispersitys.append(None)
|
|
kptscores.append(kptscore)
|
|
|
|
return keypoints, scoredispersitys, kptscores
|
|
|
|
def forward(self, scores_map, descriptor_map, sub_pixel=False):
|
|
"""
|
|
:param scores_map: Bx1xHxW
|
|
:param descriptor_map: BxCxHxW
|
|
:param sub_pixel: whether to use sub-pixel keypoint detection
|
|
:return: kpts: list[Nx2,...]; kptscores: list[N,....] normalised position: -1.0 ~ 1.0
|
|
"""
|
|
keypoints, scoredispersitys, kptscores = self.detect_keypoints(scores_map,
|
|
sub_pixel)
|
|
|
|
descriptors = sample_descriptor(descriptor_map, keypoints, sub_pixel)
|
|
|
|
# keypoints: B M 2
|
|
# descriptors: B M D
|
|
# scoredispersitys:
|
|
return keypoints, descriptors, kptscores, scoredispersitys
|