Spaces:
Sleeping
Sleeping
File size: 3,690 Bytes
20ae9ff |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 |
import torch
import torch.nn.functional as F
from .lietorch import SE3, Sim3
MIN_DEPTH = 0.2
def extract_intrinsics(intrinsics):
return intrinsics[...,None,None,:].unbind(dim=-1)
def coords_grid(ht, wd, **kwargs):
y, x = torch.meshgrid(
torch.arange(ht).to(**kwargs).float(),
torch.arange(wd).to(**kwargs).float())
return torch.stack([x, y], dim=-1)
def iproj(patches, intrinsics):
""" inverse projection """
x, y, d = patches.unbind(dim=2)
fx, fy, cx, cy = intrinsics[...,None,None].unbind(dim=2)
i = torch.ones_like(d)
xn = (x - cx) / fx
yn = (y - cy) / fy
X = torch.stack([xn, yn, i, d], dim=-1)
return X
def proj(X, intrinsics, depth=False):
""" projection """
X, Y, Z, W = X.unbind(dim=-1)
fx, fy, cx, cy = intrinsics[...,None,None].unbind(dim=2)
# d = 0.01 * torch.ones_like(Z)
# d[Z > 0.01] = 1.0 / Z[Z > 0.01]
# d = torch.ones_like(Z)
# d[Z.abs() > 0.1] = 1.0 / Z[Z.abs() > 0.1]
d = 1.0 / Z.clamp(min=0.1)
x = fx * (d * X) + cx
y = fy * (d * Y) + cy
if depth:
return torch.stack([x, y, d], dim=-1)
return torch.stack([x, y], dim=-1)
def transform(poses, patches, intrinsics, ii, jj, kk, depth=False, valid=False, jacobian=False, tonly=False):
""" projective transform """
# backproject
X0 = iproj(patches[:,kk], intrinsics[:,ii])
# transform
Gij = poses[:, jj] * poses[:, ii].inv()
if tonly:
Gij[...,3:] = torch.as_tensor([0,0,0,1], device=Gij.device)
X1 = Gij[:,:,None,None] * X0
# project
x1 = proj(X1, intrinsics[:,jj], depth)
if jacobian:
p = X1.shape[2]
X, Y, Z, H = X1[...,p//2,p//2,:].unbind(dim=-1)
o = torch.zeros_like(H)
i = torch.zeros_like(H)
fx, fy, cx, cy = intrinsics[:,jj].unbind(dim=-1)
d = torch.zeros_like(Z)
d[Z.abs() > 0.2] = 1.0 / Z[Z.abs() > 0.2]
if isinstance(Gij, SE3):
Ja = torch.stack([
H, o, o, o, Z, -Y,
o, H, o, -Z, o, X,
o, o, H, Y, -X, o,
o, o, o, o, o, o,
], dim=-1).view(1, len(ii), 4, 6)
elif isinstance(Gij, Sim3):
Ja = torch.stack([
H, o, o, o, Z, -Y, X,
o, H, o, -Z, o, X, Y,
o, o, H, Y, -X, o, Z,
o, o, o, o, o, o, o,
], dim=-1).view(1, len(ii), 4, 7)
Jp = torch.stack([
fx*d, o, -fx*X*d*d, o,
o, fy*d, -fy*Y*d*d, o,
], dim=-1).view(1, len(ii), 2, 4)
Jj = torch.matmul(Jp, Ja)
Ji = -Gij[:,:,None].adjT(Jj)
Jz = torch.matmul(Jp, Gij.matrix()[...,:,3:])
return x1, (Z > 0.2).float(), (Ji, Jj, Jz)
if valid:
return x1, (X1[...,2] > 0.2).float()
return x1
def point_cloud(poses, patches, intrinsics, ix):
""" generate point cloud from patches """
return poses[:,ix,None,None].inv() * iproj(patches, intrinsics[:,ix])
def flow_mag(poses, patches, intrinsics, ii, jj, kk, beta=0.3):
""" projective transform """
coords0 = transform(poses, patches, intrinsics, ii, ii, kk)
coords1, val = transform(poses, patches, intrinsics, ii, jj, kk, tonly=False, valid=True)
coords2 = transform(poses, patches, intrinsics, ii, jj, kk, tonly=True)
flow1 = (coords1 - coords0).norm(dim=-1)
flow2 = (coords2 - coords0).norm(dim=-1)
return beta * flow1 + (1-beta) * flow2, (val > 0.5)
|