def visualize_layers(reachable_layers):
layers = {}
for (x, y, z), label in label_matrix.voxels.items():
pos = Vec3(x, y, z)
if label == 0 or label not in reachable_layers:
continue
if label in layers:
layers[label][pos] = True
continue
layers[label] = np.zeros(label_matrix.shape, dtype=bool)
layers[label][pos] = True
color_gen = randomcolor.RandomColor()
colors = color_gen.generate(count=len(layers), format_='rgb')
for i in range(len(colors)):
values = map(lambda e: int(e), re.findall('\\d+', colors[i]))
values = list(values)
colors[i] = (values[0] / 255, values[1] / 255, values[2] / 255)
meshes = []
i = 0
for label in layers:
color = colors[i]
meshes.append(Mesh.from_voxel_grid(layers[label], colors=color))
i += 1
show(meshes)
def render_dfaust(scene, prev_renderable, seq, target):
new_renderable = Mesh.from_file(seq)
scene.remove(prev_renderable)
scene.add(new_renderable)
scene.render()
save_frame(target, scene.frame)
return new_renderable
def test_obj(self):
s = Scene(size=(1024, 1024), background=(0, 0, 0, 0))
s.add(Mesh.from_file(path.join(path.dirname(__file__), "plane.obj")))
s.rotate_x(np.pi / 2)
s.camera_position = (-1, -1, -1)
for i in range(180):
s.render()
imwrite("/tmp/frame_{:03d}.png".format(i), s.frame)
s.rotate_y(np.pi / 90)
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 test_ply(self):
s = Scene(size=(1024, 1024), background=(0, 0, 0, 0))
s.add(
Mesh.from_file(path.join(path.dirname(__file__),
"primitives.ply")))
s.up_vector = (0, -1, 0)
s.camera_position = (1, 0.5, 1)
light_initial = s.light
r_light = np.sqrt(light_initial[0]**2 + light_initial[-1]**2)
for i in range(180):
theta = i * np.pi / 90
r = np.sqrt(2)
s.camera_position = (r * np.cos(theta), 0.5, r * np.sin(theta))
s.light = (r_light * np.cos(theta), light_initial[1],
r_light * np.sin(theta))
s.render()
imwrite("/tmp/frame_{:03d}.png".format(i), s.frame)
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 test_cube(self):
points = np.array([[1, 1, 1], [1, 1, -1], [1, -1, 1], [1, -1, -1],
[-1, 1, 1], [-1, 1, -1], [-1, -1, 1], [-1, -1, -1]
]) * 0.5
hull = ConvexHull(points)
vertices = points[hull.simplices.ravel()].reshape(-1, 3)
normals = (np.ones(
(1, 3, 1)) * hull.equations[:, np.newaxis, :3]).reshape(-1, 3)
colors = np.ones((len(vertices), 3)) * [0.1, 0.5, 0.8]
s = Scene(size=(1024, 1024), background=(0, 0, 0, 0))
s.add(Mesh(vertices, normals, colors))
for i in range(180):
s.render()
imwrite("/tmp/frame_{:03d}.png".format(i), s.frame)
s.rotate_z(np.pi / 90)
s.rotate_x(np.pi / 180)
s.rotate_y(np.pi / 360)
s.camera_position = s.camera_position - 0.01
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
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
def main(argv):
parser = argparse.ArgumentParser(
description="Do the forward pass and visualize the recovered parts")
parser.add_argument(
"config_file",
help="Path to the file that contains the experiment configuration")
parser.add_argument("--output_directory",
default="../demo/output/",
help="Save the output files in that directory")
parser.add_argument(
"--weight_file",
default=None,
help=("The path to a previously trained model to continue"
" the training from"))
parser.add_argument("--prediction_file",
default=None,
help="The path to the predicted primitives")
parser.add_argument("--background",
type=lambda x: list(map(float, x.split(","))),
default="1,1,1,1",
help="Set the background of the scene")
parser.add_argument("--up_vector",
type=lambda x: tuple(map(float, x.split(","))),
default="0,0,1",
help="Up vector of the scene")
parser.add_argument("--camera_target",
type=lambda x: tuple(map(float, x.split(","))),
default="0,0,0",
help="Set the target for the camera")
parser.add_argument("--camera_position",
type=lambda x: tuple(map(float, x.split(","))),
default="-2.0,-2.0,-2.0",
help="Camera position in the scene")
parser.add_argument("--window_size",
type=lambda x: tuple(map(int, x.split(","))),
default="512,512",
help="Define the size of the scene and the window")
parser.add_argument("--with_rotating_camera",
action="store_true",
help="Use a camera rotating around the object")
parser.add_argument("--mesh",
action="store_true",
help="Visualize the target mesh")
parser.add_argument("--n_vertices",
type=int,
default=10000,
help="How many vertices to use per part")
add_dataset_parameters(parser)
args = parser.parse_args(argv)
if torch.cuda.is_available():
device = torch.device("cuda:0")
else:
device = torch.device("cpu")
print("Running code on", device)
# Check if output directory exists and if it doesn't create it
if not os.path.exists(args.output_directory):
os.makedirs(args.output_directory)
config = load_config(args.config_file)
# Extract the number of primitives
n_primitives = config["network"]["n_primitives"]
# Dictionary to keep the predictions used for the evaluation
predictions = {}
if args.prediction_file is None:
# Create a dataset instance to generate the input samples
dataset_directory = config["data"]["dataset_directory"]
dataset_type = config["data"]["dataset_type"]
train_test_splits_file = config["data"]["splits_file"]
dataset = dataset_factory(
config["data"]["dataset_factory"],
(ModelCollectionBuilder(config).with_dataset(dataset_type).
filter_category_tags(args.category_tags).filter_tags(
args.model_tags).random_subset(
args.random_subset).build(dataset_directory)),
)
assert len(dataset) == 1
# Build the network architecture to be used for training
network, _, _ = build_network(config, args.weight_file, device=device)
network.eval()
# Create the prediction input
with torch.no_grad():
for sample in dataset:
sample = [s[None] for s in sample] # make a batch dimension
X = sample[0].to(device)
targets = [yi.to(device) for yi in sample[1:]]
F = network.compute_features(X)
phi_volume, _ = network.implicit_surface(F, targets[0])
y_pred, faces = network.points_on_primitives(
F,
args.n_vertices,
random=False,
mesh=True,
union_surface=False)
predictions["phi_volume"] = phi_volume
predictions["y_prim"] = y_pred
else:
preds = torch.load(args.prediction_file, map_location="cpu")
y_pred = preds[4]
faces = preds[5]
targets = preds[0]
predictions["phi_volume"] = preds[1]
predictions["y_prim"] = y_pred
# Get the renderables from the deformed vertices and faces
vertices = y_pred.detach()
parts = range(n_primitives)
renderables = [
Mesh.from_faces(vertices[0, :, i], faces, colors=get_colors(i))
for i in parts
]
behaviours = [
SceneInit(
scene_init(
load_ground_truth(dataset) if args.mesh else None,
args.up_vector, args.camera_position, args.camera_target,
args.background)),
LightToCamera(),
]
if args.with_rotating_camera:
behaviours += [
CameraTrajectory(Circle(args.camera_target, args.camera_position,
args.up_vector),
speed=1 / 180)
]
show(renderables,
size=args.window_size,
behaviours=behaviours + [SnapshotOnKey()])
print("Saving renderables to file")
for i in range(n_primitives):
m = trimesh.Trimesh(vertices[0, :, i].detach(), faces)
m.export(os.path.join(args.output_directory,
"part_{:03d}.obj".format(i)),
file_type="obj")
def sign_unsigned(fname='./data/elephant.pwn',
K=5,
ndisc=50,
debug_viz=True,
R=15,
cache_sign=False,
use_sign_cache=False):
points, tetvertices, trivertices, is_in_band, epsilon, dist = section1(
fname, K, ndisc)
print('Chosen Epsilon:', epsilon.item())
# coarse sign estimate
print('sign estimate...')
v2v = get_v2v(tetvertices.numpy())
outside_ixes = torch.arange(len(dist))[dist >= epsilon]
is_in_band = dist < epsilon
if use_sign_cache:
print('using sign cache...')
with open('sign_cache.json', 'r') as reader:
guesses = json.load(reader, parse_int=int)
guesses = {int(k): v for k, v in guesses.items()}
else:
guesses = {}
for ix in tqdm(outside_ixes.numpy()):
guesses[ix] = estimate_sign(ix.item(), points.numpy(), v2v,
is_in_band.numpy(), R)
if cache_sign:
print('caching sign...')
guesses = {int(k): [int(v) for v in vec] for k, vec in guesses.items()}
with open('sign_cache.json', 'w') as writer:
json.dump(guesses, writer)
print(guesses)
# cross-sections of the sign plot
alldata = []
for v, changes in guesses.items():
coords = points[v]
vals = [val % 2 for val in changes if val > 0]
if len(vals) > 0:
sign = 1 if np.mean(vals) > 0.5 else 0
alldata.append(
[coords[0], coords[1], coords[2],
np.mean(vals), sign])
alldata = np.array(alldata)
spec = np.sort(np.unique(alldata[:, 2]))
for spei, spe in enumerate(spec):
kept = alldata[:, 2] == spe
fig = plt.figure()
gs = mpl.gridspec.GridSpec(1, 2)
ax = plt.subplot(gs[0, 0])
plt.scatter(alldata[kept, 0],
alldata[kept, 1],
c=alldata[kept, 3],
vmin=0,
vmax=1,
cmap='winter')
ax.set_aspect('equal')
plt.xlim((np.min(alldata[:, 0]), np.max(alldata[:, 0])))
plt.ylim((np.min(alldata[:, 1]), np.max(alldata[:, 1])))
ax = plt.subplot(gs[0, 1])
plt.scatter(alldata[kept, 0],
alldata[kept, 1],
c=alldata[kept, 4],
vmin=0,
vmax=1,
cmap='winter')
ax.set_aspect('equal')
plt.xlim((np.min(alldata[:, 0]), np.max(alldata[:, 0])))
plt.ylim((np.min(alldata[:, 1]), np.max(alldata[:, 1])))
imname = f'out{spei:04}.jpg'
print('saving', imname)
plt.savefig(imname)
plt.close(fig)
# visualize coarse mesh
kept = (dist[trivertices[:, 0]] < epsilon) & (dist[
trivertices[:, 1]] < epsilon) & (dist[trivertices[:, 2]] < epsilon)
m = Mesh.from_faces(points.numpy(),
trivertices.numpy()[kept],
colors=np.ones((len(points), 3)) * [1.0, 0.0, 0.0])
show(m)
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
import numpy as np
from simple_3dviz import Mesh, Lines
from simple_3dviz.window import show
def heart_voxel_grid(N):
"""Create a NxNxN voxel grid with True if the voxel is inside a heart
object and False otherwise."""
x = np.linspace(-1.3, 1.3, N)
y = np.linspace(-1.3, 1.3, N)
z = np.linspace(-1.3, 1.3, N)
x, y, z = np.meshgrid(x, y, z)
return (2 * x**2 + y**2 + z**2 -
1)**3 - (1 / 10) * x**2 * z**3 - y**2 * z**3 < 0
if __name__ == "__main__":
voxels = heart_voxel_grid(64)
m = Mesh.from_voxel_grid(voxels, colors=(0.8, 0, 0))
l = Lines.from_voxel_grid(voxels, colors=(0, 0, 0.), width=0.001)
show([l, m])
import numpy as np
import matplotlib.pyplot as plt
from simple_3dviz.renderables import Spherecloud
from simple_3dviz.behaviours.keyboard import SnapshotOnKey
from simple_3dviz.window import show
if __name__ == "__main__":
t = np.linspace(0, 4 * np.pi, 20)
x = np.sin(2 * t)
y = np.cos(t)
z = np.cos(2 * t)
sizes = (2 + np.sin(t)) * 0.125
centers = np.stack([x, y, z]).reshape(3, -1).T
cmap = plt.cm.copper
colors = cmap(np.random.choice(np.arange(500), centers.shape[0]))
s = Spherecloud(centers=centers, sizes=sizes, colors=colors)
from simple_3dviz import Mesh
m = Mesh.from_file("models/baby_yoda.stl", color=(0.1, 0.8, 0.1))
m.to_unit_cube()
show([s, m], camera_position=(-2.8, -2.8, 0.1), size=(512, 512))
def load_ground_truth(mesh_file):
m = Mesh.from_file(mesh_file, color=(0.8, 0.8, 0.8, 0.3))
# m.to_unit_cube()
return m
def load_ground_truth(dataset):
mesh = dataset.internal_dataset_object[0].groundtruth_mesh
N = len(mesh.vertices)
return Mesh.from_faces(mesh.vertices, mesh.faces, (0.8, 0.8, 0.8, 0.3))