From 64304879ff5a177b2c3e47908e70d61db52cfb0e Mon Sep 17 00:00:00 2001 From: Chandler Date: Mon, 8 May 2023 20:21:09 -0400 Subject: [PATCH] Add OBJ output (#20) --- shap_e/examples/sample_text_to_3d.ipynb | 7 +++++-- shap_e/rendering/mesh.py | 23 +++++++++++++++++++++++ 2 files changed, 28 insertions(+), 2 deletions(-) diff --git a/shap_e/examples/sample_text_to_3d.ipynb b/shap_e/examples/sample_text_to_3d.ipynb index 48615e8..ef39d44 100644 --- a/shap_e/examples/sample_text_to_3d.ipynb +++ b/shap_e/examples/sample_text_to_3d.ipynb @@ -92,8 +92,11 @@ "from shap_e.util.notebooks import decode_latent_mesh\n", "\n", "for i, latent in enumerate(latents):\n", + " t = decode_latent_mesh(xm, latent).tri_mesh()\n", " with open(f'example_mesh_{i}.ply', 'wb') as f:\n", - " decode_latent_mesh(xm, latent).tri_mesh().write_ply(f)" + " t.write_ply(f)\n", + " with open(f'example_mesh_{i}.obj', 'w') as f:\n", + " t.write_obj(f)" ] } ], @@ -113,7 +116,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.9.9" + "version": "3.11.3" } }, "nbformat": 4, diff --git a/shap_e/rendering/mesh.py b/shap_e/rendering/mesh.py index c95b902..7a0c2a1 100644 --- a/shap_e/rendering/mesh.py +++ b/shap_e/rendering/mesh.py @@ -86,3 +86,26 @@ class TriMesh: ), faces=self.faces, ) + + def write_obj(self, raw_f: BinaryIO): + if self.has_vertex_colors(): + vertex_colors = np.stack([self.vertex_channels[x] + for x in "RGB"], axis=1) + vertices = [ + "{} {} {} {} {} {}".format(*coord, *color) + for coord, color in zip(self.verts.tolist(), vertex_colors.tolist()) + ] + else: + vertices = [ + "{} {} {}".format(*coord) + for coord in self.verts.tolist() + ] + + faces = [ + "f {} {} {}".format(str(tri[0] + 1), str(tri[1] + 1), str(tri[2] + 1)) + for tri in self.faces.tolist() + ] + + combined_data = ["v " + vertex for vertex in vertices] + faces + + raw_f.writelines("\n".join(combined_data))