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