class ImageGrid():
def __init__(self,
width,
height=1,
cell_width=3,
cell_height=None,
margin=0.2,
create_viewer=True,
crop=True):
print("Plotting...")
self.width = width
self.height = height
cell_height = cell_height if cell_height is not None else cell_width
self.figure, self.axes = plt.subplots(height,
width,
figsize=(width * cell_width,
height * cell_height),
gridspec_kw={
'left': 0,
'right': 1,
'top': 1,
'bottom': 0,
'wspace': margin,
'hspace': margin
})
self.figure.patch.set_visible(False)
self.crop = crop
if create_viewer:
from rendering import MeshRenderer
self.viewer = MeshRenderer(start_thread=False)
else:
self.viewer = None
def set_image(self, image, x=0, y=0):
cell = self.axes[
y, x] if self.height > 1 and self.width > 1 else self.axes[x + y]
cell.imshow(image)
cell.axis('off')
cell.patch.set_visible(False)
def set_voxels(self, voxels, x=0, y=0, color=None):
if color is not None:
self.viewer.model_color = color
self.viewer.set_voxels(voxels)
image = self.viewer.get_image(crop=self.crop)
self.set_image(image, x, y)
def save(self, filename):
plt.axis('off')
extent = self.figure.get_window_extent().transformed(
self.figure.dpi_scale_trans.inverted())
plt.savefig(filename, bbox_inches=extent, dpi=400)
if self.viewer is not None:
self.viewer.delete_buffers()
for i in range(COUNT):
mesh = trimesh.load(
os.path.join(dataset_directories[index].strip(),
'model_normalized.obj'))
viewer.set_mesh(mesh, center_and_scale=True)
viewer.model_color = dataset.get_color(i)
image = viewer.get_image(crop=True)
plot.set_image(image, i, 0)
image = render_image(sdf_net,
latent_codes[i, :],
color=dataset.get_color(i),
crop=True)
plot.set_image(image, i, 1)
viewer.delete_buffers()
plot.save("plots/deepsdf-reconstruction-classes.pdf")
if "autoencoder" in sys.argv:
from dataset import dataset as dataset
dataset.load_voxels(device)
dataset.load_labels(device)
indices = random.sample(list(range(dataset.size)), 1000)
voxels = dataset.voxels[indices, :, :, :]
autoencoder = load_autoencoder(is_variational='clasic' not in sys.argv)
print("Generating codes...")
with torch.no_grad():
codes = autoencoder.encode(voxels).cpu().numpy()
create_tsne_plot(
codes, voxels, dataset.labels[indices].cpu().numpy(),