|
|
|
"""Contains utility functions for rendering.""" |
|
import torch |
|
|
|
def normalize_vecs(vectors): |
|
""" |
|
Normalize vector lengths. |
|
""" |
|
return vectors / (torch.norm(vectors, dim=-1, keepdim=True)) |
|
|
|
def truncated_normal(tensor, mean=0, std=1): |
|
""" |
|
Samples from truncated normal distribution. |
|
""" |
|
size = tensor.shape |
|
tmp = tensor.new_empty(size + (4,)).normal_() |
|
valid = (tmp < 2) & (tmp > -2) |
|
ind = valid.max(-1, keepdim=True)[1] |
|
tensor.data.copy_(tmp.gather(-1, ind).squeeze(-1)) |
|
tensor.data.mul_(std).add_(mean) |
|
return tensor |
|
|
|
def get_grid_coords(points, bounds): |
|
""" transform points from the world coordinate to the volume coordinate |
|
pts: batch_size, num_point, 3 |
|
bounds: 2, 3 |
|
""" |
|
|
|
bounds = bounds[None] |
|
min_xyz = bounds[:, :1] |
|
points = points - min_xyz |
|
|
|
size = bounds[:, 1] - bounds[:, 0] |
|
points = (points / size[:, None]) * 2 - 1 |
|
return points |
|
|
|
def grid_sample_3d(image, optical): |
|
"""grid sample images by the optical in 3D format |
|
image: batch_size, channel, D, H, W |
|
optical: batch_size, D, H, W, 3 |
|
""" |
|
N, C, ID, IH, IW = image.shape |
|
_, D, H, W, _ = optical.shape |
|
|
|
ix = optical[..., 0] |
|
iy = optical[..., 1] |
|
iz = optical[..., 2] |
|
|
|
ix = ((ix + 1) / 2) * (IW - 1) |
|
iy = ((iy + 1) / 2) * (IH - 1) |
|
iz = ((iz + 1) / 2) * (ID - 1) |
|
with torch.no_grad(): |
|
ix_tnw = torch.floor(ix) |
|
iy_tnw = torch.floor(iy) |
|
iz_tnw = torch.floor(iz) |
|
|
|
ix_tne = ix_tnw + 1 |
|
iy_tne = iy_tnw |
|
iz_tne = iz_tnw |
|
|
|
ix_tsw = ix_tnw |
|
iy_tsw = iy_tnw + 1 |
|
iz_tsw = iz_tnw |
|
|
|
ix_tse = ix_tnw + 1 |
|
iy_tse = iy_tnw + 1 |
|
iz_tse = iz_tnw |
|
|
|
ix_bnw = ix_tnw |
|
iy_bnw = iy_tnw |
|
iz_bnw = iz_tnw + 1 |
|
|
|
ix_bne = ix_tnw + 1 |
|
iy_bne = iy_tnw |
|
iz_bne = iz_tnw + 1 |
|
|
|
ix_bsw = ix_tnw |
|
iy_bsw = iy_tnw + 1 |
|
iz_bsw = iz_tnw + 1 |
|
|
|
ix_bse = ix_tnw + 1 |
|
iy_bse = iy_tnw + 1 |
|
iz_bse = iz_tnw + 1 |
|
|
|
tnw = (ix_bse - ix) * (iy_bse - iy) * (iz_bse - iz) |
|
tne = (ix - ix_bsw) * (iy_bsw - iy) * (iz_bsw - iz) |
|
tsw = (ix_bne - ix) * (iy - iy_bne) * (iz_bne - iz) |
|
tse = (ix - ix_bnw) * (iy - iy_bnw) * (iz_bnw - iz) |
|
bnw = (ix_tse - ix) * (iy_tse - iy) * (iz - iz_tse) |
|
bne = (ix - ix_tsw) * (iy_tsw - iy) * (iz - iz_tsw) |
|
bsw = (ix_tne - ix) * (iy - iy_tne) * (iz - iz_tne) |
|
bse = (ix - ix_tnw) * (iy - iy_tnw) * (iz - iz_tnw) |
|
|
|
with torch.no_grad(): |
|
torch.clamp(ix_tnw, 0, IW - 1, out=ix_tnw) |
|
torch.clamp(iy_tnw, 0, IH - 1, out=iy_tnw) |
|
torch.clamp(iz_tnw, 0, ID - 1, out=iz_tnw) |
|
|
|
torch.clamp(ix_tne, 0, IW - 1, out=ix_tne) |
|
torch.clamp(iy_tne, 0, IH - 1, out=iy_tne) |
|
torch.clamp(iz_tne, 0, ID - 1, out=iz_tne) |
|
|
|
torch.clamp(ix_tsw, 0, IW - 1, out=ix_tsw) |
|
torch.clamp(iy_tsw, 0, IH - 1, out=iy_tsw) |
|
torch.clamp(iz_tsw, 0, ID - 1, out=iz_tsw) |
|
|
|
torch.clamp(ix_tse, 0, IW - 1, out=ix_tse) |
|
torch.clamp(iy_tse, 0, IH - 1, out=iy_tse) |
|
torch.clamp(iz_tse, 0, ID - 1, out=iz_tse) |
|
|
|
torch.clamp(ix_bnw, 0, IW - 1, out=ix_bnw) |
|
torch.clamp(iy_bnw, 0, IH - 1, out=iy_bnw) |
|
torch.clamp(iz_bnw, 0, ID - 1, out=iz_bnw) |
|
|
|
torch.clamp(ix_bne, 0, IW - 1, out=ix_bne) |
|
torch.clamp(iy_bne, 0, IH - 1, out=iy_bne) |
|
torch.clamp(iz_bne, 0, ID - 1, out=iz_bne) |
|
|
|
torch.clamp(ix_bsw, 0, IW - 1, out=ix_bsw) |
|
torch.clamp(iy_bsw, 0, IH - 1, out=iy_bsw) |
|
torch.clamp(iz_bsw, 0, ID - 1, out=iz_bsw) |
|
|
|
torch.clamp(ix_bse, 0, IW - 1, out=ix_bse) |
|
torch.clamp(iy_bse, 0, IH - 1, out=iy_bse) |
|
torch.clamp(iz_bse, 0, ID - 1, out=iz_bse) |
|
|
|
image = image.view(N, C, ID * IH * IW) |
|
|
|
tnw_val = torch.gather(image, 2, |
|
(iz_tnw * IW * IH + iy_tnw * IW + |
|
ix_tnw).long().view(N, 1, |
|
D * H * W).repeat(1, C, 1)) |
|
tne_val = torch.gather(image, 2, |
|
(iz_tne * IW * IH + iy_tne * IW + |
|
ix_tne).long().view(N, 1, |
|
D * H * W).repeat(1, C, 1)) |
|
tsw_val = torch.gather(image, 2, |
|
(iz_tsw * IW * IH + iy_tsw * IW + |
|
ix_tsw).long().view(N, 1, |
|
D * H * W).repeat(1, C, 1)) |
|
tse_val = torch.gather(image, 2, |
|
(iz_tse * IW * IH + iy_tse * IW + |
|
ix_tse).long().view(N, 1, |
|
D * H * W).repeat(1, C, 1)) |
|
bnw_val = torch.gather(image, 2, |
|
(iz_bnw * IW * IH + iy_bnw * IW + |
|
ix_bnw).long().view(N, 1, |
|
D * H * W).repeat(1, C, 1)) |
|
bne_val = torch.gather(image, 2, |
|
(iz_bne * IW * IH + iy_bne * IW + |
|
ix_bne).long().view(N, 1, |
|
D * H * W).repeat(1, C, 1)) |
|
bsw_val = torch.gather(image, 2, |
|
(iz_bsw * IW * IH + iy_bsw * IW + |
|
ix_bsw).long().view(N, 1, |
|
D * H * W).repeat(1, C, 1)) |
|
bse_val = torch.gather(image, 2, |
|
(iz_bse * IW * IH + iy_bse * IW + |
|
ix_bse).long().view(N, 1, |
|
D * H * W).repeat(1, C, 1)) |
|
|
|
out_val = (tnw_val.view(N, C, D, H, W) * tnw.view(N, 1, D, H, W) + |
|
tne_val.view(N, C, D, H, W) * tne.view(N, 1, D, H, W) + |
|
tsw_val.view(N, C, D, H, W) * tsw.view(N, 1, D, H, W) + |
|
tse_val.view(N, C, D, H, W) * tse.view(N, 1, D, H, W) + |
|
bnw_val.view(N, C, D, H, W) * bnw.view(N, 1, D, H, W) + |
|
bne_val.view(N, C, D, H, W) * bne.view(N, 1, D, H, W) + |
|
bsw_val.view(N, C, D, H, W) * bsw.view(N, 1, D, H, W) + |
|
bse_val.view(N, C, D, H, W) * bse.view(N, 1, D, H, W)) |
|
|
|
return out_val |
|
|
|
def interpolate_feature(points, volume, bounds): |
|
""" |
|
points: batch_size, num_point, 3 |
|
volume: batch_size, num_channel, d, h, w |
|
bounds: 2, 3 |
|
""" |
|
grid_coords = get_grid_coords(points, bounds) |
|
grid_coords = grid_coords[:, None, None] |
|
|
|
|
|
|
|
|
|
point_features = grid_sample_3d(volume, grid_coords) |
|
point_features = point_features[:, :, 0, 0] |
|
return point_features |
