def __init__(self, Thresh=0.7, HallucinatePeeledColor=True):
super().__init__()
self.MaskLoss = nn.BCELoss(size_average=True, reduce=True)
self.Sigmoid = nn.Sigmoid()
self.Thresh = Thresh
self.HallucinatePeeledColor = HallucinatePeeledColor
self.ChamferDist = ChamferDistance()
import torch
import numpy as np
from tk3dv.extern.chamfer import ChamferDistance
chamfer_dist = ChamferDistance()
#...
# points and points_reconstructed are n_points x 3 matrices
points = torch.from_numpy(np.array([[[1, 1, 0], [2, 1, 0]]], dtype=np.float32))
#points.to(torch.device('cuda'))
points_reconstructed = torch.from_numpy(np.array([[[2, 2, 0]]], dtype=np.float32))
#points_reconstructed.to(torch.device('cuda'))
dist1, dist2 = chamfer_dist(points, points_reconstructed)
# outputs minimum squared distance for each point in the point cloud
print(dist1)
print(dist2)
loss = (torch.mean(dist1)) + (torch.mean(dist2))
print('loss: ' + str(loss))
total_dist = dist1 + dist2
cham = torch.sum(dist1) + torch.sum(dist2)
print('chamfer dist: ' + str(cham))
def test_shape_recon(model, test_loader, device, log_out, observed_steps,
unobserved_steps):
'''
Evaluates shape reconstruction from CNF.
'''
test_dataset = test_loader.dataset
chamfer_dist = ChamferDistance()
log(
log_out,
'Observed steps [%s]' % (','.join([str(idx)
for idx in observed_steps])))
log(
log_out, 'Unobserved steps [%s]' %
(','.join([str(idx) for idx in unobserved_steps])))
use_unobserved_steps = len(unobserved_steps) > 0
model.eval()
nfe_stats = []
model_ids = []
seq_ids = []
observed_stats = {'chamfer': [], 'emd': [], 'infer_time': []}
unobserved_stats = {'chamfer': [], 'emd': []}
num_batches_total = 0
for i, data in enumerate(test_loader):
print('Batch: %d / %d' % (i, len(test_loader)))
pcl_in, nocs_out = data[
0] # world point cloud, corresponding nocs point cloud
pcl_in = pcl_in.to(device) # B x T x N x 4 (x,y,z,t)
nocs_out = nocs_out.to(device) # B x T x N x 4 (x,y,z,t)
cur_model_ids = data[1]
cur_seq_ids = data[2]
model_ids.extend(cur_model_ids)
seq_ids.extend(cur_seq_ids)
# print(cur_model_ids)
# print(cur_seq_ids)
B, T, N, _ = pcl_in.size()
num_batches_total += B
T_observed = len(observed_steps)
T_unobserved = len(unobserved_steps)
if T != PROTOCOL_NUM_STEPS:
print('Test protocol requires %d steps, but %d given!' %
(PROTOCOL_NUM_STEPS, T))
exit()
if N != PROTOCOL_NUM_PTS:
print('Test protocol requires %d points, but %d given!' %
(PROTOCOL_NUM_PTS, N))
exit()
# only use the observed steps as input
observed_pcl_in = pcl_in[:, observed_steps, :, :]
elapsed = 0.0
start_t = time.time()
# reconstruct at all time steps, both observed and unobserved
_, _, pred_pcl, _ = model.reconstruct(observed_pcl_in,
num_points=N,
timestamps=nocs_out[0, :, 0, 3],
constant_in_time=False)
elapsed = time.time() - start_t
cur_nfe = model.get_nfe()
nfe_stats.append(cur_nfe)
# evaluate Chamfer and EMD
# first observed
observed_tnocs_gt = nocs_out[:, observed_steps, :, :3].view(
(B * T_observed, N, 3)) # don't need time stamp for reconstruction
observed_reconstr = pred_pcl[:, observed_steps, :, :].view(
(B * T_observed, N, 3))
mean_chamfer, cur_emd = eval_reconstr_frames(observed_reconstr,
observed_tnocs_gt,
chamfer_dist)
observed_stats['chamfer'].extend(mean_chamfer.tolist())
observed_stats['emd'].extend(cur_emd.tolist())
observed_stats['infer_time'].append(elapsed)
print('==== OBSERVED ====')
print('Shape Recon Mean Chamfer: %f' %
(np.mean(observed_stats['chamfer']) * 1000))
print('Shape Recon Median Chamfer: %f' %
(np.median(observed_stats['chamfer']) * 1000))
print('Shape Recon Mean EMD: %f' %
(np.mean(observed_stats['emd']) * 1000))
print('Shape Recon Median EMD: %f' %
(np.median(observed_stats['emd']) * 1000))
print('NFE Mean: (%f, %f)' %
tuple(np.mean(nfe_stats, axis=0).tolist()))
print('Infer time mean: %f' % (np.mean(observed_stats['infer_time'])))
if use_unobserved_steps:
unobserved_tnocs_gt = nocs_out[:, unobserved_steps, :, :3].view(
(B * T_unobserved, N,
3)) # don't need time stamp for reconstruction
unobserved_reconstr = pred_pcl[:, unobserved_steps, :, :].view(
(B * T_unobserved, N, 3))
mean_chamfer, cur_emd = eval_reconstr_frames(
unobserved_reconstr, unobserved_tnocs_gt, chamfer_dist)
unobserved_stats['chamfer'].extend(mean_chamfer.tolist())
unobserved_stats['emd'].extend(cur_emd.tolist())
print('==== UNOBSERVED ====')
print('Shape Recon Mean Chamfer: %f' %
(np.mean(unobserved_stats['chamfer']) * 1000))
print('Shape Recon Median Chamfer: %f' %
(np.median(unobserved_stats['chamfer']) * 1000))
print('Shape Recon Mean EMD: %f' %
(np.mean(unobserved_stats['emd']) * 1000))
print('Shape Recon Median EMD: %f' %
(np.median(unobserved_stats['emd']) * 1000))
print('NFE Mean: (%f, %f)' %
tuple(np.mean(nfe_stats, axis=0).tolist()))
# print(len(observed_stats['chamfer']))
# print(len(unobserved_stats['chamfer']))
# aggregate
stats_list = [observed_stats, unobserved_stats
] if use_unobserved_steps else [observed_stats]
stats_names = ['OBSERVED', 'UNOBSERVED'
] if use_unobserved_steps else ['OBSERVED']
for stat_dict, stats_name in zip(stats_list, stats_names):
mean_chamfer_err = np.mean(stat_dict['chamfer']) * 1000.0
median_chamfer_err = np.median(stat_dict['chamfer']) * 1000.0
std_chamfer_err = np.std(stat_dict['chamfer']) * 1000.0
mean_emd_err = np.mean(stat_dict['emd']) * 1000.0
median_emd_err = np.median(stat_dict['emd']) * 1000.0
std_emd_err = np.std(stat_dict['emd']) * 1000.0
log(
log_out,
'================ %s SAMPLING RECONSTR EVAL ====================='
% (stats_name))
log(
log_out, 'mean CHAMFER error (x1000): %f +- %f, median: %f' %
(mean_chamfer_err, std_chamfer_err, median_chamfer_err))
log(
log_out, 'mean EMD error (x1000): %f +- %f, median: %f' %
(mean_emd_err, std_emd_err, median_emd_err))
log(log_out,
'NFE Mean: (%f, %f)' % tuple(np.mean(nfe_stats, axis=0).tolist()))
log(log_out,
'mean Inference time: %f' % (np.mean(observed_stats['infer_time'])))
# save the evaluation data
per_seq_data_out = log_out[:-len('txt')] + 'npz'
np.savez(per_seq_data_out,
observed_chamfer=observed_stats['chamfer'],
observed_emd=observed_stats['emd'],
unobserved_chamfer=unobserved_stats['chamfer'],
unobserved_emd=unobserved_stats['emd'])
# log per-sequence performance
per_seq_log = log_out[:-len('txt')] + 'csv'
print('Per seq performance being saved to %s...' % (per_seq_log))
stats_steps = [T_observed, T_unobserved]
with open(per_seq_log, 'w', newline='') as csvfile:
# write header
csvwriter = csv.writer(csvfile,
delimiter=',',
quotechar='|',
quoting=csv.QUOTE_MINIMAL)
header = ['type', 'model_id', 'seq_id', 'chamfer', 'emd']
csvwriter.writerow(header)
for stat_dict, stats_name, stats_T in zip(stats_list, stats_names,
stats_steps):
per_seq_chamfer = np.array(stat_dict['chamfer']).reshape(
(num_batches_total, stats_T))
per_seq_chamfer = np.mean(per_seq_chamfer, axis=1)
per_seq_emd = np.array(stat_dict['emd']).reshape(
(num_batches_total, stats_T))
per_seq_emd = np.mean(per_seq_emd, axis=1)
for line_idx in range(len(model_ids)):
cur_line = [
stats_name, model_ids[line_idx], seq_ids[line_idx],
per_seq_chamfer[line_idx], per_seq_emd[line_idx]
]
csvwriter.writerow(cur_line)
return
def __init__(self, device='cuda', torch_device=None):
if torch_device is None:
self.device = torch.device(device) # 'cuda' or 'cpu'
else:
self.device = torch_device
self.ChamferDist = ChamferDistance()
def test_viz(cfg, model, test_dataset, test_loader, device):
'''
Visualize CaSPR results
'''
from tk3dv.extern.chamfer import ChamferDistance # only import if we need it
chamfer_dist = ChamferDistance()
model.eval()
for i, data in enumerate(test_loader):
print('Batch: %d / %d' % (i, len(test_loader)))
pcl_in, nocs_out = data[
0] # world point cloud, corresponding nocs point cloud
pcl_in = pcl_in.to(device) # B x T x N x 4 (x,y,z,t)
nocs_out = nocs_out.to(device) # B x T x N x 4 (x,y,z,t)
# print(nocs_out.size())
B, T, N, _ = pcl_in.size()
if B != 1:
print('batch size must be 1 to visualize!')
exit()
cur_model_id = data[1][0]
cur_seq_id = data[2][0]
print('Model %s' % (cur_model_id))
print('Seq %s' % (cur_seq_id))
#
# first show quantitative eval for context
#
pred_tnocs = None
if cfg.viz_tnocs and not (cfg.viz_observed or cfg.viz_interpolated):
_, pred_tnocs = model.encode(pcl_in)
else: # we need sampling predictions
samp_pcl, logprob_samp_pcl, pred_pcl, pred_tnocs = model.reconstruct(
pcl_in,
num_points=cfg.num_sampled_pts,
constant_in_time=cfg.constant_in_time,
max_timestamp=test_dataset.max_timestamp,
sample_contours=SAMPLE_CONTOURS_RADII
if cfg.sample_contours else None)
if cfg.viz_tnocs:
nocs_err = torch.norm(pred_tnocs[:, :, :, :3] -
nocs_out[:, :, :, :3],
dim=3).mean().to('cpu').item()
print('Cur L2 nocs spatial error: %f' % (nocs_err))
if cfg.viz_observed or cfg.viz_interpolated:
quant_num_pts = min([cfg.num_sampled_pts, N])
observed_tnocs_gt = nocs_out[:, :, :quant_num_pts, :3].view(
(B * T, quant_num_pts,
3)) # don't need time stamp for reconstruction
observed_reconstr = pred_pcl[:, :, :quant_num_pts, :].view(
(B * T, quant_num_pts, 3))
mean_chamfer, cur_emd = eval_reconstr_frames(
observed_reconstr, observed_tnocs_gt, chamfer_dist)
print('Cur Mean Chamfer: %f' % (np.mean(mean_chamfer) * 1000))
print('Cur Mean EMD: %f' % (np.mean(cur_emd) * 1000))
#
# Visualize
#
# needed by all visualizations
pcl_in_np, gt_nocs_np = torch_to_numpy([pcl_in, nocs_out])
viz_gt_nocs = np_to_list(gt_nocs_np)
viz_pcl_in = np_to_list(pcl_in_np)
gt_nocs_rgb = copy_pcl_list(viz_gt_nocs)
base_seq_to_viz = [viz_gt_nocs]
base_rgb_to_viz = [gt_nocs_rgb]
if cfg.show_input_seq:
base_seq_to_viz.append(viz_pcl_in)
base_rgb_to_viz.append(gt_nocs_rgb)
# TNOCS regression visualization
if cfg.viz_tnocs:
print('Visualizing TNOCS Regression Prediction...')
pred_nocs_np = torch_to_numpy([pred_tnocs])[0]
viz_pred_nocs = np_to_list(pred_nocs_np)
if cfg.tnocs_error_map:
pred_nocs_rgb = [
get_error_colors(viz_pred_nocs[idx], viz_gt_nocs[idx])
for idx in range(gt_nocs_np.shape[1])
]
else:
pred_nocs_rgb = copy_pcl_list(viz_pred_nocs)
# translate to be in predicted viz cube
viz_pred_nocs = shift_pcl_list(viz_pred_nocs, PRED_OFFSET)
seq_to_viz = base_seq_to_viz + [viz_pred_nocs]
rgb_to_viz = base_rgb_to_viz + [pred_nocs_rgb]
viz_pcl_seq(seq_to_viz,
rgb_seq=rgb_to_viz,
fps=T,
autoplay=True,
draw_cubes=cfg.show_nocs_cubes)
# Observed sampling visualization
if cfg.viz_observed:
print('Visualizing CaSPR Observed Reconstruction Sampling...')
viz_caspr_reconstruction(cfg, samp_pcl, logprob_samp_pcl, pred_pcl,
base_seq_to_viz, base_rgb_to_viz, T)
# Interpolated sampling visualization
if cfg.viz_interpolated:
print('Visualizing CaSPR Interpolated Reconstruction Sampling...')
timstamps = torch.linspace(0.0, 1.0,
cfg.num_sampled_steps).to(pcl_in)
# rerun reconstruction on higher-res timestamps
samp_pcl, logprob_samp_pcl, pred_pcl, _ = model.reconstruct(
pcl_in,
timestamps=timstamps,
num_points=cfg.num_sampled_pts,
constant_in_time=cfg.constant_in_time,
sample_contours=SAMPLE_CONTOURS_RADII
if cfg.sample_contours else None)
# naively subsample observations to visualize with interpolated result
subsampled_gt_nocs = []
subsampled_pcl_in = []
subsampled_times = []
subsamples_per_step = int(float(cfg.num_sampled_steps) / T)
for time_idx in range(T):
for repeat_idx in range(subsamples_per_step):
subsampled_gt_nocs.append(gt_nocs_np[0, time_idx, :, :3])
subsampled_pcl_in.append(pcl_in_np[0, time_idx, :, :3])
subsampled_times.append(gt_nocs_np[0, time_idx, 0, 3])
# fill any extras
while len(subsampled_gt_nocs) < cfg.num_sampled_steps:
subsampled_gt_nocs.append(gt_nocs_np[0, T - 1, :, :3])
subsampled_pcl_in.append(pcl_in_np[0, T - 1, :, :3])
subsampled_times.append(gt_nocs_np[0, T - 1, 0, 3])
viz_gt_nocs = subsampled_gt_nocs
viz_pcl_in = subsampled_pcl_in
gt_nocs_rgb = copy_pcl_list(viz_gt_nocs)
cur_base_seq_to_viz = [viz_gt_nocs]
cur_base_rgb_to_viz = [gt_nocs_rgb]
if cfg.show_input_seq:
cur_base_seq_to_viz.append(viz_pcl_in)
cur_base_rgb_to_viz.append(gt_nocs_rgb)
viz_caspr_reconstruction(cfg, samp_pcl, logprob_samp_pcl, pred_pcl,
cur_base_seq_to_viz, cur_base_rgb_to_viz,
cfg.num_sampled_steps)