Add OBJ output (#20)

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,

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))