diff --git a/kaolin/rep/Mesh.py b/kaolin/rep/Mesh.py index 3b1d22526..d1e92ca17 100644 --- a/kaolin/rep/Mesh.py +++ b/kaolin/rep/Mesh.py @@ -944,6 +944,24 @@ def show(self): kal.visualize.show_mesh(self) + def save_image(self, filename='mesh_image', resolution=(1080, 1080)): + r""" Saves the mesh image in png format. The extension is added automatically. + + Args: + filename: the file name to save the file under + resulution: The resolution of the image to be saved + + Example: + >>> mesh = Mesh.from_obj('model.obj') + >>> mesh.save_image('mesh_image') + + """ + + png = kal.visualize.save_mesh_image(self, filename=filename, + resolution=resolution) + + return png + def save_tensors(self, filename: (str)): r"""Saves the tensor information of the mesh in a numpy .npz format. diff --git a/kaolin/visualize/vis.py b/kaolin/visualize/vis.py index d89dec42c..123b2b9d2 100644 --- a/kaolin/visualize/vis.py +++ b/kaolin/visualize/vis.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import os from typing import Union import torch @@ -86,6 +87,32 @@ def show_mesh(input_mesh: kaolin.rep.Mesh, colors: list = [.7, .2, .2]): mesh.show() +def save_mesh_image(input_mesh: kaolin.rep.Mesh, + colors: list = [.7, .2, .2], filename='mesh_image', + resolution=(1080, 1080)): + r""" Saver for meshes + + Args: + verts (torch.Tensor): vertices of mesh to be visualized + faces (torch.Tensor): faces of mesh to be visualized + colors (list): rbg colour values for rendered mesh + filename: the file name to save the file under + resulution: The resolution of the image to be saved + """ + + mesh = trimesh.Trimesh(vertices=input_mesh.vertices.data.cpu().numpy(), + faces=input_mesh.faces.data.cpu().numpy()) + mesh.visual.vertex_colors = colors + png = mesh.scene().save_image(resolution=resolution, visible=True) + path, _ = os.path.splitext(filename) + filename = path + '.png' + with open(filename, 'wb') as f: + f.write(png) + f.close() + + return png + + def show_sdf(sdf: kaolin.rep.SDF, mode='mesh', bbox_center: float = 0., bbox_dim: float = 1., num_points: int = 100000, colors=[.7, .2, .2]):