def visualize(self, data):
''' Performs a visualization step for the data.
Args:
data (dict): data dictionary
'''
device = self.device
batch_size = data['points'].size(0)
inputs = data.get('inputs', torch.empty(batch_size, 0)).to(device)
shape = (32, 32, 32)
p = make_3d_grid([-0.5] * 3, [0.5] * 3, shape).to(device)
p = p.expand(batch_size, *p.size())
kwargs = {}
with torch.no_grad():
p_r = self.model(p, inputs, sample=self.eval_sample, **kwargs)
occ_hat = p_r.probs.view(batch_size, *shape)
voxels_out = (occ_hat >= self.threshold).cpu().numpy()
for i in trange(batch_size):
input_img_path = os.path.join(self.vis_dir, '%03d_in.png' % i)
vis.visualize_data(inputs[i].cpu(), self.input_type,
input_img_path)
vis.visualize_voxels(voxels_out[i],
os.path.join(self.vis_dir, '%03d.png' % i))
def export_points(mesh, modelname, loc, scale, args):
if not mesh.is_watertight:
print('Warning: mesh %s is not watertight!'
'Cannot sample points.' % modelname)
return
filename = os.path.join(args.points_folder, modelname + '.npz')
if not args.overwrite and os.path.exists(filename):
print('Points already exist: %s' % filename)
return
# uniform grid
total_l = 1. + args.points_padding
res = args.points_resolution
points_uniform = make_3d_grid((-total_l/2.0 + total_l/(res*2),)*3, (total_l/2.0 - total_l/(res*2),)*3, (res,)*3).numpy()
points_uniform_occupancies = check_mesh_contains(mesh, points_uniform).astype(np.float32)
points = points_uniform
occupancies = points_uniform_occupancies
# Compress
if args.float16:
dtype = np.float16
else:
dtype = np.float32
points = points.astype(dtype)
occupancies = occupancies.astype(dtype)
print('Writing points: %s' % filename)
np.savez(filename, points=points, occupancies=occupancies,
loc=loc, scale=scale)
def visualize(self, data):
""" Performs a visualization step for the data.
Args:
data (dict): data dictionary
"""
batch_size = data["points"].shape[0]
inputs = data.get("inputs", tf.zeros([batch_size, 0]))
shape = (32, 32, 32)
p = make_3d_grid([-0.5] * 3, [0.5] * 3, shape) # CHECK
p = tf.broadcast_to(p, [batch_size, *p.shape])
kwargs = {}
p_r = self.model(p,
inputs,
sample=self.eval_sample,
training=False,
**kwargs)
occ_hat = tf.reshape(p_r.probs_parameter(), [batch_size, *shape])
voxels_out = (occ_hat >= self.threshold).numpy()
for i in trange(batch_size):
input_img_path = os.path.join(self.vis_dir, "%03d_in.png" % i)
vis.visualize_data(inputs[i], self.input_type, input_img_path)
vis.visualize_voxels(voxels_out[i],
os.path.join(self.vis_dir, "%03d.png" % i))
def generate_from_latent(self,
z,
c_first=None,
Rt=None,
K=None,
stats_dict={},
**kwargs):
''' Generates mesh from latent.
Args:
z (tensor): latent code z
c (tensor): latent conditioned code c
stats_dict (dict): stats dictionary
'''
threshold = np.log(self.threshold) - np.log(1. - self.threshold)
t0 = time.time()
# Compute bounding box size
box_size = 1 + self.padding
# Shortcut
if self.upsampling_steps == 0:
nx = self.resolution0
pointsf = box_size * make_3d_grid((-0.5, ) * 3, (0.5, ) * 3,
(nx, ) * 3)
values = self.eval_points(pointsf, z, c_first,
**kwargs).cpu().numpy()
value_grid = values.reshape(nx, nx, nx)
else:
mesh_extractor = MISE(self.resolution0, self.upsampling_steps,
threshold)
points = mesh_extractor.query()
while points.shape[0] != 0:
# Query points
pointsf = torch.FloatTensor(points).to(self.device)
# Normalize to bounding box
pointsf = pointsf / mesh_extractor.resolution
pointsf = box_size * (pointsf - 0.5)
# Evaluate model and update
values = self.eval_points(pointsf, z, c_first, Rt, K,
**kwargs).cpu().numpy()
values = values.astype(np.float64)
mesh_extractor.update(points, values)
points = mesh_extractor.query()
value_grid = mesh_extractor.to_dense()
# Extract mesh
stats_dict['time (eval points)'] = time.time() - t0
mesh = self.extract_mesh(value_grid,
z,
c_first,
Rt,
K,
stats_dict=stats_dict)
return mesh
def visualize(self, data, it=0, vis_type='mesh'):
''' Visualized the data.
Args:
data (dict): data dictionary
it (int): training iteration
vis_type (string): visualization type
'''
if self.multi_gpu:
print(
"Sorry, visualizations currently not implemented when using \
multi GPU training.")
return 0
device = self.device
inputs = data.get('inputs', torch.empty(1, 0)).to(device)
batch_size = inputs.shape[0]
c = self.model.encode_inputs(inputs)
if vis_type == 'voxel':
shape = (32, 32, 32)
p = make_3d_grid([-0.5] * 3, [0.5] * 3, shape).to(device)
p = p.unsqueeze(0).repeat(batch_size, 1, 1)
with torch.no_grad():
p_r = self.model.decode(p, c=c).probs
voxels_out = (p_r >= self.threshold).cpu().numpy()
voxels_out = voxels_out.reshape(batch_size, 32, 32, 32)
for i in range(batch_size):
out_file = os.path.join(self.vis_dir, '%03d.png' % i)
vis.visualize_voxels(voxels_out[i], out_file)
elif vis_type == 'pointcloud':
p = torch.rand(batch_size, 60000, 3).to(device) - 0.5
with torch.no_grad():
occ = self.model.decode(p, c=c).probs
mask = occ > self.threshold
for i in range(batch_size):
pi = p[i][mask[i]].cpu()
out_file = os.path.join(self.vis_dir, '%03d.png' % i)
vis.visualize_pointcloud(pi, out_file=out_file)
elif vis_type == 'mesh':
try:
mesh_list = self.generator.generate_meshes(
data, return_stats=False)
for i, mesh in tqdm(enumerate(mesh_list)):
if self.overwrite_visualization:
ending = ''
else:
ending = '_%010d' % it
mesh_out_file = os.path.join(
self.vis_dir, '%03d%s.ply' % (i, ending))
mesh.export(mesh_out_file)
except Exception as e:
print("Exception occurred during visualization: ", e)
else:
print('The visualization type %s is not valid!' % vis_type)
def voxelize_interior(mesh, resolution):
shape = (resolution, ) * 3
bb_min = (0.5, ) * 3
bb_max = (resolution - 0.5, ) * 3
# Create points. Add noise to break symmetry
points = make_3d_grid(bb_min, bb_max, shape=shape).numpy()
points = points + 0.1 * (np.random.rand(*points.shape) - 0.5)
points = (points / resolution - 0.5)
occ = check_mesh_contains(mesh, points)
occ = occ.reshape(shape)
return occ
def visualize(self, data, it=0., epoch_it=0.):
''' Performs a visualization step for the data.
Args:
data (dict): data dictionary
'''
device = self.device
batch_size = data['points'].size(0)
inputs = data.get('inputs', torch.empty(batch_size, 0)).to(device)
angles = data.get('angles').to(device)
shape = (32, 32, 32)
p = make_3d_grid([-0.5] * 3, [0.5] * 3, shape).to(device)
p = p.expand(batch_size, *p.size())
kwargs = {}
with torch.no_grad():
_, _, sgn, _, _ = self.model(p * self.pnet_point_scale,
inputs,
sample=self.eval_sample,
angles=angles,
**kwargs)
if self.is_sdf:
if self.is_logits_by_min:
logits = (sgn.min(1)[0] <= 0).float()
else:
raise NotImplementedError
else:
if self.is_logits_by_max:
logits = convert_tsd_range_to_zero_to_one(sgn.max(1)[0])
elif self.is_logits_by_sign_filter:
positive = torch.relu(sgn).sum(1)
negative = torch.relu(-sgn).sum(1)
logits = torch.where(positive >= negative, positive, -negative)
else:
logits = convert_tsd_range_to_zero_to_one(sgn).sum(1)
occ_hat = logits.view(batch_size, *shape)
voxels_out = (occ_hat >= self.threshold).cpu().numpy()
input_images = []
voxels_images = []
for i in trange(batch_size):
input_img_path = os.path.join(self.vis_dir, '%03d_in.png' % i)
vis.visualize_data(inputs[i].cpu(), self.input_type,
input_img_path)
vis.visualize_voxels(voxels_out[i],
os.path.join(self.vis_dir, '%03d.png' % i))
def visualize(self, data):
''' Performs a visualization step for the data.
Args:
data (dict): data dictionary
'''
device = self.device
batch_size = data['points'].size(0)
inputs = data.get('inputs').to(device)
#gt_depth_maps = data.get('inputs.depth').to(device)
gt_mask = data.get('inputs.mask').to(device).byte()
shape = (32, 32, 32)
p = make_3d_grid([-0.5] * 3, [0.5] * 3, shape).to(device)
p = p.expand(batch_size, *p.size())
kwargs = {}
with torch.no_grad():
pr_depth_maps = self.model.predict_depth_map(inputs)
background_setting(pr_depth_maps, gt_mask)
p_r = self.model.forward_halfway(p,
pr_depth_maps,
sample=self.eval_sample,
**kwargs)
occ_hat = p_r.probs.view(batch_size, *shape)
voxels_out = (occ_hat >= self.threshold).cpu().numpy()
for i in trange(batch_size):
input_img_path = os.path.join(self.vis_dir, '%03d_in.png' % i)
vis.visualize_data(inputs[i].cpu(), 'img', input_img_path)
vis.visualize_voxels(voxels_out[i],
os.path.join(self.vis_dir, '%03d.png' % i))
depth_map_path = os.path.join(self.vis_dir,
'%03d_pr_depth.png' % i)
depth_map = pr_depth_maps[i].cpu()
depth_map = depth_to_L(depth_map, gt_mask[i].cpu())
vis.visualize_data(depth_map, 'img', depth_map_path)
def eval_step(self, data):
''' Performs an evaluation step.
Args:
data (dict): data dictionary
'''
self.model.eval()
device = self.device
threshold = self.threshold
eval_dict = {}
# Compute elbo
points = data.get('points').to(device)
occ = data.get('points.occ').to(device)
inputs = data.get('inputs', torch.empty(points.size(0), 0)).to(device)
voxels_occ = data.get('voxels')
points_iou = data.get('points_iou').to(device)
occ_iou = data.get('points_iou.occ').to(device)
angles = data.get('angles').to(device)
kwargs = {}
"""
with torch.no_grad():
elbo, rec_error, kl = self.model.compute_elbo(
points, occ, inputs, **kwargs)
eval_dict['loss'] = -elbo.mean().item()
eval_dict['rec_error'] = rec_error.mean().item()
eval_dict['kl'] = kl.mean().item()
"""
# Compute iou
batch_size = points.size(0)
with torch.no_grad():
_, _, sgn, _, _ = self.model(points_iou * self.pnet_point_scale,
inputs,
sample=self.eval_sample,
angles=angles,
**kwargs)
if self.is_sdf:
if self.is_logits_by_min:
logits = (sgn.min(1)[0] <= 0).float()
else:
raise NotImplementedError
else:
if self.is_eval_logits_by_max:
logits = (sgn >= 0.).float().max(1)[0]
elif self.is_logits_by_max:
logits = convert_tsd_range_to_zero_to_one(sgn.max(1)[0])
elif self.is_logits_by_sign_filter:
positive = torch.relu(sgn).sum(1)
negative = torch.relu(-sgn).sum(1)
logits = torch.where(positive >= negative, positive, -negative)
else:
logits = convert_tsd_range_to_zero_to_one(sgn).sum(1)
occ_iou_np = (occ_iou >= self.threshold).cpu().numpy()
occ_iou_hat_np = (logits >= threshold).cpu().numpy()
iou = compute_iou(occ_iou_np, occ_iou_hat_np).mean()
eval_dict['iou'] = iou
eval_dict['overlap'] = (torch.where(sgn >= 0, torch.ones_like(sgn),
torch.zeros_like(sgn)).sum(1) >
1).sum().detach().cpu().numpy().item()
eval_dict['overlap_mean'] = (
torch.where(sgn >= 0, torch.ones_like(sgn),
torch.zeros_like(sgn)).sum(1) >
1).float().mean().detach().cpu().numpy().item()
# Estimate voxel iou
if voxels_occ is not None:
voxels_occ = voxels_occ.to(device)
points_voxels = make_3d_grid((-0.5 + 1 / 64, ) * 3,
(0.5 - 1 / 64, ) * 3, (32, ) * 3)
points_voxels = points_voxels.expand(batch_size,
*points_voxels.size())
points_voxels = points_voxels.to(device)
with torch.no_grad():
_, _, sgn, _, _ = self.model(points_voxels *
self.pnet_point_scale,
inputs,
sample=self.eval_sample,
angles=angles,
**kwargs)
if self.is_sdf:
if self.is_logits_by_min:
logits = (sgn.min(1)[0] <= 0).float()
else:
raise NotImplementedError
else:
if self.is_eval_logits_by_max:
logits = (sgn >= 0.).float().max(1)[0]
elif self.is_logits_by_max:
logits = convert_tsd_range_to_zero_to_one(sgn.max(1)[0])
elif self.is_logits_by_sign_filter:
positive = torch.relu(sgn).sum(1)
negative = torch.relu(-sgn).sum(1)
logits = torch.where(positive >= negative, positive,
-negative)
else:
logits = convert_tsd_range_to_zero_to_one(sgn).sum(1)
voxels_occ_np = (voxels_occ >= 0.5).cpu().numpy()
occ_hat_np = (logits >= threshold).cpu().numpy()
iou_voxels = compute_iou(voxels_occ_np, occ_hat_np).mean()
eval_dict['iou_voxels'] = iou_voxels
return eval_dict
def eval_step(self, data):
''' Performs an evaluation step.
Args:
data (dict): data dictionary
'''
self.model.eval()
device = self.device
threshold = self.threshold
eval_dict = {}
# Compute elbo
points = data.get('points').to(device)
occ = data.get('points.occ').to(device)
inputs = data.get('inputs', torch.empty(points.size(0), 0)).to(device)
voxels_occ = data.get('voxels')
points_iou = data.get('points_iou').to(device)
occ_iou = data.get('points_iou.occ').to(device)
world_mat = data.get('inputs.world_mat').to(device)
camera_mat = data.get('inputs.camera_mat').to(device)
camera_args = common.get_camera_args(data,
'points.loc',
'points.scale',
device=self.device)
world_mat, camera_mat = camera_args['Rt'], camera_args['K']
kwargs = {}
with torch.no_grad():
elbo, rec_error, kl = self.model.compute_elbo(
points, occ, inputs, world_mat, camera_mat, **kwargs)
eval_dict['loss'] = -elbo.mean().item()
eval_dict['rec_error'] = rec_error.mean().item()
eval_dict['kl'] = kl.mean().item()
# Compute iou
batch_size = points.size(0)
with torch.no_grad():
p_out = self.model(points_iou,
inputs,
world_mat,
camera_mat,
sample=self.eval_sample,
**kwargs)
occ_iou_np = (occ_iou >= 0.5).cpu().numpy()
occ_iou_hat_np = (p_out.probs >= threshold).cpu().numpy()
iou = compute_iou(occ_iou_np, occ_iou_hat_np).mean()
eval_dict['iou'] = iou
# Estimate voxel iou
if voxels_occ is not None:
voxels_occ = voxels_occ.to(device)
points_voxels = make_3d_grid((-0.5 + 1 / 64, ) * 3,
(0.5 - 1 / 64, ) * 3, (32, ) * 3)
points_voxels = points_voxels.expand(batch_size,
*points_voxels.size())
points_voxels = points_voxels.to(device)
with torch.no_grad():
p_out = self.model(points_voxels,
inputs,
world_mat,
camera_mat,
sample=self.eval_sample,
**kwargs)
voxels_occ_np = (voxels_occ >= 0.5).cpu().numpy()
occ_hat_np = (p_out.probs >= threshold).cpu().numpy()
iou_voxels = compute_iou(voxels_occ_np, occ_hat_np).mean()
eval_dict['iou_voxels'] = iou_voxels
return eval_dict
def eval_step(self, data):
''' Performs an evaluation step.
Args:
data (dict): data dictionary
'''
self.model.eval()
device = self.device
threshold = self.threshold
eval_dict = {}
points = data.get('points').to(device)
sdf = data.get('points.sdf').to(device)
inputs = data.get('inputs', torch.empty(points.size(0), 0)).to(device)
voxels_occ = data.get('voxels')
points_iou = data.get('points_iou').to(device)
sdf_iou = data.get('points_iou.sdf').to(device)
kwargs = {}
# Compute loss
with torch.no_grad():
p_out = self.model(points, inputs,
sample=self.eval_sample, **kwargs)
loss = get_sdf_loss(p_out.logits, sdf, self.loss_type, ratio=self.sdf_ratio)
loss = loss.sum(-1).mean()
eval_dict['loss'] = loss.item()
# Compute iou
batch_size = points.size(0)
with torch.no_grad():
p_out = self.model(points_iou, inputs,
sample=self.eval_sample, **kwargs)
occ_iou_np = (sdf_iou <= 0.).cpu().numpy()
occ_iou_pred_sdf = p_out.logits / self.sdf_ratio
occ_iou_hat_np = (occ_iou_pred_sdf <= threshold).cpu().numpy()
iou = compute_iou(occ_iou_np, occ_iou_hat_np).mean()
eval_dict['iou'] = iou
# Estimate voxel iou
if voxels_occ is not None:
voxels_occ = voxels_occ.to(device)
points_voxels = make_3d_grid(
(-0.5 + 1/64,) * 3, (0.5 - 1/64,) * 3, (32,) * 3)
points_voxels = points_voxels.expand(
batch_size, *points_voxels.size())
points_voxels = points_voxels.to(device)
with torch.no_grad():
p_out = self.model(points_voxels, inputs,
sample=self.eval_sample, **kwargs)
voxels_occ_np = (voxels_occ >= 0.5).cpu().numpy()
occ_hat_pred_sdf = p_out.logits / self.sdf_ratio
occ_hat_np = (occ_hat_pred_sdf <= threshold).cpu().numpy()
iou_voxels = compute_iou(voxels_occ_np, occ_hat_np).mean()
eval_dict['iou_voxels'] = iou_voxels
return eval_dict
def visualize(self, data):
''' Performs a visualization step for the data.
Args:
data (dict): data dictionary
'''
device = self.device
batch_size = data['points'].size(0)
shape = (32, 32, 32)
p = make_3d_grid([-0.5] * 3, [0.5] * 3, shape).to(device)
p = p.expand(batch_size, *p.size())
encoder_inputs, raw_data = compose_inputs(
data,
mode='val',
device=self.device,
input_type=self.input_type,
use_gt_depth_map=self.use_gt_depth_map,
depth_map_mix=self.depth_map_mix,
with_img=self.with_img,
depth_pointcloud_transfer=self.depth_pointcloud_transfer,
local=self.local)
kwargs = {}
with torch.no_grad():
p_r = self.model.forward_halfway(p,
encoder_inputs,
sample=self.eval_sample,
**kwargs)
occ_hat = p_r.probs.view(batch_size, *shape)
voxels_out = (occ_hat >= self.threshold).cpu().numpy()
# visualize
if self.local:
encoder_inputs = encoder_inputs[None]
if self.input_type == 'depth_pred':
gt_mask = raw_data['mask']
if self.with_img:
encoder_inputs = encoder_inputs['depth']
for i in trange(batch_size):
if self.use_gt_depth_map:
input_img_path = os.path.join(self.vis_dir,
'%03d_in_gt.png' % i)
else:
input_img_path = os.path.join(self.vis_dir,
'%03d_in_pr.png' % i)
depth_map = encoder_inputs[i].cpu()
depth_map = depth_to_L(depth_map, gt_mask[i].cpu())
vis.visualize_data(depth_map, 'img', input_img_path)
vis.visualize_voxels(
voxels_out[i], os.path.join(self.vis_dir, '%03d.png' % i))
elif self.input_type == 'depth_pointcloud':
for i in trange(batch_size):
input_pointcloud_file = os.path.join(
self.vis_dir, '%03d_depth_pointcloud.png' % i)
pc = encoder_inputs[i].cpu()
if self.depth_pointcloud_transfer in ('view',
'view_scale_model'):
vis.visualize_pointcloud(pc,
out_file=input_pointcloud_file,
elev=15,
azim=180)
else:
vis.visualize_pointcloud(pc,
out_file=input_pointcloud_file)
vis.visualize_voxels(
voxels_out[i], os.path.join(self.vis_dir, '%03d.png' % i))
def eval_step(self, data):
''' Performs an evaluation step.
Args:
data (dict): data dictionary
'''
self.model.eval()
device = self.device
threshold = self.threshold
eval_dict = {}
# Compute elbo
points = data.get('points').to(device)
occ = data.get('points.occ').to(device)
voxels_occ = data.get('voxels')
points_iou = data.get('points_iou').to(device)
occ_iou = data.get('points_iou.occ').to(device)
kwargs = {}
encoder_inputs, _ = compose_inputs(
data,
mode='val',
device=self.device,
input_type=self.input_type,
use_gt_depth_map=self.use_gt_depth_map,
depth_map_mix=self.depth_map_mix,
with_img=self.with_img,
depth_pointcloud_transfer=self.depth_pointcloud_transfer,
local=self.local)
with torch.no_grad():
elbo, rec_error, kl = self.model.compute_elbo_halfway(
points, occ, encoder_inputs, **kwargs)
eval_dict['loss'] = -elbo.mean().item()
eval_dict['rec_error'] = rec_error.mean().item()
eval_dict['kl'] = kl.mean().item()
# Compute iou
batch_size = points.size(0)
with torch.no_grad():
p_out = self.model.forward_halfway(points_iou,
encoder_inputs,
sample=self.eval_sample,
**kwargs)
occ_iou_np = (occ_iou >= 0.5).cpu().numpy()
occ_iou_hat_np = (p_out.probs >= threshold).cpu().numpy()
iou = compute_iou(occ_iou_np, occ_iou_hat_np).mean()
eval_dict['iou'] = iou
# Estimate voxel iou
if voxels_occ is not None:
voxels_occ = voxels_occ.to(device)
points_voxels = make_3d_grid((-0.5 + 1 / 64, ) * 3,
(0.5 - 1 / 64, ) * 3, (32, ) * 3)
points_voxels = points_voxels.expand(batch_size,
*points_voxels.size())
points_voxels = points_voxels.to(device)
with torch.no_grad():
p_out = self.model.forward_halfway(points_voxels,
encoder_inputs,
sample=self.eval_sample,
**kwargs)
voxels_occ_np = (voxels_occ >= 0.5).cpu().numpy()
occ_hat_np = (p_out.probs >= threshold).cpu().numpy()
iou_voxels = compute_iou(voxels_occ_np, occ_hat_np).mean()
eval_dict['iou_voxels'] = iou_voxels
return eval_dict
def eval_step(self, data):
''' Performs an evaluation step.
Args:
data (dict): data dictionary
'''
self.model.eval()
device = self.device
threshold = self.threshold
eval_dict = {}
# Compute elbo
points = data.get('points').to(device)
occ = data.get('points.occ').to(device)
inputs = data.get('inputs').to(device)
#gt_depth_maps = data.get('inputs.depth').to(device)
gt_mask = data.get('inputs.mask').to(device).byte()
voxels_occ = data.get('voxels')
points_iou = data.get('points_iou').to(device)
occ_iou = data.get('points_iou.occ').to(device)
kwargs = {}
with torch.no_grad():
elbo, rec_error, kl = self.model.compute_elbo(
points, occ, inputs, gt_mask, **kwargs)
eval_dict['loss'] = -elbo.mean().item()
eval_dict['rec_error'] = rec_error.mean().item()
eval_dict['kl'] = kl.mean().item()
# Compute iou
batch_size = points.size(0)
with torch.no_grad():
p_out = self.model(points_iou,
inputs,
gt_mask,
sample=self.eval_sample,
**kwargs)
occ_iou_np = (occ_iou >= 0.5).cpu().numpy()
occ_iou_hat_np = (p_out.probs >= threshold).cpu().numpy()
iou = compute_iou(occ_iou_np, occ_iou_hat_np).mean()
eval_dict['iou'] = iou
# Estimate voxel iou
if voxels_occ is not None:
voxels_occ = voxels_occ.to(device)
points_voxels = make_3d_grid((-0.5 + 1 / 64, ) * 3,
(0.5 - 1 / 64, ) * 3, (32, ) * 3)
points_voxels = points_voxels.expand(batch_size,
*points_voxels.size())
points_voxels = points_voxels.to(device)
with torch.no_grad():
p_out = self.model(points_voxels,
inputs,
gt_mask,
sample=self.eval_sample,
**kwargs)
voxels_occ_np = (voxels_occ >= 0.5).cpu().numpy()
occ_hat_np = (p_out.probs >= threshold).cpu().numpy()
iou_voxels = compute_iou(voxels_occ_np, occ_hat_np).mean()
eval_dict['iou_voxels'] = iou_voxels
return eval_dict
def eval_step(self, data):
""" Performs an evaluation step.
Args:
data (dict): data dictionary
"""
threshold = self.threshold
eval_dict = {}
# Compute elbo
points = data.get("points")
occ = data.get("points.occ")
inputs = data.get("inputs", tf.zeros([points.shape[0], 0]))
voxels_occ = data.get("voxels")
points_iou = data.get("points_iou")
occ_iou = data.get("points_iou.occ")
kwargs = {}
elbo, rec_error, kl = self.model.compute_elbo(points,
occ,
inputs,
training=False,
**kwargs)
eval_dict["loss"] = -float(tf.reduce_mean(elbo))
eval_dict["rec_error"] = float(tf.reduce_mean(rec_error))
eval_dict["kl"] = float(tf.reduce_mean(kl))
# Compute iou
batch_size = points.shape[0]
p_out = self.model(points_iou,
inputs,
sample=self.eval_sample,
training=False,
**kwargs)
occ_iou_np = (occ_iou >= 0.5).numpy()
occ_iou_hat_np = (p_out.probs_parameter() >= threshold).numpy()
iou = compute_iou(occ_iou_np, occ_iou_hat_np).mean()
eval_dict["iou"] = float(iou)
# Estimate voxel iou
if voxels_occ is not None:
points_voxels = make_3d_grid((-0.5 + 1 / 64, ) * 3,
(0.5 - 1 / 64, ) * 3, (32, ) * 3)
points_voxels = tf.broadcast_to(points_voxels,
[batch_size, *points_voxels.shape])
p_out = self.model(points_voxels,
inputs,
sample=self.eval_sample,
training=False,
**kwargs)
voxels_occ_np = (voxels_occ >= 0.5).numpy()
occ_hat_np = (p_out.probs_parameter() >= threshold).numpy()
iou_voxels = compute_iou(voxels_occ_np, occ_hat_np).mean()
eval_dict["iou_voxels"] = float(iou_voxels)
return eval_dict
def generate_from_latent(self,
z,
c=None,
stats_dict={},
data=None,
**kwargs):
''' Generates mesh from latent.
Args:
z (tensor): latent code z
c (tensor): latent conditioned code c
stats_dict (dict): stats dictionary
'''
assert data is not None
if self.is_explicit_mesh:
normal_faces = data.get('angles.normal_face').to(self.device)
normal_angles = data.get('angles.normal_angles').to(self.device)
if self.is_skip_surface_mask_generation_time:
t0 = time.time()
output = self.model.decode(None,
z,
c,
angles=normal_angles,
only_return_points=True,
**kwargs)
stats_dict['time (eval points)'] = time.time() - t0
t0 = time.time()
if not self.is_just_measuring_time:
output = self.model.decode(None,
z,
c,
angles=normal_angles,
**kwargs)
if not self.is_skip_surface_mask_generation_time:
stats_dict['time (eval points)'] = time.time() - t0
normal_vertices, normal_mask, _, _, _ = output
t0 = time.time()
B, N, P, D = normal_vertices.shape
normal_faces_all = torch.cat([(normal_faces + idx * P)
for idx in range(N)],
axis=1)
assert B == 1
mem_t = time.time()
verts = normal_vertices.view(
N * P, D).to('cpu').detach().numpy() / self.pnet_point_scale
faces = normal_faces_all.view(-1, 3).to('cpu').detach().numpy()
if self.is_just_measuring_time:
visbility = None
else:
visbility = (normal_mask > 0.5).view(
N * P).to('cpu').detach().numpy()
skip_t = time.time() - mem_t
mesh = trimesh.Trimesh(
verts,
faces,
process=False,
vertex_attributes={'vertex_visibility': visbility})
stats_dict['time (copy to trimesh)'] = time.time() - t0 - skip_t
else:
t0 = time.time()
threshold = 0.
#threshold = np.log(self.threshold) - np.log(1. - self.threshold)
# Compute bounding box size
box_size = 1 + self.padding
# Shortcut
if self.upsampling_steps == 0:
nx = self.resolution0
pointsf = box_size * make_3d_grid((-0.5, ) * 3, (0.5, ) * 3,
(nx, ) * 3)
values = self.eval_points(pointsf, z, c, data=data,
**kwargs).cpu().numpy()
value_grid = values.reshape(nx, nx, nx)
else:
mesh_extractor = MISE(self.resolution0, self.upsampling_steps,
threshold)
points = mesh_extractor.query()
while points.shape[0] != 0:
# Query points
pointsf = torch.FloatTensor(points).to(self.device)
# Normalize to bounding box
pointsf = pointsf / mesh_extractor.resolution
pointsf = box_size * (pointsf - 0.5)
# Evaluate model and update
values = self.eval_points(pointsf,
z,
c,
data=data,
**kwargs).cpu().numpy()
values = values.astype(np.float64)
mesh_extractor.update(points, values)
points = mesh_extractor.query()
value_grid = mesh_extractor.to_dense()
# Extract mesh
stats_dict['time (eval points)'] = time.time() - t0
mesh = self.extract_mesh(value_grid, z, c, stats_dict=stats_dict)
return mesh
def eval_step(self, data):
''' Performs an evaluation step.
Args:
data (dict): data dictionary
'''
self.model.eval()
device = self.device
threshold = self.threshold
eval_dict = {}
# Compute elbo
points = data.get('points').to(device)
#occ = data.get('points.occ').to(device)
inputs = data.get('inputs', torch.empty(points.size(0), 0)).to(device)
voxels_occ = data.get('voxels')
#occ_iou = data.get('points_iou.occ').to(device)
kwargs = {}
#with torch.no_grad():
# elbo, rec_error, kl = self.model.compute_elbo(
# points, occ, inputs, **kwargs)
target = data.get('points.point_lab').to(device)
# eval_dict['rec_error'] = rec_error.mean().item()
# eval_dict['kl'] = kl.mean().item()
# Compute iou
batch_size = points.size(0)
with torch.no_grad():
p_out, vote = self.model(points,
inputs,
sample=self.eval_sample,
**kwargs)
instance_loss = True
if instance_loss:
centers = data.get('points.centers').to(device)
vote_loss = (torch.max(vote.float() - centers.T.float(),
torch.tensor([0.]).cuda())**2).sum()
logits = p_out.logits
if self.loss_type == 'cross_entropy':
loss_i = F.cross_entropy(logits, target, reduction='none')
occ_iou_np = (target >= 0).cpu().numpy()
occ_iou_hat_np = p_out.probs.argmax(dim=1).cpu().numpy()
#iou = compute_iou(occ_iou_np, occ_iou_hat_np).mean()
#occ_iou_hat_np_04 = (p_out.probs >= 0.4).cpu().numpy()
iou = compute_iou(occ_iou_np, occ_iou_hat_np).mean()
eval_dict['iou'] = iou
#loss = loss_i.sum(-1).mean()
loss = loss_i.sum(-1).mean() + vote_loss * 1000
eval_dict['loss'] = loss.cpu().numpy()
# Estimate voxel iou
if voxels_occ is not None:
voxels_occ = voxels_occ.to(device)
points_voxels = make_3d_grid((-0.5 + 1 / 64, ) * 3,
(0.5 - 1 / 64, ) * 3, (32, ) * 3)
points_voxels = points_voxels.expand(batch_size,
*points_voxels.size())
points_voxels = points_voxels.to(device)
with torch.no_grad():
p_out, vote = self.model(points_voxels,
inputs,
sample=self.eval_sample,
**kwargs)
voxels_occ_np = (voxels_occ >= 0.5).cpu().numpy()
occ_hat_np = p_out.probs.argmax(dim=1).cpu().numpy()
iou_voxels = compute_iou(voxels_occ_np, occ_hat_np).mean()
eval_dict['iou_voxels'] = iou_voxels
return eval_dict
def generate_from_conditional(self, c, stats_dict={}, **kwargs):
''' Generates mesh from latent.
Args:
c (tensor): latent conditioned code c
stats_dict (dict): stats dictionary
'''
threshold = np.log(self.threshold) - np.log(1. - self.threshold)
t0 = time.time()
# Compute bounding box size
box_size = 1 + self.padding
# Shortcut
if self.upsampling_steps == 0:
nx = self.resolution0
pointsf = box_size * make_3d_grid((-0.5, ) * 3, (0.5, ) * 3,
(nx, ) * 3).to(self.device)
out_dict = self.eval_points(pointsf, c, **kwargs)
if self.tf_type is None:
value_minimal_grid = out_dict['values_minimal'].detach().cpu(
).numpy().reshape([nx, nx, nx])
value_cloth_grid = out_dict['values_cloth'].detach().cpu(
).numpy().reshape([nx, nx, nx])
label_grid = out_dict['labels'].detach().cpu().numpy().reshape(
[nx, nx, nx])
p_hat_grid = None
else:
value_minimal_grid = out_dict['values_minimal'].detach().cpu(
).numpy().reshape([nx, nx, nx])
value_cloth_grid = out_dict['values_cloth'].detach().cpu(
).numpy().reshape([nx, nx, nx])
label_grid = out_dict['labels'].detach().cpu().numpy().reshape(
[nx, nx, nx])
p_hat_grid = out_dict['p_hats'].detach().cpu().numpy().reshape(
[nx, nx, nx, 3])
else:
raise ValueError('We do not support MISE for double layer')
# Extract mesh
stats_dict['time (eval points)'] = time.time() - t0
mesh_all = {}
# Generate body under cloth
mesh_minimal = self.extract_mesh(value_minimal_grid,
label_grid,
p_hat_grid,
c,
stats_dict=stats_dict,
**kwargs)
for k, v in mesh_minimal.items():
mesh_all['minimal_' + k] = v
# Generate body with cloth
mesh_cloth = self.extract_mesh(value_cloth_grid,
label_grid,
p_hat_grid,
c,
stats_dict=stats_dict,
**kwargs)
for k, v in mesh_cloth.items():
mesh_all['cloth_' + k] = v
return mesh_all
def eval_step(self, data):
''' Performs an evaluation step.
Args:
data (dict): data dictionary
'''
self.model.eval()
device = self.device
threshold = self.threshold
eval_dict = {}
# Compute elbo
points = data.get('points').to(device)
occ = data.get('points.occ').to(device)
inputs = data.get('inputs', torch.empty(points.size(0), 0)).to(device)
voxels_occ = data.get('voxels')
points_iou = data.get('points_iou').to(device)
occ_iou = data.get('points_iou.occ').to(device)
'''TRANSFORMATION TO CAMERA SPACE'''
if self.camera_space:
transform = data.get('inputs.world_mat').to(device)
R = transform[:, :, :3]
points = transform_points(points, transform)
points_iou = transform_points(points_iou, R)
'''END'''
kwargs = {}
with torch.no_grad():
elbo, rec_error, kl = self.model.compute_elbo(
points, occ, inputs, **kwargs)
eval_dict['loss'] = -elbo.mean().item()
eval_dict['rec_error'] = rec_error.mean().item()
eval_dict['kl'] = kl.mean().item()
# Compute iou
batch_size = points.size(0)
with torch.no_grad():
p_out = self.model(points_iou,
inputs,
sample=self.eval_sample,
**kwargs)
occ_iou_np = (occ_iou >= 0.5).cpu().numpy()
occ_iou_hat_np = (p_out.probs >= threshold).cpu().numpy()
iou = compute_iou(occ_iou_np, occ_iou_hat_np).mean()
eval_dict['iou'] = iou
'''
with torch.no_grad():
p_out_r = self.model(points_iou_r, inputs,
sample=self.eval_sample, **kwargs)
occ_iou_r_hat_np = (p_out_r.probs >= threshold).cpu().numpy()
iou_r = compute_iou(occ_iou_np, occ_iou_r_hat_np).mean()
eval_dict['iou_r'] = iou_r
data['iou_r'] = iou_r
data['iou'] = iou
data['occ_iou_r_hat_np'] = occ_iou_r_hat_np
data['occ_iou_hat_np'] = occ_iou_hat_np
SAVE IOU DATA
# Create npz and save
im_path = str(data.get('inputs.image_path')[0]).split('/')
model=im_path[3]
im_nr=im_path[5].split('.')[0]
out_dir = 'IOU'
np.savez(os.path.join(out_dir, model+'_'+im_nr), iou_r=iou_r,
occ_iou_hat_np=occ_iou_hat_np,
occ_iou_r_hat_np=occ_iou_r_hat_np,
iou=iou,
occ_iou=occ_iou.cpu().numpy(),
transform=transform.cpu().numpy(),
points_iou=points_iou.cpu().numpy(),
points_iou_r=points_iou_r.cpu().numpy(),
model=model,
im_nr=im_nr,
image=inputs.cpu().numpy()
)
END'''
# Estimate voxel iou
if voxels_occ is not None:
voxels_occ = voxels_occ.to(device)
points_voxels = make_3d_grid((-0.5 + 1 / 64, ) * 3,
(0.5 - 1 / 64, ) * 3, (32, ) * 3)
points_voxels = points_voxels.expand(batch_size,
*points_voxels.size())
points_voxels = points_voxels.to(device)
with torch.no_grad():
p_out = self.model(points_voxels,
inputs,
sample=self.eval_sample,
**kwargs)
voxels_occ_np = (voxels_occ >= 0.5).cpu().numpy()
occ_hat_np = (p_out.probs >= threshold).cpu().numpy()
iou_voxels = compute_iou(voxels_occ_np, occ_hat_np).mean()
eval_dict['iou_voxels'] = iou_voxels
return eval_dict
