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 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
