neighbor_name = neighbor_names[ni]
neighbor_fn = 'neighbor_%02d_%s' % (ni, neighbor_name)
neighbor_obj = neighbor_objs[ni]
neighbor_obj.to(device)
neighbor_diff = neighbor_diffs[ni]
neighbor_diff.to(device)
with open(os.path.join(cur_res_dir, neighbor_fn+'.orig.diff'), 'w') as fout:
fout.write(str(neighbor_diff))
root_code = encoder.encode_tree_diff(obj, neighbor_diff)
recon_diff = decoder.decode_tree_diff(root_code, obj)
with open(os.path.join(cur_res_dir, neighbor_fn+'.recon.diff'), 'w') as fout:
fout.write(str(recon_diff))
recon_neighbor = Tree(Tree.apply_shape_diff(obj.root, recon_diff))
PartNetDataset.save_object(recon_neighbor, os.path.join(cur_res_dir, neighbor_fn+'.recon.json'))
cd = geometry_dist(neighbor_obj, recon_neighbor)
sd = struct_dist(neighbor_obj.root, recon_neighbor.root)
sd = sd / len(neighbor_obj.root.boxes())
with open(os.path.join(cur_res_dir, neighbor_fn+'.stats'), 'w') as fout:
fout.write('cd: %f\nsd: %f\n' % (cd, sd))
print('computing stats ...')
compute_recon_numbers(
in_dir=result_dir, baseline_dir=conf.baseline_dir, shapediff_topk=conf.shapediff_topk)
dirname=os.path.join(conf.ckpt_path, conf.exp_name),
epoch=conf.model_epoch,
strict=True)
# set models to evaluation mode
for m in models:
m.eval()
# test over all test shapes
with torch.no_grad():
objA = PartNetDataset.load_object(
os.path.join(conf.data_path, '%s.json' % conf.shapeA))
objB = PartNetDataset.load_object(
os.path.join(conf.data_path, '%s.json' % conf.shapeB))
objC = PartNetDataset.load_object(
os.path.join(conf.data_path, '%s.json' % conf.shapeC))
encoder.encode_tree(objA)
diffAB = Tree.compute_shape_diff(objA.root, objB.root)
code = encoder.encode_tree_diff(objA, diffAB)
encoder.encode_tree(objC)
recon_obj_diff = decoder.decode_tree_diff(code, objC)
recon_obj = Tree(Tree.apply_shape_diff(objC.root, recon_obj_diff))
PartNetDataset.save_object(objA, os.path.join(result_dir, 'shapeA.json'))
PartNetDataset.save_object(objB, os.path.join(result_dir, 'shapeB.json'))
PartNetDataset.save_object(objC, os.path.join(result_dir, 'shapeC.json'))
PartNetDataset.save_object(recon_obj,
os.path.join(result_dir, 'output.json'))
cur_res_dir = os.path.join(result_dir, obj_name)
if not os.path.exists(cur_res_dir):
os.mkdir(cur_res_dir)
obj.to(device)
PartNetDataset.save_object(obj,
os.path.join(cur_res_dir, 'cond_obj.json'))
encoder.encode_tree(obj)
zs = torch.randn(conf.num_gen, conf.latent_size).to(device)
for i in range(conf.num_gen):
recon_diff = decoder.decode_tree_diff(zs[i:i + 1], obj)
recon_obj2 = Tree(Tree.apply_shape_diff(obj.root, recon_diff))
PartNetDataset.save_object(
recon_obj2, os.path.join(cur_res_dir, 'obj2-%03d.json' % i))
with open(os.path.join(cur_res_dir, 'diff-%03d.txt' % i),
'w') as fout:
fout.write(str(recon_diff))
print('computing chamfer distance stats ...')
compute_gen_cd_numbers(in_dir=result_dir,
data_path=conf.data_path,
object_list=conf.test_dataset,
shapediff_topk=conf.shapediff_topk,
shapediff_metric=conf.shapediff_metric,
self_is_neighbor=conf.self_is_neighbor,
obj_D = PartNetDataset.load_object(
os.path.join(conf.data_path2, name_D + '.json')).to(device)
PartNetDataset.save_object(obj_D,
os.path.join(cur_res_dir, 'obj_D.json'))
# compute e := A --> B
diff_AB = Tree.compute_shape_diff(obj_A.root,
obj_B.root,
device=device)
encoder.encode_tree(obj_A)
e = encoder.encode_tree_diff(obj_A, diff_AB)
# compute Dp := e(C)
encoder.encode_tree(obj_C)
recon_diff_CD = decoder.decode_tree_diff(e, obj_C)
recon_D = Tree(Tree.apply_shape_diff(obj_C.root, recon_diff_CD))
PartNetDataset.save_object(recon_D,
os.path.join(cur_res_dir, 'recon_D.json'))
# compute dist between D and Dp
cd = geometry_dist(obj_D, recon_D)
sd = struct_dist(obj_D.root, recon_D.root) / len(obj_D.root.boxes())
with open(os.path.join(cur_res_dir, 'stats.txt'), 'w') as fout:
fout.write('cd: %f\nsd: %f\n' % (cd, sd))
tot_cd += cd
tot_sd += sd
tot_cnt += 1
avg_cd = tot_cd / tot_cnt
avg_sd = tot_sd / tot_cnt
