def ch_iter(self):
sampled_source = mesh_utils.sample_on_mesh(self.source_mesh, self.opt.num_samples[self.level])
# sampled_target = mesh_utils.sample_on_sphere(self.target_mesh[0], self.opt.num_samples[self.level], self.device)
sampled_target = mesh_utils.sample_on_mesh(self.target_mesh, self.opt.num_samples[self.level])
chamfer_loss = self.chamfer(sampled_source, sampled_target)
ch_loss = 0
to_log = []
for gamma, loss, name in zip(self.opt.chamfer_weights,
chamfer_loss, ['dis_s2t', 'dis_t2s', 'n_s2t', 'n_t2s']):
if gamma != 0:
ch_loss += gamma * loss
to_log += [name, loss]
return ch_loss, to_log
def sample_sub_points(self, data: mesh_utils.MeshWrap):
p = self.opt.partial_range[0] + np.random.random() * (
self.opt.partial_range[1] - self.opt.partial_range[0])
sub_mesh, sub_areas = mesh_utils.split_mesh_side(
data.mesh(), data['face_centers'], data['face_areas'],
data['total_area'], p)
sub_points = mesh_utils.sample_on_mesh(sub_mesh, sub_areas,
self.opt.partial_samples[1])
return sub_points
def __getitem__(self, idx):
data = self.load_mesh(idx)
mesh, face_areas = data.mesh(), data['face_areas']
points = mesh_utils.sample_on_mesh(mesh, face_areas,
self.opt.partial_samples[0])
if self.recon or len(self.transforms) > 0:
pc_trans, transforms = self.get_transformed_pc(idx, data, points)
return [points] + [pc.astype(np.float32) for pc in pc_trans] + [
tr.astype(np.float32) for tr in transforms
]
else:
return points