def fix_cube_count(lines):
P = Program()
cube_count = 0
switches = {}
new_lines = []
for l in lines:
if "Cuboid" in l and not ("bbox" in l):
parse = P.parseCuboid(l)
switches[parse[0]] = f"cube{cube_count}"
cube_count += 1
for l in lines:
locations = [i for (i, x) in enumerate(l) if l[i:i + 4] == "cube"]
new_line = ""
prev_idx = 0
for i in locations:
new_line += l[prev_idx:i]
try:
cube_num = int(l[i + 4:i + 6])
name_len = 6
if not len(l[i:i + name_len].strip()) == len(
l[i:i + name_len]):
name_len = 5
except Exception as e:
cube_num = int(l[i + 4:i + 5])
name_len = 5
if not l[i:i + name_len] in switches:
return None
new_line += switches[l[i:i + name_len]]
prev_idx = i + name_len
new_line += l[prev_idx:]
new_lines.append(new_line)
return new_lines
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 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 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 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 propsToProgram(gt):
T = Program()
for i, a in enumerate(gt):
c = Cuboid(str(i))
c.dims = torch.stack([
a['xd'],
a['yd'],
a['zd'],
])
c.pos = a['center']
c.rfnorm = a['xdir']
c.tfnorm = a['ydir']
c.ffnorm = a['zdir']
T.cuboids[str(i)] = c
return T
def canonical(lines):
P = Program()
def order(l):
if "Cuboid(" in l:
name = P.parseCuboid(l)[0]
if name == "bbox":
return 0
else:
return int(name[4:]) + 1
elif ("reflect" in l) or ("translate" in l):
return 1000
else:
return 100
lines.sort(key=order)
return lines
def lines_to_prog(lines):
P = Program()
num_children = len([l for l in lines if "Cuboid(" in l])
def order(l):
if "Cuboid(" in l:
name = P.parseCuboid(l)[0]
if name == "bbox":
return 0
else:
return int(name[4:]) + 1
else:
return 100
lines.sort(key=order)
prog = {"name": 0, "prog": lines, "children": [{}] * num_children}
return prog
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)
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
from ShapeAssembly import Program
import os
import torch
P = Program()
files = os.listdir("data/chair")
dim_list = []
src_avgs = []
for f in files:
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])
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
# parse = P.parseCuboid(l[1:-1])
# 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]})\n")
# 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]})\n")
# else:
# new_lines.append(l)
# with open(f"random_data_fixed/{f}", "w") as file:
# for l in new_lines:
# file.write(l)
files = os.listdir("random_hier_data")
P = Program()
for f in files:
prog = loadHPFromFile(f"random_hier_data/{f}")
def fix_lines(prog):
def order(l):
if "Cuboid(" in l:
name = P.parseCuboid(l)[0]
if name == "bbox":
return 0
else:
return int(name[4:]) + 1
elif ("reflect" in l) or ("translate" in l):
return 1000
else:
