def create_tsne_plot(codes, voxels=None, labels=None, filename="plot.pdf", indices=None): from sklearn.manifold import TSNE from matplotlib.offsetbox import OffsetImage, AnnotationBbox width, height = 40, 52 print("Calculating t-sne embedding...") tsne = TSNE(n_components=2) embedded = tsne.fit_transform(codes) print("Plotting...") fig, ax = plt.subplots() plt.axis('off') margin = 0.0128 plt.margins(margin * height / width, margin) x = embedded[:, 0] y = embedded[:, 1] x = np.interp(x, (x.min(), x.max()), (0, 1)) y = np.interp(y, (y.min(), y.max()), (0, 1)) ax.scatter(x, y, c=labels, s=40, cmap='Set1') fig.set_size_inches(width, height) if voxels is not None: print("Creating images...") from rendering import MeshRenderer viewer = MeshRenderer(start_thread=False) for i in tqdm(range(voxels.shape[0])): viewer.set_voxels(voxels[i, :, :, :].cpu().numpy()) viewer.model_color = dataset.get_color(labels[i]) image = viewer.get_image(crop=True, output_size=128) box = AnnotationBbox(OffsetImage(image, zoom=0.5, cmap='gray'), (x[i], y[i]), frameon=True) ax.add_artist(box) if indices is not None: print("Creating images...") dataset_directories = open('data/models.txt', 'r').readlines() from rendering import MeshRenderer viewer = MeshRenderer(start_thread=False) import trimesh import logging logging.getLogger('trimesh').setLevel(1000000) for i in tqdm(range(len(indices))): 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(labels[i]) image = viewer.get_image(crop=True, output_size=128) box = AnnotationBbox(OffsetImage(image, zoom=0.5, cmap='gray'), (x[i], y[i]), frameon=True) ax.add_artist(box) print("Saving PDF...") extent = ax.get_window_extent().transformed(fig.dpi_scale_trans.inverted()) plt.savefig(filename, bbox_inches=extent, dpi=200)
for label in range(COUNT): objects = (dataset.labels == label).nonzero() indices.append(objects[random.randint(0, objects.shape[0] - 1)].item()) latent_codes = latent_codes[indices, :] plot = ImageGrid(COUNT, 2, create_viewer=False) dataset_directories = directories = open('data/models.txt', 'r').readlines() 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)