def get_scores_once(feats, points, values): # sample points inside voxels sampled_xyz = offset_points(points, self.voxel_size / 2.0, bits=bits) sampled_idx = torch.arange( points.size(0), device=points.device)[:, None].expand(*sampled_xyz.size()[:2]) sampled_xyz, sampled_idx = sampled_xyz.reshape( -1, 3), sampled_idx.reshape(-1) field_inputs = self.forward( { 'sampled_point_xyz': sampled_xyz, 'sampled_point_voxel_idx': sampled_idx, 'sampled_point_ray_direction': None }, { 'voxel_vertex_idx': feats, 'voxel_center_xyz': points, 'voxel_vertex_emb': values }) # get field inputs # evaluation with density field_outputs = field_fn(field_inputs, outputs=['sigma']) free_energy = -torch.relu(field_outputs['sigma']).reshape( -1, bits**3) # return scores return torch.exp(free_energy)
def forward(self, samples, encoder_states): # encoder states point_feats = encoder_states['voxel_vertex_idx'] point_xyz = encoder_states['voxel_center_xyz'] values = encoder_states['voxel_vertex_emb'] # ray point samples sampled_idx = samples['sampled_point_voxel_idx'].long() sampled_xyz = samples['sampled_point_xyz'].requires_grad_(True) sampled_dir = samples['sampled_point_ray_direction'] # prepare inputs for implicit field inputs = {'pos': sampled_xyz, 'ray': sampled_dir} if values is not None: # resample point features point_xyz = F.embedding(sampled_idx, point_xyz) point_feats = F.embedding(F.embedding(sampled_idx, point_feats), values).view(point_xyz.size(0), -1) # tri-linear interpolation p = ((sampled_xyz - point_xyz) / self.voxel_size + .5).unsqueeze(1) q = offset_points(p, .5, offset_only=True).unsqueeze(0) + .5 # BUG (FIX) inputs.update({'emb': trilinear_interp(p, q, point_feats)}) return inputs
def forward(self, samples, encoder_states): # encoder states point_feats = encoder_states['voxel_vertex_idx'] point_xyz = encoder_states['voxel_center_xyz'] values = encoder_states['voxel_vertex_emb'] # ray point samples sampled_idx = samples['sampled_point_voxel_idx'].long() sampled_xyz = samples['sampled_point_xyz'].requires_grad_(True) sampled_dir = samples['sampled_point_ray_direction'] sampled_dis = samples['sampled_point_distance'] # prepare inputs for implicit field # / self.scene_scale inputs = { 'pos': sampled_xyz, 'ray': sampled_dir, 'dists': sampled_dis} # --- just for debugging ---- # # r = inputs['pos'].norm(p=2, dim=-1, keepdim=True) # inputs['pos'] = torch.cat([inputs['pos'] / (r + 1e-8), r / (1 + r)], dim=-1) if values is not None: # resample point features #print('sample-->',sampled_idx.shape,sampled_xyz.shape,point_xyz.shape) print('unique torch -->',torch.unique(sampled_idx)) point_xyz = F.embedding(sampled_idx, point_xyz) point_feats = F.embedding(F.embedding(sampled_idx, point_feats), values).view(point_xyz.size(0), -1) # tri-linear interpolation p = ((sampled_xyz - point_xyz) / self.voxel_size + .5).unsqueeze(1) q = offset_points(p, .5, offset_only=True).unsqueeze(0) + .5 # BUG (FIX) temb = trilinear_interp(p, q, point_feats) #print('interpolate-->',temb.shape,point_xyz.shape) inputs.update({'emb': temb,'originalpoints':point_xyz}) return inputs
def __init__(self, args, voxel_path=None, bbox_path=None, shared_values=None): super().__init__(args) # read initial voxels or learned sparse voxels self.voxel_path = voxel_path if voxel_path is not None else args.voxel_path self.bbox_path = bbox_path if bbox_path is not None else getattr(args, "initial_boundingbox", None) assert (self.bbox_path is not None) or (self.voxel_path is not None), \ "at least initial bounding box or pretrained voxel files are required." self.voxel_index = None self.scene_scale = getattr(args, "scene_scale", 1.0) if self.voxel_path is not None: # read voxel file assert os.path.exists(self.voxel_path), "voxel file must exist" if Path(self.voxel_path).suffix == '.ply': from plyfile import PlyData, PlyElement plyvoxel = PlyData.read(self.voxel_path) elements = [x.name for x in plyvoxel.elements] assert 'vertex' in elements plydata = plyvoxel['vertex'] fine_points = torch.from_numpy( np.stack([plydata['x'], plydata['y'], plydata['z']]).astype('float32').T) if 'face' in elements: # read voxel meshes... automatically detect voxel size faces = plyvoxel['face']['vertex_indices'] t = fine_points[faces[0].astype('int64')] voxel_size = torch.abs(t[0] - t[1]).max() # indexing voxel vertices fine_points = torch.unique(fine_points, dim=0) # vertex_ids, _ = discretize_points(fine_points, voxel_size) # vertex_ids_offset = vertex_ids + 1 # # simple hashing # vertex_ids = vertex_ids[:, 0] * 1000000 + vertex_ids[:, 1] * 1000 + vertex_ids[:, 2] # vertex_ids_offset = vertex_ids_offset[:, 0] * 1000000 + vertex_ids_offset[:, 1] * 1000 + vertex_ids_offset[:, 2] # vertex_ids = {k: True for k in vertex_ids.tolist()} # vertex_inside = [v in vertex_ids for v in vertex_ids_offset.tolist()] # # get voxel centers # fine_points = fine_points[torch.tensor(vertex_inside)] + voxel_size * .5 # fine_points = fine_points + voxel_size * .5 --> use all corners as centers else: # voxel size must be provided assert getattr(args, "voxel_size", None) is not None, "final voxel size is essential." voxel_size = args.voxel_size if 'quality' in elements: self.voxel_index = torch.from_numpy(plydata['quality']).long() else: # supporting the old style .txt voxel points fine_points = torch.from_numpy(np.loadtxt(self.voxel_path)[:, 3:].astype('float32')) else: # read bounding-box file bbox = np.loadtxt(self.bbox_path) voxel_size = bbox[-1] if getattr(args, "voxel_size", None) is None else args.voxel_size fine_points = torch.from_numpy(bbox2voxels(bbox[:6], voxel_size)) half_voxel = voxel_size * .5 # transform from voxel centers to voxel corners (key/values) fine_coords, _ = discretize_points(fine_points, half_voxel) fine_keys0 = offset_points(fine_coords, 1.0).reshape(-1, 3) fine_keys, fine_feats = torch.unique(fine_keys0, dim=0, sorted=True, return_inverse=True) fine_feats = fine_feats.reshape(-1, 8) num_keys = torch.scalar_tensor(fine_keys.size(0)).long() # ray-marching step size if getattr(args, "raymarching_stepsize_ratio", 0) > 0: step_size = args.raymarching_stepsize_ratio * voxel_size else: step_size = args.raymarching_stepsize # register parameters (will be saved to checkpoints) self.register_buffer("points", fine_points) # voxel centers self.register_buffer("keys", fine_keys.long()) # id used to find voxel corners/embeddings self.register_buffer("feats", fine_feats.long()) # for each voxel, 8 voxel corner ids self.register_buffer("num_keys", num_keys) #self.register_buffer("points_labels", fine_points) # voxel centers self.register_buffer("keep", fine_feats.new_ones(fine_feats.size(0)).long()) # whether the voxel will be pruned self.register_buffer("voxel_size", torch.scalar_tensor(voxel_size)) self.register_buffer("step_size", torch.scalar_tensor(step_size)) self.register_buffer("max_hits", torch.scalar_tensor(args.max_hits)) logger.info("loaded {} voxel centers, {} voxel corners".format(fine_points.size(0), num_keys)) # set-up other hyperparameters and initialize running time caches self.embed_dim = getattr(args, "voxel_embed_dim", None) self.deterministic_step = getattr(args, "deterministic_step", False) self.use_octree = getattr(args, "use_octree", False) self.track_max_probs = getattr(args, "track_max_probs", False) self._runtime_caches = { "flatten_centers": None, "flatten_children": None, "max_voxel_probs": None } # sparse voxel embeddings if shared_values is None and self.embed_dim > 0: self.values = Embedding(num_keys, self.embed_dim, None) else: self.values = shared_values
def __init__(self, args, voxel_path=None, bbox_path=None, shared_values=None): super().__init__(args) # read initial voxels or learned sparse voxels self.voxel_path = voxel_path if voxel_path is not None else args.voxel_path self.bbox_path = bbox_path if bbox_path is not None else getattr( args, "initial_boundingbox", None) assert (self.bbox_path is not None) or (self.voxel_path is not None), \ "at least initial bounding box or pretrained voxel files are required." self.voxel_index = None if self.voxel_path is not None: assert os.path.exists(self.voxel_path), "voxel file must exist" assert getattr(args, "voxel_size", None) is not None, "final voxel size is essential." voxel_size = args.voxel_size if Path(self.voxel_path).suffix == '.ply': from plyfile import PlyData, PlyElement plydata = PlyData.read(self.voxel_path)['vertex'] fine_points = torch.from_numpy( np.stack([plydata['x'], plydata['y'], plydata['z']]).astype('float32').T) try: self.voxel_index = torch.from_numpy( plydata['quality']).long() except ValueError: pass else: # supporting the old version voxel points fine_points = torch.from_numpy( np.loadtxt(self.voxel_path)[:, 3:].astype('float32')) else: bbox = np.loadtxt(self.bbox_path) voxel_size = bbox[-1] fine_points = torch.from_numpy(bbox2voxels(bbox[:6], voxel_size)) half_voxel = voxel_size * .5 # transform from voxel centers to voxel corners (key/values) fine_coords, _ = discretize_points(fine_points, half_voxel) fine_keys0 = offset_points(fine_coords, 1.0).reshape(-1, 3) fine_keys, fine_feats = torch.unique(fine_keys0, dim=0, sorted=True, return_inverse=True) fine_feats = fine_feats.reshape(-1, 8) num_keys = torch.scalar_tensor(fine_keys.size(0)).long() # ray-marching step size if getattr(args, "raymarching_stepsize_ratio", 0) > 0: step_size = args.raymarching_stepsize_ratio * voxel_size else: step_size = args.raymarching_stepsize # register parameters (will be saved to checkpoints) self.register_buffer("points", fine_points) # voxel centers self.register_buffer( "keys", fine_keys.long()) # id used to find voxel corners/embeddings self.register_buffer( "feats", fine_feats.long()) # for each voxel, 8 voxel corner ids self.register_buffer("num_keys", num_keys) self.register_buffer( "keep", fine_feats.new_ones( fine_feats.size(0)).long()) # whether the voxel will be pruned self.register_buffer("voxel_size", torch.scalar_tensor(voxel_size)) self.register_buffer("step_size", torch.scalar_tensor(step_size)) self.register_buffer("max_hits", torch.scalar_tensor(args.max_hits)) # set-up other hyperparameters and initialize running time caches self.embed_dim = getattr(args, "voxel_embed_dim", None) self.deterministic_step = getattr(args, "deterministic_step", False) self.use_octree = getattr(args, "use_octree", False) self.track_max_probs = getattr(args, "track_max_probs", False) self._runtime_caches = { "flatten_centers": None, "flatten_children": None, "max_voxel_probs": None } # sparse voxel embeddings if shared_values is None and self.embed_dim > 0: self.values = Embedding(num_keys, self.embed_dim, None) else: self.values = shared_values