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