def forward(self, xy, depth, cam2world, intersection_net, intrinsics):
self.logs = list()
batch_size, _, _ = cam2world.shape
intersections = geometry.world_from_xy_depth(xy=xy,
depth=depth,
cam2world=cam2world,
intrinsics=intrinsics)
depth = geometry.depth_from_world(intersections, cam2world)
if self.training:
print(depth.min(), depth.max())
return intersections, depth
def forward(self,
cam2world, # pose
phi,
uv,
intrinsics):
batch_size, num_samples, _ = uv.shape
log = list()
ray_dirs = geometry.get_ray_directions(uv,
cam2world=cam2world,
intrinsics=intrinsics)
initial_depth = torch.zeros((batch_size, num_samples, 1)).normal_(mean=0.05, std=5e-4).cuda()
init_world_coords = geometry.world_from_xy_depth(uv,
initial_depth,
intrinsics=intrinsics,
cam2world=cam2world)
world_coords = [init_world_coords]
depths = [initial_depth]
states = [None]
for step in range(self.steps):
v = phi(world_coords[-1])
state = self.lstm(v.view(-1, self.n_feature_channels), states[-1])
if state[0].requires_grad:
state[0].register_hook(lambda x: x.clamp(min=-20, max=20))
signed_distance = self.out_layer(state[0]).view(batch_size, num_samples, 1)
new_world_coords = world_coords[-1] + ray_dirs * signed_distance
states.append(state)
world_coords.append(new_world_coords)
depth = geometry.depth_from_world(world_coords[-1], cam2world)
# if self.training:
# print("Raymarch step %d/%d: Min depth %0.6f, max depth %0.6f" %
# (step, self.steps, depths[-1].min().detach().cpu().numpy(), depths[-1].max().detach().cpu().numpy()))
depths.append(depth)
if not self.counter % 100:
# Write tensorboard summary for each step of ray-marcher.
drawing_depths = torch.stack(depths, dim=0)[:, 0, :, :]
drawing_depths = util.lin2img(drawing_depths).repeat(1, 3, 1, 1)
log.append(('image', 'raycast_progress',
torch.clamp(torchvision.utils.make_grid(drawing_depths, scale_each=False, normalize=True), 0.0,
5),
100))
# Visualize residual step distance (i.e., the size of the final step)
fig = util.show_images([util.lin2img(signed_distance)[i, :, :, :].detach().cpu().numpy().squeeze()
for i in range(batch_size)])
log.append(('figure', 'stopping_distances', fig, 100))
self.counter += 1
return world_coords[-1], depths[-1], log