def __call__(self, logdepth, sampled_rld, pts2ddict):
sample_rld = sampled_rld.expand(
[self.bz, -1, self.featureh, self.featurew])
sample_ld = logdepth.expand([-1, self.nedges, -1, -1]) + sample_rld
sample_d = torch.exp(sample_ld)
infkey = "{}_{}_{}".format(self.bz, self.featureh, self.featurew)
xx, yy, ones = pts2ddict[infkey]
xx = xx.expand([-1, self.nedges, -1, -1])
yy = yy.expand([-1, self.nedges, -1, -1])
ones = ones.expand([-1, self.nedges, -1, -1])
pts3d = torch.stack([xx * sample_d, yy * sample_d, sample_d, ones],
dim=-1).unsqueeze(-1)
pts2dp = self.projMimg.unsqueeze(1).expand(
[-1, self.nedges, -1, -1, -1, -1]) @ pts3d
pxx, pyy, pzz, _ = torch.split(pts2dp, 1, dim=4)
sign = pzz.sign()
sign[sign == 0] = 1
pzz = torch.clamp(torch.abs(pzz), min=1e-20) * sign
pxx = (pxx / pzz).squeeze(-1).squeeze(-1)
pyy = (pyy / pzz).squeeze(-1).squeeze(-1)
supressval = torch.ones_like(pxx) * (-100)
inboundmask = (pzz > 1e-20).squeeze(-1).squeeze(-1).float()
pxx = inboundmask * pxx + supressval * (1 - inboundmask)
pyy = inboundmask * pyy + supressval * (1 - inboundmask)
sample_px = (pxx / self.featurew - 0.5) * 2
sample_py = (pyy / self.featureh - 0.5) * 2
sample_px = sample_px.permute([0, 2, 3, 1]).contiguous().view(
self.bz * self.featureh * self.featurew, 1, self.nedges)
sample_py = sample_py.permute([0, 2, 3, 1]).contiguous().view(
self.bz * self.featureh * self.featurew, 1, self.nedges)
local_sample_pts2d = torch.stack([sample_px, sample_py], dim=-1)
out = bilinear_sampler(self.corr, local_sample_pts2d)
out = out.view(self.bz, self.featureh, self.featurew, -1)
return out.permute(0, 3, 1,
2).contiguous().float(), local_sample_pts2d.view([
self.bz, self.featureh, self.featurew,
self.nedges, 2
])
def __call__(self, coords, flowmag, ratio):
r = self.radius
coords = coords.permute(0, 2, 3, 1)
batch, h1, w1, _ = coords.shape
corr = self.corr
dx = torch.linspace(-r, r, 2*r+1)
dy = torch.linspace(-r, r, 2*r+1)
delta = torch.stack(torch.meshgrid(dy, dx), axis=-1).to(coords.device)
centroid_lvl = coords.reshape(batch*h1*w1, 1, 1, 2)
delta_lvl = delta.view(1, 2*r+1, 2*r+1, 2).expand([batch*h1*w1, -1, -1, -1]) * flowmag.view(batch*h1*w1, 1, 1, 1).expand([-1, 2*r+1, 2*r+1, 2]) * ratio
coords_lvl = centroid_lvl + delta_lvl
corr = bilinear_sampler(corr, coords_lvl)
corr = corr.view(batch, h1, w1, -1)
corr = corr.permute(0, 3, 1, 2).contiguous().float()
coords_lvl = coords_lvl.view(batch, h1, w1, 2*r+1, 2*r+1, 2)
delta_lvl = delta_lvl.view(batch, h1, w1, (2*r+1) * (2*r+1), 2)
return corr, coords_lvl, delta_lvl
def __call__(self, coords):
r = self.radius
coords = coords.permute(0, 2, 3, 1)
batch, h1, w1, _ = coords.shape
out_pyramid = []
for i in range(self.num_levels):
corr = self.corr_pyramid[i]
dx = torch.linspace(-r, r, 2 * r + 1)
dy = torch.linspace(-r, r, 2 * r + 1)
delta = torch.stack(torch.meshgrid(dy, dx),
axis=-1).to(coords.device)
centroid_lvl = coords.reshape(batch * h1 * w1, 1, 1, 2) / 2**i
delta_lvl = delta.view(1, 2 * r + 1, 2 * r + 1, 2)
coords_lvl = centroid_lvl + delta_lvl
corr = bilinear_sampler(corr, coords_lvl)
corr = corr.view(batch, h1, w1, -1)
out_pyramid.append(corr)
out = torch.cat(out_pyramid, dim=-1)
return out.permute(0, 3, 1, 2).contiguous().float()