242 lines
6.9 KiB
Python
242 lines
6.9 KiB
Python
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)
|