def geometry_dist(obj, recon_obj):
ori_obbs_np = torch.cat([item.view(1, -1) for item in obj.boxes(leafs_only=True)], dim=0).cpu().numpy()
ori_mesh_v, ori_mesh_f = utils.gen_obb_mesh(ori_obbs_np)
ori_pc_sample = utils.sample_pc(ori_mesh_v, ori_mesh_f)
recon_obbs_np = torch.cat([item.view(1, -1) for item in recon_obj.boxes(leafs_only=True)], dim=0).cpu().numpy()
recon_mesh_v, recon_mesh_f = utils.gen_obb_mesh(recon_obbs_np)
recon_pc_sample = utils.sample_pc(recon_mesh_v, recon_mesh_f)
cd1, cd2 = chamfer_loss(
torch.tensor(ori_pc_sample, dtype=torch.float32).view(1, -1, 3),
torch.tensor(recon_pc_sample, dtype=torch.float32).view(1, -1, 3))
cd = ((cd1.sqrt().mean() + cd2.sqrt().mean()) / 2).item()
return cd
def compute_binary_diff(pred_node):
if pred_node.is_leaf:
return 0, 0
else:
binary_diff = 0; binary_tot = 0;
# all children
for cnode in pred_node.children:
cur_binary_diff, cur_binary_tot = compute_binary_diff(cnode)
binary_diff += cur_binary_diff
binary_tot += cur_binary_tot
# current node
if pred_node.edges is not None:
for edge in pred_node.edges:
pred_part_a_id = edge['part_a']
obb1 = pred_node.children[pred_part_a_id].box.cpu().numpy()
obb_quat1 = pred_node.children[pred_part_a_id].get_box_quat().cpu().numpy()
mesh_v1, mesh_f1 = utils.gen_obb_mesh(obb1)
pc1 = utils.sample_pc(mesh_v1, mesh_f1, n_points=500)
pc1 = torch.tensor(pc1, dtype=torch.float32, device=device)
pred_part_b_id = edge['part_b']
obb2 = pred_node.children[pred_part_b_id].box.cpu().numpy()
obb_quat2 = pred_node.children[pred_part_b_id].get_box_quat().cpu().numpy()
mesh_v2, mesh_f2 = utils.gen_obb_mesh(obb2)
pc2 = utils.sample_pc(mesh_v2, mesh_f2, n_points=500)
pc2 = torch.tensor(pc2, dtype=torch.float32, device=device)
if edge['type'] == 'ADJ':
dist1, dist2 = chamferLoss(pc1.view(1, -1, 3), pc2.view(1, -1, 3))
binary_diff += (dist1.sqrt().min().item() + dist2.sqrt().min().item()) / 2
elif 'SYM' in edge['type']:
if edge['type'] == 'TRANS_SYM':
mat1to2, _ = compute_sym.compute_trans_sym(obb_quat1.reshape(-1), obb_quat2.reshape(-1))
elif edge['type'] == 'REF_SYM':
mat1to2, _ = compute_sym.compute_ref_sym(obb_quat1.reshape(-1), obb_quat2.reshape(-1))
elif edge['type'] == 'ROT_SYM':
mat1to2, _ = compute_sym.compute_rot_sym(obb_quat1.reshape(-1), obb_quat2.reshape(-1))
else:
assert 'ERROR: unknown symmetry type: %s' % edge['type']
mat1to2 = torch.tensor(mat1to2, dtype=torch.float32, device=device)
transformed_pc1 = pc1.matmul(torch.transpose(mat1to2[:, :3], 0, 1)) + \
mat1to2[:, 3].unsqueeze(dim=0).repeat(pc1.size(0), 1)
dist1, dist2 = chamferLoss(transformed_pc1.view(1, -1, 3), pc2.view(1, -1, 3))
loss = (dist1.sqrt().mean() + dist2.sqrt().mean()) / 2
binary_diff += loss.item()
else:
assert 'ERROR: unknown symmetry type: %s' % edge['type']
binary_tot += 1
return binary_diff, binary_tot
def convert_to_structurenet_format(csg, boxes):
global idx
sem_name = csg['name']
if 'box_id' in csg:
out = dict({
'box_id': csg['box_id'],
'id': idx,
'label': sem_name,
'box': boxes[csg['box_id']].flatten().tolist(),
})
idx += 1
v, f = utils.gen_obb_mesh(np.reshape(boxes[csg['box_id']], [1, -1]))
pc = utils.sample_pc(v, f)
else:
children = []
cpcs = []
for citem in csg['parts']:
cnode, cpc = convert_to_structurenet_format(citem, boxes)
children.append(cnode)
cpcs.append(cpc)
pc = np.concatenate(cpcs, axis=0)
box = utils.fit_box(pc)
out = dict({
'id': idx,
'label': sem_name,
'box': box.tolist(),
'children': children,
})
return out, pc
def compute_gen_cd_numbers(in_dir, data_path, object_list, shapediff_topk,
shapediff_metric, self_is_neighbor, tot_shape):
chamfer_loss = ChamferDistance()
data_features = [
'object', 'name', 'neighbor_diffs', 'neighbor_objs', 'neighbor_names'
]
dataset = PartNetShapeDiffDataset(data_path, object_list, data_features,
shapediff_topk, shapediff_metric,
self_is_neighbor)
tot_gen = 100
bar = ProgressBar()
quality = 0.0
coverage = 0.0
for i in bar(range(tot_shape)):
obj, obj_name, neighbor_diffs, neighbor_objs, neighbor_names = dataset[
i]
mat = np.zeros((shapediff_topk, tot_gen), dtype=np.float32)
gt_pcs = []
for ni in range(shapediff_topk):
obbs_np = torch.cat([
item.view(1, -1)
for item in neighbor_objs[ni].boxes(leafs_only=True)
],
dim=0).cpu().numpy()
mesh_v, mesh_f = utils.gen_obb_mesh(obbs_np)
pc_sample = utils.sample_pc(mesh_v, mesh_f)
gt_pcs.append(np.expand_dims(pc_sample, axis=0))
gt_pcs = np.concatenate(gt_pcs, axis=0)
gt_pcs = torch.from_numpy(gt_pcs).float().cuda()
for i in range(tot_gen):
obj = PartNetDataset.load_object(
os.path.join(in_dir, obj_name, 'obj2-%03d.json' % i))
obbs_np = torch.cat(
[item.view(1, -1) for item in obj.boxes(leafs_only=True)],
dim=0).cpu().numpy()
mesh_v, mesh_f = utils.gen_obb_mesh(obbs_np)
gen_pc = utils.sample_pc(mesh_v, mesh_f)
gen_pc = np.tile(np.expand_dims(gen_pc, axis=0),
[shapediff_topk, 1, 1])
gen_pc = torch.from_numpy(gen_pc).float().cuda()
d1, d2 = chamfer_loss(gt_pcs.cuda(), gen_pc)
mat[:, i] = (d1.sqrt().mean(dim=1) +
d2.sqrt().mean(dim=1)).cpu().numpy() / 2
quality += mat.min(axis=0).mean()
coverage += mat.min(axis=1).mean()
np.save(os.path.join(in_dir, obj_name, 'cd_stats.npy'), mat)
quality /= tot_shape
coverage /= tot_shape
print('mean cd quality: %.5f' % quality)
print('mean cd coverage: %.5f' % coverage)
print('q + c: %.5f' % (quality + coverage))
with open(
os.path.join(in_dir,
'neighbor_%s_cd_stats.txt' % shapediff_metric),
'w') as fout:
fout.write('mean cd quality: %.5f\n' % quality)
fout.write('mean cd coverage: %.5f\n' % coverage)
fout.write('q + c: %.5f\n' % (quality + coverage))
# structure diff and edge accuracy
sd = compute_struct_diff(obj.root, recon_obj.root)
sd = sd / len(obj.root.boxes())
structure_dists.append(sd)
# save original and reconstructed object
os.mkdir(os.path.join(result_dir, obj_name))
orig_output_filename = os.path.join(result_dir, obj_name, 'orig.json')
recon_output_filename = os.path.join(result_dir, obj_name, 'recon.json')
PartNetDataset.save_object(obj=obj, fn=orig_output_filename)
PartNetDataset.save_object(obj=recon_obj, fn=recon_output_filename)
# chamfer distance
ori_obbs_np = torch.cat([item.view(1, -1) for item in obj.boxes(leafs_only=True)], dim=0).cpu().numpy()
ori_mesh_v, ori_mesh_f = utils.gen_obb_mesh(ori_obbs_np)
ori_pc_sample = utils.sample_pc(ori_mesh_v, ori_mesh_f)
print(ori_pc_sample.shape)
recon_obbs_np = torch.cat([item.view(1, -1) for item in recon_obj.boxes(leafs_only=True)], dim=0).cpu().numpy()
recon_mesh_v, recon_mesh_f = utils.gen_obb_mesh(recon_obbs_np)
recon_pc_sample = utils.sample_pc(recon_mesh_v, recon_mesh_f)
print(recon_pc_sample.shape)
cd1, cd2 = chamferLoss(torch.tensor(ori_pc_sample, dtype=torch.float32).view(1, -1, 3),
torch.tensor(recon_pc_sample, dtype=torch.float32).view(1, -1, 3))
cd = ((cd1.sqrt().mean() + cd2.sqrt().mean()) / 2).item()
chamfer_dists.append(cd)
stat_filename = os.path.join(result_dir, obj_name, 'stats.txt')
with open(stat_filename, 'w') as stat_file:
print(f'box pc chamfer distance: {cd}', file=stat_file)
print(f'structure distance: {sd}', file=stat_file)
# enumerate over all training shapes
num_shape = len(dataset)
n_points = 2048
print('Creating point clouds ...')
pcs = np.zeros((num_shape, n_points, 3), dtype=np.float32)
objs = []
names = []
bar = ProgressBar()
for i in bar(range(num_shape)):
obj, name = dataset[i]
objs.append(obj)
names.append(name)
obbs = torch.cat([item.view(1, -1) for item in obj.boxes(leafs_only=True)], dim=0).cpu().numpy()
mesh_v, mesh_f = utils.gen_obb_mesh(obbs)
pcs[i] = utils.sample_pc(mesh_v, mesh_f, n_points=n_points)
pcs = torch.tensor(pcs, dtype=torch.float32, device=device)
#np.save(os.path.join(out_dir, 'sd-%d.npy'%shape_id), sd_mat)
print('Computing distance ...')
bar = ProgressBar()
for i in bar(range(num_shape)):
pc1 = pcs[i:i+1].repeat(num_shape, 1, 1)
cd1, cd2 = chamferLoss(pc1, pcs)
dists = ((cd1.sqrt().mean(1) + cd2.sqrt().mean(1)) / 2).cpu().numpy()
np.save(
os.path.join(out_dir, names[i]+'.npy'),
{'dists': dists, 'names': names})
def compute_struct_diff(gt_node, pred_node):
if gt_node.is_leaf:
if pred_node.is_leaf:
return 0, 0, 0, 0, 0, 0
else:
return len(pred_node.boxes())-1, 0, 0, pred_node.get_subtree_edge_count(), 0, 0
else:
if pred_node.is_leaf:
return len(gt_node.boxes())-1, 0, gt_node.get_subtree_edge_count() * 2, 0, 0, 0
else:
gt_sem = set([node.label for node in gt_node.children])
pred_sem = set([node.label for node in pred_node.children])
intersect_sem = set.intersection(gt_sem, pred_sem)
gt_cnodes_per_sem = dict()
for node_id, gt_cnode in enumerate(gt_node.children):
if gt_cnode.label in intersect_sem:
if gt_cnode.label not in gt_cnodes_per_sem:
gt_cnodes_per_sem[gt_cnode.label] = []
gt_cnodes_per_sem[gt_cnode.label].append(node_id)
pred_cnodes_per_sem = dict()
for node_id, pred_cnode in enumerate(pred_node.children):
if pred_cnode.label in intersect_sem:
if pred_cnode.label not in pred_cnodes_per_sem:
pred_cnodes_per_sem[pred_cnode.label] = []
pred_cnodes_per_sem[pred_cnode.label].append(node_id)
matched_gt_idx = []; matched_pred_idx = []; matched_gt2pred = np.zeros((conf.max_child_num), dtype=np.int32)
for sem in intersect_sem:
gt_boxes = torch.cat([gt_node.children[cid].get_box_quat() for cid in gt_cnodes_per_sem[sem]], dim=0).to(device)
pred_boxes = torch.cat([pred_node.children[cid].get_box_quat() for cid in pred_cnodes_per_sem[sem]], dim=0).to(device)
num_gt = gt_boxes.size(0)
num_pred = pred_boxes.size(0)
if num_gt == 1 and num_pred == 1:
cur_matched_gt_idx = [0]
cur_matched_pred_idx = [0]
else:
gt_boxes_tiled = gt_boxes.unsqueeze(dim=1).repeat(1, num_pred, 1)
pred_boxes_tiled = pred_boxes.unsqueeze(dim=0).repeat(num_gt, 1, 1)
dmat = boxLoss(gt_boxes_tiled.view(-1, 10), pred_boxes_tiled.view(-1, 10)).view(-1, num_gt, num_pred).cpu()
_, cur_matched_gt_idx, cur_matched_pred_idx = utils.linear_assignment(dmat)
for i in range(len(cur_matched_gt_idx)):
matched_gt_idx.append(gt_cnodes_per_sem[sem][cur_matched_gt_idx[i]])
matched_pred_idx.append(pred_cnodes_per_sem[sem][cur_matched_pred_idx[i]])
matched_gt2pred[gt_cnodes_per_sem[sem][cur_matched_gt_idx[i]]] = pred_cnodes_per_sem[sem][cur_matched_pred_idx[i]]
struct_diff = 0.0; edge_both = 0; edge_gt = 0; edge_pred = 0;
gt_binary_diff = 0.0; gt_binary_tot = 0;
for i in range(len(gt_node.children)):
if i not in matched_gt_idx:
struct_diff += len(gt_node.children[i].boxes())
edge_gt += gt_node.children[i].get_subtree_edge_count() * 2
for i in range(len(pred_node.children)):
if i not in matched_pred_idx:
struct_diff += len(pred_node.children[i].boxes())
edge_pred += pred_node.children[i].get_subtree_edge_count()
for i in range(len(matched_gt_idx)):
gt_id = matched_gt_idx[i]
pred_id = matched_pred_idx[i]
cur_struct_diff, cur_edge_both, cur_edge_gt, cur_edge_pred, cur_gt_binary_diff, cur_gt_binary_tot = compute_struct_diff(gt_node.children[gt_id], pred_node.children[pred_id])
gt_binary_diff += cur_gt_binary_diff
gt_binary_tot += cur_gt_binary_tot
struct_diff += cur_struct_diff
edge_both += cur_edge_both
edge_gt += cur_edge_gt
edge_pred += cur_edge_pred
pred_node.children[pred_id].part_id = gt_node.children[gt_id].part_id
if pred_node.edges is not None:
edge_pred += len(pred_node.edges)
if gt_node.edges is not None:
edge_gt += len(gt_node.edges) * 2
pred_edges = np.zeros((conf.max_child_num, conf.max_child_num, len(conf.edge_types)), dtype=np.bool)
for edge in pred_node.edges:
pred_part_a_id = edge['part_a']
pred_part_b_id = edge['part_b']
edge_type_id = conf.edge_types.index(edge['type'])
pred_edges[pred_part_a_id, pred_part_b_id, edge_type_id] = True
for edge in gt_node.edges:
gt_part_a_id = edge['part_a']
gt_part_b_id = edge['part_b']
edge_type_id = conf.edge_types.index(edge['type'])
if gt_part_a_id in matched_gt_idx and gt_part_b_id in matched_gt_idx:
pred_part_a_id = matched_gt2pred[gt_part_a_id]
pred_part_b_id = matched_gt2pred[gt_part_b_id]
edge_both += pred_edges[pred_part_a_id, pred_part_b_id, edge_type_id]
edge_both += pred_edges[pred_part_b_id, pred_part_a_id, edge_type_id]
# gt edges eval
obb1 = pred_node.children[pred_part_a_id].box.cpu().numpy()
obb_quat1 = pred_node.children[pred_part_a_id].get_box_quat().cpu().numpy()
mesh_v1, mesh_f1 = utils.gen_obb_mesh(obb1)
pc1 = utils.sample_pc(mesh_v1, mesh_f1, n_points=500)
pc1 = torch.tensor(pc1, dtype=torch.float32, device=device)
obb2 = pred_node.children[pred_part_b_id].box.cpu().numpy()
obb_quat2 = pred_node.children[pred_part_b_id].get_box_quat().cpu().numpy()
mesh_v2, mesh_f2 = utils.gen_obb_mesh(obb2)
pc2 = utils.sample_pc(mesh_v2, mesh_f2, n_points=500)
pc2 = torch.tensor(pc2, dtype=torch.float32, device=device)
if edge_type_id == 0: # ADJ
dist1, dist2 = chamferLoss(pc1.view(1, -1, 3), pc2.view(1, -1, 3))
gt_binary_diff += (dist1.sqrt().min().item() + dist2.sqrt().min().item()) / 2
else: # SYM
if edge_type_id == 2: # TRANS_SYM
mat1to2, _ = compute_sym.compute_trans_sym(obb_quat1.reshape(-1), obb_quat2.reshape(-1))
elif edge_type_id == 3: # REF_SYM
mat1to2, _ = compute_sym.compute_ref_sym(obb_quat1.reshape(-1), obb_quat2.reshape(-1))
elif edge_type_id == 1: # ROT_SYM
mat1to2, _ = compute_sym.compute_rot_sym(obb_quat1.reshape(-1), obb_quat2.reshape(-1))
else:
assert 'ERROR: unknown symmetry type: %s' % edge['type']
mat1to2 = torch.tensor(mat1to2, dtype=torch.float32, device=device)
transformed_pc1 = pc1.matmul(torch.transpose(mat1to2[:, :3], 0, 1)) + \
mat1to2[:, 3].unsqueeze(dim=0).repeat(pc1.size(0), 1)
dist1, dist2 = chamferLoss(transformed_pc1.view(1, -1, 3), pc2.view(1, -1, 3))
loss = (dist1.sqrt().mean() + dist2.sqrt().mean()) / 2
gt_binary_diff += loss.item()
gt_binary_tot += 1
return struct_diff, edge_both, edge_gt, edge_pred, gt_binary_diff, gt_binary_tot