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)
# 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: # overall_max_dim = max_dim # if not not_child: print(float(src_count)) src_avgs.append(float(src_count)) dim_tensor = torch.tensor(dim_list)