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