def getSqueezeProg(prog): P = Program() cuboid_lines = [] sym_lines = [] attach_lines = [] face_info = {} for line in prog: if "Cuboid(" in line: cuboid_lines.append(line) elif "reflect(" in line or "translate(" in line: sym_lines.append(line) elif "attach(" in line: attach_lines.append(line) parse = P.parseAttach(line) face, uv = getFace(torch.stack(parse[2:5]), torch.stack(parse[5:8]), parse[1]) if face is not None: spec_add(face_info, parse[0], {face: (uv, parse[1])}) squeeze_lines = [] for att_line in attach_lines: parse = P.parseAttach(att_line) face, uv = getFace(torch.stack(parse[2:5]), torch.stack(parse[5:8]), parse[1]) # not a face att so just skip if face is None: squeeze_lines.append(att_line) continue oface = getOppFace(face) cube = parse[0] # if this is missing it must have been removed previously by a squeeze attach if face not in face_info[cube]: continue if oface in face_info[cube] and \ (face_info[cube][oface][0] - uv).abs().sum() < OTHRESH: squeeze_lines.append( make_function('\tsqueeze', (cube, parse[1], face_info[cube][oface][1], face, uv[0].item(), uv[1].item()))) face_info[cube].pop(oface) face_info[cube].pop(face) else: squeeze_lines.append(att_line) ord_squeeze_lines = orderSqueeze(squeeze_lines) return cuboid_lines + ord_squeeze_lines + sym_lines
def orderSqueeze(lines): ord_lines = [] q = [] at_q = [] grounded = set(['bbox']) P = Program() for line in lines: if "attach(" in line: parse = P.parseAttach(line) if parse[1] in grounded: grounded.add(parse[0]) ord_lines.append(line) else: at_q.append((parse[0], parse[1], line)) elif "squeeze(" in line: parse = P.parseSqueeze(line) if parse[1] in grounded and parse[2] in grounded: ord_lines.append(line) grounded.add(parse[0]) else: q.append((parse[0], parse[1], parse[2], line)) torm = [] for i, (c, o1, o2, sl) in enumerate(q): if o1 in grounded and o2 in grounded: grounded.add(c) torm.append(i) ord_lines.append(sl) torm.sort(reverse=True) for i in torm: q.pop(i) atorm = [] for i, (c, o, al) in enumerate(at_q): if o in grounded: grounded.add(c) atorm.append(i) ord_lines.append(al) atorm.sort(reverse=True) for i in atorm: at_q.pop(i) assert len(q) == 0, 'Some squeezes are ungrounded' assert len(at_q) == 0, 'Some attaches are ungrounded' return ord_lines
def clean_prog(prog): P = Program() new_lines = [] for l in prog['prog']: if "Cuboid" in l: parse = P.parseCuboid(l) new_num = [round(x.item(), 3) for x in parse[1:4]] new_lines.append( f"{parse[0]} = Cuboid({new_num[0]}, {new_num[1]}, {new_num[2]}, {parse[4]})" ) elif "attach" in l: parse = P.parseAttach(l) new_num = [round(x.item(), 3) for x in parse[2:]] new_lines.append( f"attach({parse[0]}, {parse[1]}, {new_num[0]}," + f" {new_num[1]}, {new_num[2]}, {new_num[3]}, {new_num[4]}, {new_num[5]})" ) elif "squeeze" in l: parse = P.parseSqueeze(l) new_num = [round(x.item(), 3) for x in parse[-2:]] new_lines.append(f"squeeze({parse[0]}, {parse[1]}, {parse[2]}," + f" {parse[3]}, {new_num[0]}, {new_num[1]})") elif "translate" in l: parse = P.parseTranslate(l) new_num = [round(x.item(), 3) for x in parse[-1:]] new_lines.append( f"translate({parse[0]}, {parse[1]}, {parse[2]}, {new_num[0]})\n" ) elif "<END>" in l: pass else: new_lines.append(l) prog['prog'] = new_lines for c in prog["children"]: if not c == {}: clean_prog(c)
with open(f"data/chair/{f}", "r") as file: lines = file.readlines() src_count = 0 seen_cuboids = set() not_child = True for l in lines: # if "Assembly Program_1" in l: # not_child = False # if not_child: # continue # if "Cuboid" in l and ("bbox" in l): # src_count += 1 # parse = P.parseCuboid(l) # dim_list += parse[1:4] if "attach" in l and "bbox" in l: parse = P.parseAttach(l) # dim_list += [float(parse[-1])] if parse[0] not in seen_cuboids: src_count += 1 seen_cuboids.add(parse[0]) if "squeeze" in l and "bbox" in l: parse = P.parseSqueeze(l) if parse[0] not in seen_cuboids: src_count += 1 seen_cuboids.add(parse[0]) if "Assembly Program_1" in l: break # mean = torch.mean(torch.stack(parse[1:4])).item() # overall_mean += mean # num_lines += 1 # if max_dim > overall_max_dim: