def eval_mesh2(self, mesh, pointcloud_tgt, normals_tgt,
points_iou, occ_tgt):
''' Evaluates a mesh.
Args:
mesh (trimesh): mesh which should be evaluated
pointcloud_tgt (numpy array): target point cloud
normals_tgt (numpy array): target normals
points_iou (numpy_array): points tensor for IoU evaluation
occ_tgt (numpy_array): GT occupancy values for IoU points
'''
if len(mesh.vertices) != 0 and len(mesh.faces) != 0:
pointcloud, idx = mesh.sample(self.n_points, return_index=True)
pointcloud = pointcloud.astype(np.float32)
normals = mesh.face_normals[idx]
else:
pointcloud = np.empty((0, 3))
normals = np.empty((0, 3))
out_dict = self.eval_pointcloud(
pointcloud, pointcloud_tgt, normals, normals_tgt)
if len(mesh.vertices) != 0 and len(mesh.faces) != 0:
occ = check_mesh_contains(mesh, points_iou)
out_dict['iou'] = compute_iou(occ, occ_tgt)
else:
out_dict['iou'] = 0.
return out_dict
def eval_step_iou(self, data, c_s=None, c_t=None):
''' Calculates the IoU score for an evaluation test set item.
Args:
data (dict): data dictionary
c_s (tensor): spatial conditioned code
c_t (tensor): temporal conditioned code
z (tensor): latent shape code
z_t (tensor): latent motion code
'''
device = self.device
threshold = np.log(self.threshold) - np.log(1. - self.threshold)
eval_dict = {}
pts_iou = data.get('points').to(device)
occ_iou = data.get('points.occ').squeeze(0)
pts_iou_t = data.get('points.time').to(device)
batch_size, n_steps, n_pts, dim = pts_iou.shape
####unevn
index = torch.arange(n_steps).float().to(device)
pts_iou_t = (index / float(n_steps - 1))[None, :].expand(
batch_size, n_steps)
print('eval uneven', pts_iou.shape, pts_iou_t.shape)
logits_pts_iou_nsteps = []
for i in range(n_steps):
pts_iou_at_t0 = self.model.transform_to_t0(pts_iou_t[:, i],
pts_iou[:, i],
c_t=c_t)
logits_pts_iou_t = self.model.decode(pts_iou_at_t0,
c=c_s[:, 0, :]).logits
logits_pts_iou_nsteps.append(logits_pts_iou_t)
logits_t0 = torch.stack(logits_pts_iou_nsteps, dim=1)
# Calculate predicted occupancy values
rec_error = F.binary_cross_entropy_with_logits(
logits_t0.view(-1, n_pts),
occ_iou.to(device).view(-1, n_pts),
reduction='none')
rec_error = rec_error.mean(-1)
rec_error = rec_error.view(batch_size, n_steps).mean(0)
occ_pred = (logits_t0 > threshold).view(batch_size, n_steps,
n_pts).cpu().numpy()
# Calculate IoU
occ_gt = (occ_iou >= 0.5).numpy()
iou = compute_iou(occ_pred.reshape(-1, n_pts),
occ_gt.reshape(-1, n_pts))
iou = iou.reshape(batch_size, n_steps).mean(0)
eval_dict['iou'] = iou.sum() / len(iou)
eval_dict['rec_error'] = rec_error.sum().item() / len(rec_error)
for i in range(len(iou)):
eval_dict['iou_t%d' % i] = iou[i]
eval_dict['rec_error_t%d' % i] = rec_error[i].item()
return eval_dict
def eval_step(self, data, model_dict=None):
''' Performs an evaluation step.
Args:
data (dict): data dictionary
'''
self.model.eval()
device = self.device
threshold = self.threshold
eval_dict = {}
points_iou = data.get('points_iou').to(device)
occ_iou = data.get('points_iou.occ').to(device)
root_locs = data.get('points_iou.root_loc').to(device)
trans = data.get('points_iou.trans').to(device)
loc = data.get('points_iou.loc').to(device)
# bone_transforms = data.get('points_iou.bone_transforms').to(device)
# bone_transforms_inv = data.get('points_iou.bone_transforms_inv').to(device) # B x num_joints x 4 x 4
batch_size, T, D = points_iou.size()
occ_iou = occ_iou[:, :]
kwargs = {}
scale = data.get('points_iou.scale').to(device)
kwargs.update({'scale': scale.view(-1, 1, 1)}) #, 'bone_transforms_inv': bone_transforms_inv})
with torch.no_grad():
# Encoder inputs
inputs = data.get('inputs', torch.empty(1, 1, 0)).to(device)
mask = torch.ones(batch_size, T, dtype=points_iou.dtype, device=points_iou.device)
# Decode occupancies
out_dict = self.model(points_iou, inputs, **kwargs)
logits = out_dict['logits']
if len(logits.shape) == 4:
# PTF-piecewise predictions
logits = torch.max(logits, dim=1)[0]
p_out = dist.Multinomial(logits=logits.transpose(1, 2))
elif len(logits.shape) == 3:
# IPNet/PTF predictions
p_out = dist.Multinomial(logits=logits.transpose(1, 2))
else:
raise ValueError('Wrong logits shape')
# Compute iou
occ_iou_np = ((occ_iou >= 0.5) * mask).cpu().numpy()
if len(logits.shape) == 3:
# IoU for outer surface; we just want an easy-to-compute indicator for model selection
occ_iou_hat_np = ((p_out.probs[:, :, 1:].sum(-1) >= threshold) * mask).cpu().numpy()
else:
raise ValueError('Wrong logits shape')
iou = compute_iou(occ_iou_np, occ_iou_hat_np).mean()
eval_dict['iou'] = iou
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
occ = data.get('voxels').to(device)
inputs = data.get('inputs').to(device)
points = data.get('points_iou')
points_occ = data.get('points_iou.occ')
with torch.no_grad():
occ_logits = self.model(inputs).squeeze(1)
eval_dict = {}
# Compute loss
occ_hat = torch.sigmoid(occ_logits)
loss = F.binary_cross_entropy_with_logits(occ_logits, occ)
eval_dict['loss'] = loss.item()
# Compute discretized IOU
occ_np = (occ >= 0.5).cpu().numpy()
occ_hat_np = (occ_hat >= threshold).cpu().numpy()
iou_voxels = compute_iou(occ_np, occ_hat_np).mean()
eval_dict['iou_voxels'] = iou_voxels
# Compute continuous IOU (if possible)
if points is not None:
voxel_grids = [
VoxelGrid(occ_hat_np_i) for occ_hat_np_i in occ_hat_np
]
points_np = points.cpu().numpy()
points_occ_np = (points_occ >= 0.5).cpu().numpy()
points_occ_hat_np = np.stack(
[vg.contains(p) for p, vg in zip(points_np, voxel_grids)])
iou = compute_iou(points_occ_np, points_occ_hat_np).mean()
eval_dict['iou'] = iou
return eval_dict
def eval_step_iou(self, data, c_s=None, c_t=None, z=None, z_t=None):
''' Calculates the IoU score for an evaluation test set item.
Args:
data (dict): data dictionary
c_s (tensor): spatial conditioned code
c_t (tensor): temporal conditioned code
z (tensor): latent shape code
z_t (tensor): latent motion code
'''
device = self.device
threshold = self.threshold
eval_dict = {}
pts_iou = data.get('points').to(device)
occ_iou = data.get('points.occ').squeeze(0)
pts_iou_t = data.get('points.time').to(device)
batch_size, n_steps, n_pts, dim = pts_iou.shape
# Transform points from later time steps back to t=0
pts_iou_t0 = torch.stack([
self.model.transform_to_t0(pts_iou_t[:, i], pts_iou[:, i], z_t,
c_t) for i in range(n_steps)
],
dim=1)
# Reshape latent codes and predicted points tensor
c_s = c_s.unsqueeze(1).repeat(1, n_steps,
1).view(batch_size * n_steps, -1)
z = z.unsqueeze(1).repeat(1, n_steps, 1).view(batch_size * n_steps, -1)
pts_iou_t0 = pts_iou_t0.view(batch_size * n_steps, n_pts, dim)
# Calculate predicted occupancy values
p_r = self.model.decode(pts_iou_t0, z, c_s)
rec_error = -p_r.log_prob(occ_iou.to(device).view(-1, n_pts)).mean(-1)
rec_error = rec_error.view(batch_size, n_steps).mean(0)
occ_pred = (p_r.probs > threshold).view(batch_size, n_steps,
n_pts).cpu().numpy()
# Calculate IoU
occ_gt = (occ_iou >= 0.5).numpy()
iou = compute_iou(occ_pred.reshape(-1, n_pts),
occ_gt.reshape(-1, n_pts))
iou = iou.reshape(batch_size, n_steps).mean(0)
eval_dict['iou'] = iou.sum() / len(iou)
eval_dict['rec_error'] = rec_error.sum().item() / len(rec_error)
for i in range(len(iou)):
eval_dict['iou_t%d' % i] = iou[i]
eval_dict['rec_error_t%d' % i] = rec_error[i].item()
return eval_dict
def calc_mask_intersection(self, mask_gt, mask_pred, loss={}):
''' Calculates th intersection and IoU of provided mask tensors.
Args:
mask_gt (tensor): GT mask
mask_pred (tensor): predicted mask
loss (dict): loss dictionary
'''
mask_intersection = (mask_gt == mask_pred).float().mean()
mask_iou = compute_iou(mask_gt.cpu().float(),
mask_pred.cpu().float()).mean()
loss['mask_intersection'] = mask_intersection
loss['mask_iou'] = mask_iou
def eval_step_iou(self, data, c_t=None, z=None):
''' Calculates the IoU for the evaluation step.
Args:
data (tensor): training data
c_t (tensor): temporal conditioned latent code
z (tensor): latent code
'''
device = self.device
threshold = self.threshold
eval_dict = {}
pts_iou = data.get('points_iou').to(device)
occ_iou = data.get('points_iou.occ')
pts_iou_t = data.get('points_iou.time').to(device)
batch_size, n_steps, n_pts, dim = pts_iou.shape
c_t = c_t.unsqueeze(1).repeat(1, n_steps, 1)
z = z.unsqueeze(1).repeat(1, n_steps, 1)
t_axis = pts_iou_t.unsqueeze(2).unsqueeze(3).repeat(1, 1, n_pts, 1)
p = torch.cat([pts_iou, t_axis], dim=-1)
# Reshape network inputs
p = p.view(batch_size * n_steps, n_pts, -1)
z = z.view(batch_size * n_steps, z.shape[-1])
c_t = c_t.view(batch_size * n_steps, c_t.shape[-1])
occ_iou = occ_iou.view(batch_size * n_steps, n_pts)
occ_pred = self.model.decode(p, z, c_t)
occ_pred = (occ_pred.probs > threshold).cpu().numpy()
occ_gt = (occ_iou >= 0.5).numpy()
iou = compute_iou(occ_pred, occ_gt)
iou = iou.reshape(batch_size, -1).mean(0)
eval_dict['iou'] = iou.sum() / len(iou)
for i in range(len(iou)):
eval_dict['iou_t%d' % i] = iou[i]
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 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
'''
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 compute_loss(self, data):
''' Computes the loss.
Args:
data (dict): data dictionary
'''
device = self.device
p = data.get('points').to(device)
n_p_occ = p.size(1)
p_corr = data.get('points.points_corr')
if p_corr is not None:
p_corr = p_corr.to(device)
n_p_corr = p_corr.size(1)
p = torch.cat([p, p_corr], dim=1)
else:
n_p_corr = 0
batch_size = p.size(0)
occ = data.get('points.occ').to(device)
# bone_transforms = data.get('points.bone_transforms').to(device) # B x num_joints x 4 x 4
# bone_transforms_inv = data.get('points.bone_transforms_inv').to(device) # B x num_joints x 4 x 4
root_locs = data.get('points.root_loc').to(device)
trans = data.get('points.trans').to(device)
loc = data.get('points.loc').to(device)
kwargs = {}
scale = data.get('points.scale').to(device)
kwargs.update({'scale': scale.view(-1, 1, 1)}) #, 'bone_transforms_inv': bone_transforms_inv})
inputs = data.get('inputs', torch.empty(1, 1, 0)).to(device)
z = None # no prior
loss_dict = {}
if self.model_type in ['ipnet', 'deformable']:
# Models for IP-Net and PTF
# Get labels for skinning loss
s_targets = data.get('points.skinning_inds').to(self.device).long()
s_targets_ipnet = data.get('points.skinning_inds_ipnet').to(self.device).long()
s_targets_smpl = data.get('points.skinning_inds_smpl').to(self.device).long()
# Network computation
out_dict = self.model(p, inputs, **kwargs)
logits = out_dict['logits']
p_hat = out_dict['p_hat']
p_hat_template = out_dict.get('p_hat_template', None) # get predicted offset in canonical space, if available
if n_p_corr > 0:
logits = logits[:, :, :n_p_occ]
p_hat = p_hat[:, :, n_p_occ:]
if p_hat_template is not None:
p_hat_template = p_hat_template[:, :, n_p_occ:]
# Occupancy loss
if self.max_operator == 'lse':
occ_logits = torch.logsumexp(logits, dim=1) # smoothed-max over logits, for PTF-piecewise only
elif self.max_operator == 'max':
occ_logits = torch.max(logits, dim=1)[0] # max over logits, for PTF-piecewise only
elif self.max_operator == 'softmax':
occ_logits = logits # do nothing, softmax was already applied inside IP-Net/PTF decoder
else:
raise ValueError('Max operator type {} is not supported'.format(self.max_operator))
if self.occ_loss_type == 'ce':
if len(occ_logits.shape) == 3:
# Double-layer prediction
loss_occ = F.cross_entropy(
occ_logits, occ, reduction='none').sum(-1).mean()
else:
# Single-layer prediction
loss_occ = F.binary_cross_entropy_with_logits(
occ_logits, occ, reduction='none').sum(-1).mean()
else:
raise ValueError('Occupancy loss type {} is not supported'.format(self.occ_loss_type))
# Compute IoUs for this batch
if len(occ_logits.shape) == 3:
# Double-layer prediction
occ_hat_np = ((F.softmax(occ_logits, dim=1)[:, 1:, :].sum(1)) > 0.5).detach().cpu().numpy()
else:
# Single-layer prediction
occ_hat_np = (torch.sigmoid(occ_logits) > 0.5).detach().cpu().numpy()
occ_np = (occ > 0.5).detach().cpu().numpy()
ious = compute_iou(occ_np, occ_hat_np).flatten()
loss_dict['iou'] = ious
# Skinning loss with PTFs predictions
if self.max_operator == 'softmax':
# For IP-Net/PTF
out_cls = out_dict['out_cls']
if n_p_corr > 0:
out_cls = out_cls[:, :, n_p_occ:]
if self.skin_loss_type == 'ce':
loss_skin = F.cross_entropy(out_cls, s_targets, reduction='none').sum(-1).mean()
else:
raise ValueError('Skinning loss type {} is not supported'.format(self.skin_loss_type))
elif self.max_operator in ['lse']:
# For PTF-piecewise
if len(logits.shape) == 4:
# Double layer prediction
skin_logits = torch.logsumexp(logits, dim=2) # smoothed-max over class (i.e. inside, in-between, outside) dimension
else:
# Single layer prediction
skin_logits = logits
if self.skin_loss_type == 'ce':
loss_skin = F.cross_entropy(skin_logits, s_targets, reduction='none').sum(-1).mean()
else:
raise ValueError('Skinning loss type {} is not supported'.format(self.skin_loss_type))
else:
raise ValueError('Max operator type {} is not supported'.format(self.max_operator))
if self.use_corr_loss_pred: # and self.tf_type is not None:
parts_softmax = out_dict.get('parts_softmax', None)
# Compute p_hat
if parts_softmax is not None:
# For IP-Net/PTF, who has a separate branch to predict part probabilities
if n_p_corr > 0:
parts_softmax = parts_softmax[:, :, n_p_occ:]
p_hat = p_hat.view(batch_size, -1, 3, n_p_corr if n_p_corr > 0 else n_p_occ)
p_hat = p_hat * parts_softmax.view(batch_size, -1, 1, n_p_corr if n_p_corr > 0 else n_p_occ)
if p_hat_template is not None:
p_hat_template = p_hat_template.view(batch_size, -1, 3, n_p_corr if n_p_corr > 0 else n_p_occ)
p_hat_template = p_hat_template * parts_softmax.view(batch_size, -1, 1, n_p_corr if n_p_corr > 0 else n_p_occ)
else:
# For PTF-piecewise, who does not predict part probabilities directly
p_hat = p_hat.view(batch_size, -1, 3, n_p_corr if n_p_corr > 0 else n_p_occ)
p_hat = p_hat * F.softmax(skin_logits, dim=1).view(batch_size, -1, 1, n_p_corr if n_p_corr > 0 else n_p_occ)
if p_hat_template is not None:
p_hat_template = p_hat_template.view(batch_size, -1, 3, n_p_corr if n_p_corr > 0 else n_p_occ)
p_hat_template = p_hat_template * F.softmax(skin_logits, dim=1).view(batch_size, -1, 1, n_p_corr if n_p_corr > 0 else n_p_occ)
p_hat = p_hat.sum(1).transpose(1, 2)
if p_hat_template is not None:
p_hat_template = p_hat_template.sum(1).transpose(1, 2)
# Correspondence loss with predicted skinning inds
loss_corr_dict = {}
if self.use_corr_loss_pred: # and self.tf_type is not None:
p_a_pose = data.get('points.pts_a_pose').to(device)
p_template = data.get('points.pts_template')
if p_template is not None:
p_template = p_template.to(device)
loss_corr_pred_all = torch.norm(p_hat - p_a_pose, 2, dim=-1)
loss_corr_pred = loss_corr_pred_all.sum(-1).mean()
if p_hat_template is not None:
loss_template_pred = torch.norm(p_hat_template - p_template, 2, dim=-1).sum(-1).mean()
else:
loss_template_pred = torch.zeros(1, device=self.device)
else:
loss_corr_pred = torch.zeros(1, device=self.device)
loss_template_pred = torch.zeros(1, device=self.device)
loss_dict['occ_loss'] = loss_occ
loss_dict['skinning_loss'] = loss_skin
loss_dict['corr_loss'] = loss_corr_pred
loss_dict['template_loss'] = loss_template_pred
loss_dict.update(loss_corr_dict)
# Total weighted sum
loss = loss_occ + self.skinning_weight * loss_skin + self.corr_weight * (loss_corr_pred + loss_template_pred)
loss_dict['total_loss'] = loss
else:
raise ValueError('Supported model type: ipnet, deformable, got {}'.format(self.model_type))
return loss_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)
'''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
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 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
