def get_train_pc(prog, ns): full_pc = prog_to_pc(prog, ns // 2) partial_prog = Program() pc_list = [get_line_pc(full_pc, partial_prog, ns)] for i in range(len(prog['prog'])): partial_prog.execute(prog['prog'][i]) pc_list.append(get_line_pc(full_pc, partial_prog, ns)) train_pc = torch.cat(pc_list, 0) return train_pc
def main(ind): sa = ShapeAssembly() root_lines, has_mid = make_skel() prog_lines = ['Assembly Program_0 {'] for b in root_lines: prog_lines.append('\t'+b) prog_lines.append('}') RP = Program() for l in root_lines: RP.execute(l) base_par = 'mid' if has_mid else 'top' base_lines = make_base_prog( RP.cuboids['base'], RP.cuboids[base_par] ) for i in range(len(prog_lines)): prog_lines[i] = prog_lines[i].replace('base', 'Program_1') prog_lines += base_lines #for b in base_lines: # prog_lines.append('\t'+b) #prog_lines.append('}') cube_count = -1 switches = [] for line in prog_lines: if 'Cuboid' in line: if not ('Program_' in line or "bbox" in line): switches.append(( f'cube{cube_count}', line.split()[0] )) if "bbox" in line: cube_count = -1 cube_count += 1 for a, b in switches: prog_lines = [line.replace(b,a) for line in prog_lines] hier_prog = make_hier_prog(prog_lines) verts, faces = hier_execute(hier_prog) if check_rooted(verts, faces) and check_stability(verts, faces): # writeObj(verts, faces, f'out_{ind}.obj') writeHierProg(hier_prog, f"random_hier_data/{ind}.txt") return True return False
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 rand_program(prog, max_cuboids, bbox_dims, hier_index): size_max = max(bbox_dims) num_cuboids = int(samp((max_cuboids / 1.5), 1, 3, max_cuboids)) num_source = int(samp(1.75, 0.5, 1, min(max_sources, num_cuboids))) cdim_mean = np.mean(bbox_dims) / 2 cdim_std = cdim_mean / 2.5 leaf_prob = 0.5**(hier_index + 1) print(f"NUM SOURCE: {num_source}") print(f"NUM CUBOIDS: {num_cuboids}") def attach_source(name, P): counter = 0 while counter < 100: old_prog = deepcopy(P) aligned = random.random() < align_prob cuboid_line, dims = make_cuboid(name, cdim_mean, cdim_std, size_max, aligned) # check if sqeezing into bbox is a good fit good_squeeze = ((dims[1] / bbox_dims[1]) > 0.9) if random.random() < squeeze_prob and (not aligned or good_squeeze): attach_line = make_squeeze(name, "bbox", "bot") else: # bottom of bbox and bottom of (unattached) source cuboid attach_line = make_attach(name, "bbox", "bot", "bot") lines = [cuboid_line, attach_line] P.execute(cuboid_line) P.execute(attach_line) if valid(P): break P = old_prog counter += 1 if counter >= 100: print("GAVE UP ON SOURCE") return P, [], None return P, lines, dims def extend_cuboid(src_cuboid, cuboids, P): faces = ['right', 'left', 'top', 'bot', 'front', 'back'] lines = [] new_cuboids = [] next_q = [] num_new_cuboids = int(samp(1, 0.5, 1, max_out)) for _ in range(num_new_cuboids): if new_cuboids == []: index = max([x["id"] for x in cuboids]) + 1 else: index = max([x["id"] for x in new_cuboids]) + 1 name = f"cube{index}" # make initial attachment to src cuboid if (len(cuboids) + len(new_cuboids)) < (num_cuboids - 1): counter = 0 while counter < 100: old_prog = deepcopy(P) old_new_cuboids = deepcopy(new_cuboids) old_next_q = deepcopy(next_q) aligned = random.random() < align_prob cuboid_line, dims = make_cuboid(name, cdim_mean, cdim_std, size_max, aligned) nc = { "name": f"cube{index}", "id": index, "ancestor": src_cuboid["name"], "dims": dims } new_cuboids.append(nc) attach_lines = [] if random.random() < squeeze_prob: # faces = ['top', 'bot'] # face = random.choice(faces) attach_line = make_squeeze(name, src_cuboid["name"], 'top') else: next_q.append(nc) attach_line = make_attach(name, src_cuboid["name"], random.choice(faces), random.choice(faces)) P.execute(cuboid_line) P.execute(attach_line) if valid(P): lines += [cuboid_line, attach_line] break P = old_prog new_cuboids = old_new_cuboids next_q = old_next_q counter += 1 if counter >= 100: print("GAVE UP ON EXTENSION") return [], [], P, [] # if cuboid is not aligned potentially add some more attachments counter = 0 while counter < 100 and not aligned: old_prog = deepcopy(P) num_extra_attaches = int(samp(1, 0.5, 0, 3)) print(f"EXTRA ATTACHES {num_extra_attaches}") attach_lines = [] for _ in range(num_extra_attaches): possible_cuboids = [ c["name"] for c in cuboids if not c["name"] == c["ancestor"] ] attach_cuboid = random.choice(possible_cuboids) attach_lines.append( make_attach(name, attach_cuboid, random.choice(faces), random.choice(faces))) for l in attach_lines: P.execute(l) if valid(P): lines += attach_lines break P = old_prog counter += 1 if counter >= 100: print("GAVE UP ON EXTENSION II") return new_cuboids, next_q, P, lines P = Program() bbox_line = f"bbox = Cuboid({bbox_dims[0]}, {bbox_dims[1]}, {bbox_dims[2]}, True)" print(bbox_line) P.execute(bbox_line) lines = [bbox_line] q = [] src_count = 0 for _ in range(num_source): P, new_lines, dims = attach_source(f"cube{src_count}", P) if len(new_lines) > 0: lines += new_lines for l in new_lines: print(l) q.append({ "name": f"cube{src_count}", "id": src_count, "ancestor": "bbox", "dims": dims }) src_count += 1 if len(q) == 0: print("COULDNT FIT ANY CUBOIDS") return None cuboids = deepcopy(q) while len(q) > 0 and len(cuboids) < (num_cuboids - 1): c = q.pop(0) new_cuboids, next_q, P, new_lines = extend_cuboid(c, cuboids, P) lines += new_lines for l in new_lines: print(l) q += next_q cuboids += new_cuboids # add some symmetry macros num_sym = int(samp(1, 0.5, 0, 3)) for _ in range(num_sym): counter = 0 while counter < 100: old_prog = deepcopy(P) sym_cuboid = random.choice([c["name"] for c in cuboids]) if random.random() < reflect_prob: new_line = make_reflect(sym_cuboid) else: new_line = make_translate(sym_cuboid) P.execute(new_line) if valid(P): lines.append(new_line) print(new_line) break P = old_prog counter += 1 if counter >= 100: print("GAVE UP ON MACRO") # correct dimensions since cuboids might have been scaled during execution for c in cuboids: new_dims = [round(x, 3) for x in P.cuboids[c['name']].dims.tolist()] c['dims'] = new_dims new_lines = [] for l in lines: if (c['name'] in l) and ("Cuboid" in l): aligned = P.parseCuboid(l)[-1] new_lines.append( f"{c['name']} = Cuboid({new_dims[0]}, {new_dims[1]}, {new_dims[2]}, {aligned})" ) else: new_lines.append(l) lines = new_lines # choose from the largest cuboids to expand non_bbox_cuboids = [ x for x in cuboids if not x['name'] == "bbox" and np.prod(x["dims"]) > 0.02 ] sorted_cuboids = sorted(non_bbox_cuboids, key=lambda x: -np.prod(x["dims"])) num_sub = len([ _ for _ in range(len(non_bbox_cuboids)) if random.random() < leaf_prob ]) # num_sub = len(sorted_cuboids) sub_cuboids = sorted_cuboids[:num_sub] # prog['prog'] = canonical(lines) prog['prog'] = lines next_q = [] # start of with bbox child children = [{}] for c in cuboids: if c in sub_cuboids: cprog = {"prog": None, "children": None} children.append(cprog) next_q.append((cprog, hier_index + 1, c["dims"])) else: children.append({}) prog['children'] = children return next_q
def semValidGen(pc, prog, encoder, decoder, max_lines, input_dim, device, gt_prog, rejection_sample): q = [] prog_out = [] preds = [] children = [] out = torch.zeros((1, 1, input_dim), dtype = torch.float).to(device) out[0][0][0] = 1.0 bb_dims = prog["bb_dims"] gt_nc_ind = 0 gt_children = [] if gt_prog is not None and len(gt_prog) > 0: gt_children = gt_prog["children"] P = Program() P.mode = 'start' P.grounded = set([0]) P.atts = {} c = 0 stop = False loops = 0 num_rejects = 0 h = decoder.init_gru full_pc = pc[0, :, :(num_samps // 2)] while(not stop and loops < max_lines): loops += 1 prev_out = out.clone().detach() reencoding = reencode_prog(full_pc, P, num_samps, encoder) if loops == 1: bb_pred = decoder.bbdimNet(reencoding.squeeze()) prog["bb_dims"] = bb_pred bb_dims = bb_pred out, h, valid = clean_forward( decoder, out, reencoding, h, bb_dims, input_dim, device, P ) if not valid: if rejection_sample and DO_REJECT: assert False, "Couldn't clean line" num_rejects += 1 if MASK_BAD_OUT and num_rejects < MAX_REJECT: out = prev_out else: num_rejects = 0 continue prog_out.append(out) line = out.clone().detach().squeeze() preds.append(line) command = torch.argmax(line[:7]) pline = None if command == 1: P.mode = 'cuboid' pline = getCuboidLine(line, c) c += 1 elif command == 2: P.mode = 'attach' cub1 = torch.argmax(line[7:18]).item() cub2 = torch.argmax(line[18:29]).item() P.grounded.add(cub1) if cub2 in P.atts: P.atts[cub2].append(cub1) else: P.atts[cub2] = [cub1] if cub1 in P.atts: P.atts[cub1].append(cub2) else: P.atts[cub1] = [cub2] pline = getAttachLines(line) elif command == 3: P.mode = 'sym' pline = getReflectLine(line) elif command == 4: P.mode = 'sym' pline = getTranslateLine(line) elif command == 5: P.mode = 'attach' cub1 = torch.argmax(line[7:18]).item() cub2 = torch.argmax(line[18:29]).item() cub3 = torch.argmax(line[29:40]).item() P.grounded.add(cub1) if cub2 in P.atts: P.atts[cub2].append(cub1) else: P.atts[cub2] = [cub1] if cub3 in P.atts: P.atts[cub3].append(cub1) else: P.atts[cub3] = [cub1] if cub1 in P.atts: P.atts[cub1].append(cub2) P.atts[cub1].append(cub3) else: P.atts[cub1] = [cub2, cub3] pline = getSqueezeLine(line) try: if pline is not None: P.execute(pline) except Exception: if VERBOSE: print("Unexpectedly, failed to execute line") pass # Stop at end token or when we have gone past max lines if command == 6: stop = True # If make a new Cuboid, use l to decide if it should have a child program or be a leaf if command == 1: children.append({}) fc_preds, fc_prog_out, fchildren = cuboid_line_clean(preds, prog_out, children, P) prog["children"] = fchildren return fc_preds, fc_prog_out
def rand_program(prog, max_cuboids, bbox_dims, hier_index): size_max = max(bbox_dims) num_cuboids = int(samp((max_cuboids / 3), 1, 2, max_cuboids)) num_source = int(samp(1.5, 1, 1, min(max_sources, num_cuboids))) cdim_mean = np.mean(bbox_dims) / 2 cdim_std = np.std(bbox_dims) leaf_prob = 0.5**(hier_index + 1) print(f"NUM SOURCE: {num_source}") print(f"NUM CUBOIDS: {num_cuboids}") def attach_source(name, P): counter = 0 while counter < 100: old_prog = deepcopy(P) cuboid_line = make_cuboid(name, cdim_mean, cdim_std, size_max, align_prob) # bottom of bbox and bottom of (unattached) source cuboid attach_line = make_attach(name, "bbox", "bot", "bot") lines = [cuboid_line, attach_line] P.execute(cuboid_line) P.execute(attach_line) if valid(P): break P = old_prog counter += 1 if counter >= 100: print("GAVE UP ON SOURCE") return P, [] return P, lines def create_edge(src_cuboid, cuboids, P, new_src_cuboids, prev_attach, new_cuboid): lines = [] faces = ['right', 'left', 'top', 'bot', 'front', 'back'] new_attach = None possible_attachments = [] for c in cuboids: # don't attach to ancestors, sources, or cuboids already attached to in this cycle if not (c["id"] in src_cuboid["ancestors"] or c["level"] == 1 or c["id"] in prev_attach): if c['name'] == 'bbox' and src_cuboid['level'] == 1: continue possible_attachments.append(c) if new_cuboid or possible_attachments == []: if len(P.cuboids) >= num_cuboids: print("TOO MANY CUBOIDS") return P, new_src_cuboids, [], new_attach if new_src_cuboids == []: index = max([x["id"] for x in cuboids]) + 1 else: index = max([x["id"] for x in new_src_cuboids]) + 1 new_src_cuboids.append({ "name": f"cube{index}", "id": index, "ancestors": src_cuboid["ancestors"] | set([index]), "level": src_cuboid["level"] + 1 }) cuboid_line = make_cuboid(f"cube{index}", cdim_mean, cdim_std, size_max, align_prob) lines = [cuboid_line] P.execute(cuboid_line) c_new = f"cube{index}" face1 = faces[random.choice(range(6))] face2 = faces[random.choice(range(6))] attach_line = make_attach(c_new, src_cuboid["name"], face1, face2) else: # pos_levels = [c['level'] for c in possible_attachments] # p = [-1*(x - max(pos_levels))+1.0 for x in pos_levels] # p = p / np.sum(p) selected_cuboid = np.random.choice(possible_attachments) # selected_cuboid['ancestors'] = selected_cuboid['ancestors'] | src_cuboid['ancestors'] src_cuboid['ancestors'] = selected_cuboid[ 'ancestors'] | src_cuboid['ancestors'] c_new = selected_cuboid["name"] new_attach = selected_cuboid["id"] if not c_new == "bbox": face1 = faces[random.choice(range(6))] face2 = faces[random.choice(range(6))] attach_line = make_attach(src_cuboid["name"], c_new, face1, face2) # if bbox only attach to top face and reverse direction of edge else: attach_line = make_attach(src_cuboid["name"], c_new, "top", "top") lines.append(attach_line) P.execute(attach_line) return P, new_src_cuboids, lines, new_attach def add_attachment(src_cuboid, cuboids, P, num_sources): p = np.full((max_out + 1, ), 1.0) p[0] = p[0] / 5 p[3] = p[3] / 5 p[2] = p[2] / 2 p = p / np.sum(p) out = np.random.choice(list(range(max_out + 1)), p=p) print(f"OUT: {out}") new_src_cuboids = [] lines = [] prev_attach = [] for edge in range(out): edge_counter = 0 new_cuboid = random.random() < new_cuboid_prob if len(P.cuboids) == num_cuboids: new_cuboid = False while edge_counter < 100: old_prog = deepcopy(P) old_new_src_cuboids = deepcopy(new_src_cuboids) P, new_src_cuboids, new_lines, new_attach = create_edge(src_cuboid, cuboids, \ P, new_src_cuboids, prev_attach, new_cuboid) if valid(P): if not new_cuboid: print("MADE ATTACHMENT WITH PREV CUBOID") else: print("MADE NEW CUBOID") prev_attach.append(new_attach) lines += new_lines break P = old_prog new_src_cuboids = old_new_src_cuboids edge_counter += 1 if not out == 0 and edge_counter >= 100: print("GAVE UP ON EDGE") return new_src_cuboids, P, lines P = Program() bbox_line = f"bbox = Cuboid({bbox_dims[0]}, {bbox_dims[1]}, {bbox_dims[2]}, True)" P.execute(bbox_line) lines = [bbox_line] current_cuboids = [] src_count = 0 for _ in range(num_source): P, new_lines = attach_source(f"cube{src_count}", P) if len(new_lines) > 0: lines += new_lines current_cuboids.append({ "name": f"cube{src_count}", "id": src_count, "ancestors": set([src_count]), "level": 1 }) src_count += 1 if len(current_cuboids) <= 1: print("COULDNT FIT MORE THAN ONE CUBOID") return None cuboids = [{ "name": "bbox", "id": -1, "ancestors": set([-1]), "level": 0 }] + current_cuboids while len(current_cuboids) > 0: c = current_cuboids.pop(0) new_cuboids, P, new_lines = add_attachment(c, cuboids, P, num_source) lines += new_lines current_cuboids += new_cuboids cuboids += new_cuboids # choose from the largest cuboids to expand non_bbox_cuboids = [ x for x in P.cuboids if not x == "bbox" and torch.prod(P.cuboids[x].dims) > 0.02 ] sorted_cuboids = sorted(non_bbox_cuboids, key=lambda x: -torch.prod(P.cuboids[x].dims)) # num_sub = len([_ for _ in range(len(non_bbox_cuboids)) if random.random() < leaf_prob]) num_sub = len(sorted_cuboids) sub_cuboids = sorted_cuboids[:num_sub] print(sub_cuboids) prog['prog'] = canonical(lines) next_q = [] children = [] for c in P.cuboids: if c in sub_cuboids: cprog = {"prog": None, "children": None} children.append(cprog) next_q.append((cprog, hier_index + 1, P.cuboids[c].dims.tolist())) else: children.append({}) prog['children'] = children return next_q
def semValidGen(prog, rnn, h, hier_ind, max_lines, input_dim, device, gt_prog, rejection_sample): q = [] prog_out = [] preds = [] children = [] out = torch.zeros((1, 1, input_dim), dtype=torch.float).to(device) out[0][0][0] = 1.0 h_start = h.clone() bb_dims = prog["bb_dims"] gt_nc_ind = 0 gt_children = [] if gt_prog is not None and len(gt_prog) > 0: gt_children = gt_prog["children"] P = Program() P.mode = 'start' P.grounded = set([0]) P.atts = {} c = 0 stop = False loops = 0 num_rejects = 0 while (not stop and loops < max_lines): loops += 1 prev_out = out.clone().detach() out, pnext, pleaf, h, valid = clean_forward(rnn, out, h, h_start, bb_dims, hier_ind, input_dim, device, P) if not valid: if rejection_sample and DO_REJECT: assert False, "Couldn't clean line" num_rejects += 1 if MASK_BAD_OUT and num_rejects < MAX_REJECT: out = prev_out else: num_rejects = 0 continue prog_out.append(out) line = out.clone().detach().squeeze() preds.append(line) command = torch.argmax(line[:7]) pline = None if command == 1: P.mode = 'cuboid' pline = getCuboidLine(line, c) c += 1 elif command == 2: P.mode = 'attach' cub1 = torch.argmax(line[7:18]).item() cub2 = torch.argmax(line[18:29]).item() P.grounded.add(cub1) if cub2 in P.atts: P.atts[cub2].append(cub1) else: P.atts[cub2] = [cub1] if cub1 in P.atts: P.atts[cub1].append(cub2) else: P.atts[cub1] = [cub2] pline = getAttachLines(line) elif command == 3: P.mode = 'sym' pline = getReflectLine(line) elif command == 4: P.mode = 'sym' pline = getTranslateLine(line) elif command == 5: P.mode = 'attach' cub1 = torch.argmax(line[7:18]).item() cub2 = torch.argmax(line[18:29]).item() cub3 = torch.argmax(line[29:40]).item() P.grounded.add(cub1) if cub2 in P.atts: P.atts[cub2].append(cub1) else: P.atts[cub2] = [cub1] if cub3 in P.atts: P.atts[cub3].append(cub1) else: P.atts[cub3] = [cub1] if cub1 in P.atts: P.atts[cub1].append(cub2) P.atts[cub1].append(cub3) else: P.atts[cub1] = [cub2, cub3] pline = getSqueezeLine(line) try: if pline is not None: P.execute(pline) except Exception: if VERBOSE: print("Unexpectedly, failed to execute line") pass # Stop at end token or when we have gone past max lines if command == 6: stop = True # If make a new Cuboid, use l to decide if it should have a child program or be a leaf if command == 1: gt_child = None if gt_nc_ind < len(gt_children): gt_child = gt_children[gt_nc_ind] gt_nc_ind += 1 # Skip BBox line if pleaf.squeeze().item() < 0 and len(preds) > 1: d = {"children": [], "bb_dims": line[40:43]} children.append(d) q.append((pnext, d, hier_ind + 1, gt_child)) else: children.append({}) fc_preds, fc_prog_out, fchildren = cuboid_line_clean( preds, prog_out, children, P) prog["children"] = fchildren return fc_preds, fc_prog_out, q