import matplotlib matplotlib.use("agg") import os from pathlib import Path import matplotlib.pyplot as plt import numpy as np import trimesh from mpl_toolkits.mplot3d import Axes3D """ Custom visualization """ def export_to_pc_batch(dir, pcs, colors=None): Path(dir).mkdir(parents=True, exist_ok=True) for i, xyz in enumerate(pcs): if colors is None: color = None else: color = colors[i] pcwrite(os.path.join(dir, "sample_" + str(i) + ".ply"), xyz, color) def export_to_obj(dir, meshes, transform=lambda v, f: (v, f)): """ transform: f(vertices, faces) --> transformed (vertices, faces) """ Path(dir).mkdir(parents=True, exist_ok=True) for i, data in enumerate(meshes): v, f = transform(data[0], data[1]) if len(data) > 2: v_color = data[2] else: v_color = None mesh = trimesh.Trimesh(v, f, vertex_colors=v_color) out = trimesh.exchange.obj.export_obj(mesh) with open(os.path.join(dir, "sample_" + str(i) + ".obj"), "w") as f: f.write(out) f.close() def export_to_obj_single(path, data, transform=lambda v, f: (v, f)): """ transform: f(vertices, faces) --> transformed (vertices, faces) """ v, f = transform(data[0], data[1]) if len(data) > 2: v_color = data[2] else: v_color = None mesh = trimesh.Trimesh(v, f, vertex_colors=v_color) out = trimesh.exchange.obj.export_obj(mesh) with open(path, "w") as f: f.write(out) f.close() def meshwrite(filename, verts, faces, norms, colors): """Save a 3D mesh to a polygon .ply file.""" # Write header ply_file = open(filename, "w") ply_file.write("ply\n") ply_file.write("format ascii 1.0\n") ply_file.write("element vertex %d\n" % (verts.shape[0])) ply_file.write("property float x\n") ply_file.write("property float y\n") ply_file.write("property float z\n") ply_file.write("property float nx\n") ply_file.write("property float ny\n") ply_file.write("property float nz\n") ply_file.write("property uchar red\n") ply_file.write("property uchar green\n") ply_file.write("property uchar blue\n") ply_file.write("element face %d\n" % (faces.shape[0])) ply_file.write("property list uchar int vertex_index\n") ply_file.write("end_header\n") # Write vertex list for i in range(verts.shape[0]): ply_file.write( "%f %f %f %f %f %f %d %d %d\n" % ( verts[i, 0], verts[i, 1], verts[i, 2], norms[i, 0], norms[i, 1], norms[i, 2], colors[i, 0], colors[i, 1], colors[i, 2], ) ) # Write face list for i in range(faces.shape[0]): ply_file.write("3 %d %d %d\n" % (faces[i, 0], faces[i, 1], faces[i, 2])) ply_file.close() def pcwrite(filename, xyz, rgb=None): """Save a point cloud to a polygon .ply file.""" if rgb is None: rgb = np.ones_like(xyz) * 128 rgb = rgb.astype(np.uint8) # Write header ply_file = open(filename, "w") ply_file.write("ply\n") ply_file.write("format ascii 1.0\n") ply_file.write("element vertex %d\n" % (xyz.shape[0])) ply_file.write("property float x\n") ply_file.write("property float y\n") ply_file.write("property float z\n") ply_file.write("property uchar red\n") ply_file.write("property uchar green\n") ply_file.write("property uchar blue\n") ply_file.write("end_header\n") # Write vertex list for i in range(xyz.shape[0]): ply_file.write( "%f %f %f %d %d %d\n" % ( xyz[i, 0], xyz[i, 1], xyz[i, 2], rgb[i, 0], rgb[i, 1], rgb[i, 2], ) ) """ Matplotlib Visualization """ def visualize_voxels(out_file, voxels, num_shown=16, threshold=0.5): r"""Visualizes voxel data. show only first num_shown """ batch_size = voxels.shape[0] voxels = voxels.squeeze(1) > threshold num_shown = min(num_shown, batch_size) n = int(np.sqrt(num_shown)) fig = plt.figure(figsize=(20, 20)) for idx, pc in enumerate(voxels[:num_shown]): if idx >= n * n: break pc = voxels[idx] ax = fig.add_subplot(n, n, idx + 1, projection="3d") ax.voxels(pc, edgecolor="k", facecolors="green", linewidth=0.1, alpha=0.5) ax.view_init() ax.axis("off") plt.savefig(out_file, bbox_inches="tight") plt.close() def visualize_pointcloud(points, normals=None, out_file=None, show=False, elev=30, azim=225): r"""Visualizes point cloud data. Args: points (tensor): point data normals (tensor): normal data (if existing) out_file (string): output file show (bool): whether the plot should be shown """ # Create plot fig = plt.figure() ax = fig.gca(projection=Axes3D.name) ax.scatter(points[:, 2], points[:, 0], points[:, 1]) if normals is not None: ax.quiver( points[:, 2], points[:, 0], points[:, 1], normals[:, 2], normals[:, 0], normals[:, 1], length=0.1, color="k" ) ax.set_xlabel("Z") ax.set_ylabel("X") ax.set_zlabel("Y") # ax.set_xlim(-0.5, 0.5) # ax.set_ylim(-0.5, 0.5) # ax.set_zlim(-0.5, 0.5) ax.view_init(elev=elev, azim=azim) if out_file is not None: plt.savefig(out_file) if show: plt.show() plt.close(fig) def visualize_pointcloud_batch( path, pointclouds, pred_labels, labels, categories, vis_label=False, target=None, elev=30, azim=225 ): batch_size = len(pointclouds) fig = plt.figure(figsize=(20, 20)) ncols = int(np.sqrt(batch_size)) nrows = max(1, (batch_size - 1) // ncols + 1) for idx, pc in enumerate(pointclouds): if vis_label: label = categories[labels[idx].item()] pred = categories[pred_labels[idx]] colour = "g" if label == pred else "r" elif target is None: colour = "g" else: colour = target[idx] pc = pc.cpu().numpy() ax = fig.add_subplot(nrows, ncols, idx + 1, projection="3d") ax.scatter(pc[:, 0], pc[:, 2], pc[:, 1], c=colour, s=5) ax.view_init(elev=elev, azim=azim) ax.axis("off") if vis_label: ax.set_title("GT: {0}\nPred: {1}".format(label, pred)) plt.savefig(path) plt.close(fig) """ Plot stats """ def plot_stats(output_dir, stats, interval): content = stats.keys() # f = plt.figure(figsize=(20, len(content) * 5)) f, axs = plt.subplots(len(content), 1, figsize=(20, len(content) * 5)) for j, (k, v) in enumerate(stats.items()): axs[j].plot(interval, v) axs[j].set_ylabel(k) f.savefig(os.path.join(output_dir, "stat.pdf"), bbox_inches="tight") plt.close(f)