def compute_object_intersect_tensors(name, ray_batch, scene_info, far, object2padding, swap_object_yz, **kwargs): """Compute intersecting rays.""" rays_o = ray_utils.extract_slice_from_ray_batch( # [R, 3] ray_batch=ray_batch, # [R, M] key='origin') rays_d = ray_utils.extract_slice_from_ray_batch( # [R, 3] ray_batch=ray_batch, # [R, M] key='direction') rays_far = ray_utils.extract_slice_from_ray_batch( # [R, 3] ray_batch=ray_batch, # [R, M] key='far') rays_sid = ray_utils.extract_slice_from_ray_batch( # [R, 1] ray_batch=ray_batch, # [R, M] key='metadata') ( box_dims, # [R, 3], [R, 3], [R, 3, 3] box_center, box_rotation) = scene_utils.extract_object_boxes_for_scenes( name=name, scene_info=scene_info, sids=rays_sid, # [R, 1] padding=object2padding[name], swap_yz=swap_object_yz, box_delta_t=kwargs['box_delta_t']) # Compute ray-bbox intersections. intersect_bounds, intersect_indices, intersect_mask = ( # [R', 2],[R',],[R,] box_utils.compute_ray_bbox_bounds_pairwise( rays_o=rays_o, # [R, 3] rays_d=rays_d, # [R, 3] rays_far=rays_far, # [R, 1] box_length=box_dims[:, 0], # [R,] box_width=box_dims[:, 1], # [R,] box_height=box_dims[:, 2], # [R,] box_center=box_center, # [R, 3] box_rotation=box_rotation, # [R, 3, 3] far_limit=far)) # Apply the intersection mask to the ray batch. intersect_ray_batch = apply_intersect_mask_to_tensors( # [R', M] intersect_mask=intersect_mask, # [R,] tensors=[ray_batch])[0] # [R, M] # Update the near and far bounds of the ray batch with the intersect bounds. intersect_ray_batch = ray_utils.update_ray_batch_bounds( # [R', M] ray_batch=intersect_ray_batch, # [R', M] bounds=intersect_bounds) # [R', 2] return intersect_ray_batch, intersect_indices # [R', M], [R', 1]
def network_query_fn_helper_nodirs(pts, ray_batch, network, network_query_fn, use_viewdirs, use_lightdirs, use_lightdir_norm, scene_info, use_random_lightdirs, light_origins=None, **kwargs): """Same as network_query_fn_helper, but without input directions.""" _ = kwargs if not use_viewdirs: viewdirs = None if not use_lightdirs: lightdirs = None # Extract unit-normalized viewing direction. if use_viewdirs: viewdirs = ray_batch[:, 8:11] # [R, 3] viewdirs = broadcast_samples_dim(x=viewdirs, target=pts) # [R, S, 3] else: viewdirs = None # Compute the light directions. # if use_lightdirs: light_ray_batch, lightdirs = compute_lightdirs( # [R, S, 3] pts=pts, metadata=ray_utils.extract_slice_from_ray_batch( ray_batch, key='metadata', use_viewdirs=use_viewdirs), scene_info=scene_info, use_lightdir_norm=use_lightdir_norm, use_random_lightdirs=use_random_lightdirs, light_pos=kwargs['light_pos'], light_origins=light_origins) # else: # light_ray_batch = None # lightdirs = None # Extract additional per-ray metadata. rays_data = ray_utils.extract_slice_from_ray_batch( ray_batch, key='example_id', use_viewdirs=use_viewdirs) # Query NeRF for the corresponding densities for the light points. raw = network_query_fn(pts, viewdirs, lightdirs, rays_data, network) return light_ray_batch, raw
def create_w2o_transformations_tensors(name, scene_info, ray_batch, use_viewdirs, box_delta_t): """Create transformation tensor from world to object space.""" metadata = ray_utils.extract_slice_from_ray_batch( ray_batch, key='metadata', use_viewdirs=use_viewdirs) # [R, 1] w2o_rt_per_scene, w2o_r_per_scene = ( scene_utils.extract_w2o_transformations_per_scene( name=name, scene_info=scene_info, box_delta_t=box_delta_t)) w2o_rt = tf.gather_nd( # [R, 4, 4] params=w2o_rt_per_scene, # [N_scenes, 4, 4] indices=metadata) # [R, 1] w2o_r = tf.gather_nd( # [R, 4, 4] params=w2o_r_per_scene, # [N_scenes, 4, 4] indices=metadata) # [R, 1] return w2o_rt, w2o_r
def compute_view_light_dirs(ray_batch, pts, scene_info, use_viewdirs, use_lightdir_norm, use_random_lightdirs, light_pos): """Compute viewing and lighting directions.""" viewdirs = ray_batch[:, 8:11] # [R, 3] metadata = ray_utils.extract_slice_from_ray_batch( ray_batch, key='metadata', use_viewdirs=use_viewdirs) # [R, 1] viewdirs = broadcast_samples_dim(x=viewdirs, target=pts) # [R, S, 3] light_ray_batch, lightdirs = compute_lightdirs( pts=pts, metadata=metadata, scene_info=scene_info, use_lightdir_norm=use_lightdir_norm, use_random_lightdirs=use_random_lightdirs, light_pos=light_pos) # [R, S, 3] return viewdirs, light_ray_batch, lightdirs
def network_query_fn_helper(pts, ray_batch, network, network_query_fn, viewdirs, lightdirs, use_viewdirs, use_lightdirs): """Query the NeRF network.""" if not use_viewdirs: viewdirs = None if not use_lightdirs: lightdirs = None # Extract unit-normalized viewing direction. # [n_rays, 3] # viewdirs = ray_batch[:, 8:11] if use_viewdirs else None # Extract additional per-ray metadata. # [n_rays, metadata_channels] rays_data = ray_utils.extract_slice_from_ray_batch( ray_batch, key='example_id', use_viewdirs=use_viewdirs) # Query NeRF for the corresponding densities for the light points. raw = network_query_fn(pts, viewdirs, lightdirs, rays_data, network) return raw
def create_shadow_ray_batch(ray_batch, pts, scene_info, light_pos): """Create batch for shadow rays. Args: ray_batch: [?, M] tf.float32. Primary ray batch. pts: [?, S, 3] tf.float32. Primary points. scene_info: Dict. Scene information. light_pos: Returns: shadow_ray_batch: [?S, M] """ # num_rays = tf.shape(pts)[0] # ? num_primary_samples = pts.shape[1] # S pts = tf.reshape(pts, [-1, 3]) # [?S, 3] # Prepare light positions. light_positions = ray_utils.extract_light_positions_for_rays( ray_batch=ray_batch, scene_info=scene_info, light_pos=light_pos) # Get ray scene IDs. rays_sid = ray_utils.extract_slice_from_ray_batch( ray_batch=ray_batch, key='metadata') # Repeat ray-level information by the number of primary samples per ray. light_positions = tf.tile( light_positions[:, None, :], # [?, S, 3] [1, num_primary_samples, 1]) rays_sid = tf.tile( rays_sid[:, None, :], # [?, S, 1] [1, num_primary_samples, 1]) light_positions = tf.reshape(light_positions, [-1, 3]) # [?S, 3] rays_sid = tf.reshape(rays_sid, [-1, 1]) # [?S, 1] # Create the ray batch. shadow_ray_batch = ray_utils.create_ray_batch( rays_o=light_positions, rays_dst=pts, rays_sid=rays_sid) return shadow_ray_batch