def get_renderables(P, group_indices, vertex_count=100):
padding = 0.3
order = "xyz"
color_inactive = (0.8, 0.8, 0.8, 0.4)
max_depth = int(np.log(P.n_primitives)/np.log(2)) + 1
def get_center(depth, index):
width = (2.**depth) * (1+padding)
tree = dict(
x=0,
y=float(index)*(1+padding)-(width-1-padding)/2,
z=-float(depth)*(1+padding)
)
return np.array([[tree[axis] for axis in order]])
def children(d, i):
if d+1 < max_depth:
return children(d+1, 2*i) + children(d+1, 2*i+1)
else:
return [i]
def colors(d, i):
if d+1 < max_depth:
return colors(d+1, 2*i) + colors(d+1, 2*i+1)
else:
return [(i, get_color(d, i, 2**(max_depth-1)))]
meshes = []
lines = []
for d, i in group_indices:
meshes.append(Mesh.from_superquadrics(
*_unpack(
P,
get_center(d, i),
children(d, i),
colors(d, i),
color_inactive
),
vertex_count=vertex_count
))
if d > 0:
lines.extend([
get_center(d-1, i//2)[0],
get_center(d, i)[0]
])
meshes.append(Lines(lines, (0.1, 0.1, 0.1), width=0.05))
return meshes
def _renderables_from_flat_partition(y_hat, args):
_, P = y_hat.space_partition
# Collect the sqs that have prob larger than threshold
indices = [
i for i in range(y_hat.n_primitives)
if y_hat.probs_r[0, i] >= args.prob_threshold
]
active_prims_map = {(0, j): i for i, j in enumerate(indices)}
return [
Mesh.from_superquadrics(
*_unpack(
PrimitiveParameters.with_keys(
translations=P[-1].translations_r[0, indices],
rotations=P[-1].rotations_r[0, indices],
sizes=P[-1].sizes_r[0, indices] / 2.0,
shapes=P[-1].shapes_r[0, indices])), get_color(0, indices))
], indices
def _renderables_from_partition(y_hat, args):
[C, P] = y_hat.space_partition
active_prims_map = {
p: i
for (i, p) in enumerate(get_primitives_indices(P))
}
return [
Mesh.from_superquadrics(
*_unpack(
PrimitiveParameters.with_keys(
translations=P[depth].translations_r[:, idx],
rotations=P[depth].rotations_r[:, idx],
sizes=P[depth].sizes_r[:, idx],
shapes=P[depth].shapes_r[:, idx])),
get_color_groupped(depth, idx, args.max_depth)
if args.group_color else get_color(depth, idx))
for depth, idx in active_prims_map
], get_primitives_indices(P)
def _renderables_from_fit(y_hat, args):
F = y_hat.fit
active_prims = filter_primitives(
F,
qos_less(args.qos_threshold),
volume_larger(args.vol_threshold),
)
active_prims_map = {p: i for (i, p) in enumerate(active_prims)}
return [
Mesh.from_superquadrics(
*_unpack(
PrimitiveParameters.with_keys(
translations=F[depth].translations_r[:, idx],
rotations=F[depth].rotations_r[:, idx],
sizes=F[depth].sizes_r[:, idx],
shapes=F[depth].shapes_r[:, idx])),
get_color_groupped(depth, idx, args.max_depth)
if args.group_color else get_color(depth, idx))
for depth, idx in active_prims_map
], active_prims
def _renderables_from_flat_primitives(y_hat, args):
# Collect the sqs that have prob larger than threshold
indices = [
i for i in range(y_hat.n_primitives)
if y_hat.probs_r[0, i] >= args.prob_threshold
]
active_prims_map = {(0, j): i for i, j in enumerate(indices)}
if args.with_post_processing:
indices = get_non_overlapping_primitives(y_hat, indices)
return [
Mesh.from_superquadrics(
*_unpack(
PrimitiveParameters.with_keys(
translations=y_hat.translations_r[0, indices],
rotations=y_hat.rotations_r[0, indices],
sizes=y_hat.sizes_r[0, indices],
shapes=y_hat.shapes_r[0, indices])), get_color(0, indices))
], indices
X = sample[0].to(device)
y_hat = network(X)
F = y_hat.fit
[C, P] = y_hat.space_partition
n_primitives = y_hat.n_primitives
colors = torch.tensor(np.array(
plt.cm.jet(np.linspace(0, 1, 16))
))
if args.from_fit:
max_depth = 2**(len(F) - 1)
else:
max_depth = -1
meshes = [
[Mesh.from_superquadrics(
*_unpack(p, i, max_depth, colors, args.color_siblings)
)]
for i, p in enumerate(F if args.from_fit else P)
]
behaviours = [
CycleThroughObjects(meshes, interval=args.interval),
LightToCamera()
]
if args.save_frames:
behaviours += [
SaveFrames(args.save_frames, args.save_frequency)
]
if args.with_rotating_camera:
behaviours += [
CameraTrajectory(
Circle([0, 0, 1], [4, 0, 1], [0, 0, 1]), 0.001
def get_filtered_renderables(tree, group_indices, qos_threshold, vol_threshold,
padding=0.3, order="xyz", no_lines=False,
visualize_as_tree=False,
color_inactive=(0.8, 0.8, 0.8, 1.0),
group_color=False, vertex_count=100,
line_offset=[0, 0, 0]):
max_depth = max(d for (d, i) in group_indices)
def get_center(depth, index):
if visualize_as_tree:
max_width = (2.**max_depth) * (1+padding)
width = (2.**(max_depth - depth)) * (1+padding)
tree = dict(
x=0,
y=float(index)*(width)-(max_width-width)/2,
z=-float(depth)*(1+padding)
)
else:
width = (2.**depth) * (1+padding)
tree = dict(
x=0,
y=float(index)*(1+padding)-(width-1-padding)/2,
z=-float(depth)*(1+padding)
)
return np.array([[tree[axis] for axis in order]])
def get_ancestors(depth, index, storage):
storage.add((depth, index))
pdepth = depth-1
pindex = index//2
parent = pdepth, pindex
if parent not in storage:
get_ancestors(pdepth, pindex, storage)
active_prims = filter_primitives(
tree,
qos_less(qos_threshold),
volume_larger(vol_threshold)
)
active_prims_map = {p: i for (i, p) in enumerate(active_prims)}
valid_nodes = set([(0, 0)])
for d, i in active_prims:
get_ancestors(d, i, valid_nodes)
def children(d, i):
if (d, i) in active_prims_map:
return [active_prims_map[d, i]]
elif d < len(tree):
return children(d+1, 2*i) + children(d+1, 2*i+1)
else:
return []
def colors(d, i):
max_d = len(tree)
if (d, i) in active_prims_map:
return [(active_prims_map[d, i], get_color(d, i, 2**(max_d-1)))]
elif d < max_d:
return colors(d+1, 2*i) + colors(d+1, 2*i+1)
else:
return []
for d, i in group_indices:
if d > 0 and (d, i) in valid_nodes and len(children(d-1, i//2)) == 1:
valid_nodes.remove((d, i))
P = primitive_parameters_from_indices(tree, active_prims)
meshes = []
lines = []
for d, i in group_indices:
if (d, i) not in valid_nodes:
continue
sq_colors = colors(d, i)
if group_color:
sq_colors = [
(j, get_color(d, i))
for j, _ in sq_colors
]
if len(group_indices) == 1:
meshes.append(Mesh.from_superquadrics(
*_unpack(
P,
get_center(0, 0),
children(d, i),
sq_colors,
color_inactive
),
vertex_count=vertex_count,
offset=get_center(0, 0)
))
else:
meshes.append(Mesh.from_superquadrics(
*_unpack(
P,
get_center(d, i),
children(d, i),
sq_colors,
color_inactive
),
vertex_count=vertex_count,
offset=get_center(d, i)
))
if d > 0 and not no_lines:
lines.extend([
get_center(d-1, i//2)[0] + np.asarray(line_offset),
get_center(d, i)[0] + np.asarray(line_offset)
])
meshes.append(Lines(lines, (0.1, 0.1, 0.1), width=0.05))
if no_lines:
meshes.pop()
return meshes