def main(): sa = ShapeAssembly() lines = sa.load_lines(sys.argv[1]) # should be shape N x 3 target_pc = load_point_cloud(sys.argv[2]) out_file = sys.argv[3] hier, param_dict, param_list = sa.make_hier_param_dict(lines) opt = torch.optim.Adam(param_list, 0.001) start = torch.cat(param_list).clone() for iter in range(400): verts, faces = sa.diff_run(hier, param_dict) samps = sample_surface(faces, verts.unsqueeze(0), 10000) closs = cham_loss(samps.squeeze().T.unsqueeze(0).cuda(), target_pc.T.unsqueeze(0).cuda(), 0.0) ploss = (torch.cat(param_list) - start).abs().sum() loss = closs + ploss.cuda() * 0.001 opt.zero_grad() loss.backward() opt.step() if iter % 10 == 0: writeObj(verts, faces, f'{iter}_' + out_file + '.obj') writeObj(verts, faces, out_file + '.obj') sa.fill_hier(hier, param_dict) writeHierProg(hier, out_file + '.txt')
def create_point_cloud(in_file, out_file): sa = ShapeAssembly() lines = sa.load_lines(sys.argv[1]) hier, param_dict, _ = sa.make_hier_param_dict(lines) verts, faces = sa.diff_run(hier, param_dict) tsamps = sample_surface(faces, verts.unsqueeze(0), 10000).squeeze() writePC(tsamps, out_file)
def main(): sa = ShapeAssembly() lines = sa.load_lines(sys.argv[1]) # should be shape N x 3 tverts, tfaces = loadObj(sys.argv[2]) tverts = torch.tensor(tverts) tfaces = torch.tensor(tfaces).long() tsamps = sample_surface(tfaces, tverts.unsqueeze(0), 10000) out_file = sys.argv[3] hier, param_dict, param_list = sa.make_hier_param_dict(lines) start_time = time.time() opt = torch.optim.Adam(param_list, 0.001) start = torch.cat(param_list).clone() for iter in range(400): verts, faces = sa.diff_run(hier, param_dict) samps = sample_surface(faces, verts.unsqueeze(0), 10000) # tsamps = sample_surface(tfaces, tverts.unsqueeze(0), 10000) closs = cham_loss(samps.squeeze().T.unsqueeze(0).cuda(), tsamps.squeeze().T.unsqueeze(0).cuda(), 0.0) ploss = (torch.cat(param_list) - start).abs().sum() loss = closs + ploss.cuda() * 0.001 print(float(loss)) opt.zero_grad() loss.backward() opt.step() # if iter % 10 == 0: # writeObj(verts, faces, f'{iter}_' + out_file + '.obj') end_time = time.time() print(f"TIME: {end_time-start_time}") writeObj(verts, faces, out_file + '.obj') sa.fill_hier(hier, param_dict) writeHierProg(hier, out_file + '.txt')
def prog_to_pc(prog): verts, faces = hier_execute(prog) for i in range(3): verts[:, i] = verts[:, i] - verts[:, i].mean() pc = utils.sample_surface(faces, verts.unsqueeze(0), num_samps, return_normals=False)[0] return pc
def __getitem__(self, index): verts, faces = self.datas[index] for i in range(3): verts[:, i] = verts[:, i] - verts[:, i].mean() inputs = sample_surface(faces, verts.unsqueeze(0), 2500, return_normals=False)[0] return inputs
def prog_to_pc(prog, ns): if prog['prog'] == []: return torch.zeros((ns, 3)) verts, faces = hier_execute(prog) for i in range(3): verts[:, i] = verts[:, i] - verts[:, i].mean() pc = utils.sample_surface(faces, verts.unsqueeze(0), ns, return_normals=False)[0] return pc
def lines_to_pc(lines): P = Program() for i, line in enumerate(lines): P.execute(line) verts, faces = P.getShapeGeo() for i in range(3): verts[:, i] = verts[:, i] - verts[:, i].mean() pc = utils.sample_surface(faces, verts.unsqueeze(0), num_samps, return_normals=False)[0] return pc
def reencode_prog(full_pc, partial_prog, ns, encoder): full_labeled = full_pc.transpose(1, 0) pverts, pfaces = partial_prog.getShapeGeo() if pverts == None: partial_pc = torch.zeros((ns // 2, 3)) else: for i in range(3): pverts[:, i] = pverts[:, i] - pverts[:, i].mean() partial_pc = utils.sample_surface(pfaces, pverts.unsqueeze(0), ns // 2, return_normals=False)[0] partial_labels = torch.ones((ns // 2, 1)) partial_labeled = torch.cat([partial_pc, partial_labels], 1).to(device) pc = torch.cat([full_labeled, partial_labeled], 0) pc = pc.transpose(1, 0).unsqueeze(0) return encoder(pc).unsqueeze(0)
def eval_recon(outdir, data_inds): decoder = torch.load(shapeAssembly_decoder).to(device) decoder.eval() encoder = PCEncoder() encoder.load_state_dict(torch.load(point_cloud_encoder)) encoder.eval() encoder.to(device) os.system(f'mkdir {outdir}') count = 0. tdist = 0. for ind in tqdm(data_inds): pc_samp = torch.load(f'{point_cloud_folder}/{ind}.pts').to(device) enc = encoder(pc_samp.unsqueeze(0)) prog, _ = run_eval_decoder(enc.unsqueeze(0), decoder, False) verts, faces = hier_execute(prog) utils.writeObj(verts, faces, f"{outdir}/{ind}.obj") utils.writeHierProg(prog, f"{outdir}/{ind}.txt") verts = verts.to(device) faces = faces.to(device) pred_samp = utils.sample_surface(faces, verts.unsqueeze(0), 10000, True) # Center PC offset = (pc_samp.max(dim=0).values + pc_samp.min(dim=0).values) / 2 pc_samp -= offset #utils.writeSPC(pc_samp,f'tar_pc_{ind}.obj') #utils.writeSPC(pred_samp[0,:,:3],f'scripts/output/pred_pc_{ind}.obj') pc_samp = pc_samp.repeat(1, 2).unsqueeze(0) tdist += fscore.score(pred_samp.squeeze().T.unsqueeze(0), pc_samp.squeeze().T.unsqueeze(0)) count += 1 print(f"Average F-score: {tdist/count}")
def getFScore(verts, faces, gt_verts, gt_faces): p_samps = utils.sample_surface(faces, verts.unsqueeze(0), NUM_SAMPS) t_samps = utils.sample_surface(gt_faces, gt_verts.unsqueeze(0), NUM_SAMPS) return fscore.score(p_samps.squeeze().T.unsqueeze(0), t_samps.squeeze().T.unsqueeze(0))
os.system(f'mkdir {outdir}') for ind in inds: count += 1. try: hier = jp.parseJsonToHier(ind, CATEGORY) nshier = jp.parseJsonToHier(ind, CATEGORY, True) gen.generate_program(hier) pverts, pfaces = hier_execute(hier) tverts, tfaces = get_gt_geom(nshier, False) tsamps = utils.sample_surface(tfaces, tverts.unsqueeze(0), 10000) try: psamps = utils.sample_surface(pfaces, pverts.unsqueeze(0), 10000) pfs = fscore.score(psamps.squeeze().T.unsqueeze(0), tsamps.squeeze().T.unsqueeze(0)) except Exception: pfs = 0. if pfs >= 50: gpfsv += 1.
def fit(prog_path, obj_path, out_path): progs = os.listdir(prog_path) objs = os.listdir(obj_path) fitted_progs = [] for i, prg in enumerate(progs): print(f"fitting program {i}") sa = ShapeAssembly() p_no_e = prg.split("_")[1] index = int(p_no_e.split(".")[0]) # should be shape N x 3 tverts, tfaces = loadObj(f"{obj_path}/{index}.obj") tverts = torch.tensor(tverts) tfaces = torch.tensor(tfaces).long() out_file = f"{out_path}/{index}" with open(f"{prog_path}/{prg}") as file: lines = file.readlines() hier, param_dict, param_list = sa.make_hier_param_dict(lines) opt = torch.optim.Adam(param_list, 0.001) start = torch.cat(param_list).clone() for iter in range(400): verts, faces = sa.diff_run(hier, param_dict) samps = sample_surface(faces, verts.unsqueeze(0), 10000) tsamps = sample_surface(tfaces, tverts.unsqueeze(0), 10000) closs = cham_loss(samps.squeeze().T.unsqueeze(0).cuda(), tsamps.squeeze().T.unsqueeze(0).cuda(), 0.0) ploss = (torch.cat(param_list) - start).abs().sum() loss = closs + ploss.cuda() * 0.001 opt.zero_grad() loss.backward() opt.step() # # prevent cuboids from having 0 dimensions new_param_dict = {} for p in param_dict: new_p = [] for param in param_dict[p]: if param[0] == "Cuboid": new_attrs = [] for attr in param[1]: if torch.is_tensor(attr): new_attr = torch.clamp(attr, min=0.01).detach() new_attrs.append(new_attr) else: new_attrs.append(attr) new_p.append((param[0], new_attrs)) else: new_p.append(param) new_param_dict[p] = new_p sa.fill_hier(hier, new_param_dict) verts, faces = hier_execute(hier) writeObj(verts, faces, out_file + '.obj') writeHierProg(hier, out_file + '.txt') fitted_progs.append((hier, index)) return fitted_progs