Add glb output

setup.py

@@ -30,6 +30,7 @@
  "scipy",
  "numpy",
  "blobfile",
+ "pygltflib",
  "clip @ git+https://github.com/openai/CLIP.git",
  ],
  author="OpenAI",

shap_e/examples/sample_text_to_3d.ipynb

@@ -96,7 +96,9 @@
  " with open(f'example_mesh_{i}.ply', 'wb') as f:\n",
  " t.write_ply(f)\n",
  " with open(f'example_mesh_{i}.obj', 'w') as f:\n",
- " t.write_obj(f)"
+ " t.write_obj(f)\n",
+ " with open(f'example_mesh_{i}.glb', 'wb') as f:\n",
+ " t.write_glb(f)"
  ]
  }
  ],

shap_e/rendering/glb_util.py

@@ -0,0 +1,124 @@
+import numpy as np
+import pygltflib
+
+from typing import BinaryIO, Optional
+
+from shap_e.util.io import buffered_writer
+
+
+def write_glb(
+ raw_f: BinaryIO,
+ coords: np.ndarray,
+ rgb: Optional[np.ndarray] = None,
+ faces: Optional[np.ndarray] = None,
+):
+ coords = np.asarray(coords, dtype=np.float32)
+ coords_binary_blob = coords.tobytes()
+
+ # setting gltf bufferviews and accessors and mesh primitives
+ primitive = pygltflib.Primitive(
+ attributes=pygltflib.Attributes(POSITION=0),
+ mode=0,
+ )
+ bufferviews = [
+ pygltflib.BufferView(
+ buffer=0,
+ byteLength=len(coords_binary_blob),
+ target=pygltflib.ARRAY_BUFFER,
+ ),
+ ]
+ accessors = [
+ pygltflib.Accessor(
+ bufferView=0,
+ componentType=pygltflib.FLOAT,
+ count=len(coords),
+ type=pygltflib.VEC3,
+ max=coords.max(axis=0).tolist(),
+ min=coords.min(axis=0).tolist(),
+ ),
+ ]
+
+ # adding faces and rgb if exists
+ indx = 1
+ faces_binary_blob = b''
+ if faces is not None:
+ faces = np.asarray(faces, dtype=np.uint32)
+ faces_binary_blob = faces.flatten().tobytes()
+ primitive.indices = indx
+ primitive.mode = 4
+ accessors.append(
+ pygltflib.Accessor(
+ bufferView=indx,
+ componentType=pygltflib.UNSIGNED_INT,
+ count=faces.size,
+ type=pygltflib.SCALAR,
+ max=[int(faces.max())],
+ min=[int(faces.min())],
+ )
+ )
+ bufferviews.append(
+ pygltflib.BufferView(
+ buffer=0,
+ byteOffset=len(coords_binary_blob),
+ byteLength=len(faces_binary_blob),
+ target=pygltflib.ELEMENT_ARRAY_BUFFER,
+ )
+ )
+ indx += 1
+
+ rgb_binary_blob = b''
+ if rgb is not None:
+ rgb = np.asarray(rgb, dtype=np.float32)
+ rgb_binary_blob = rgb.tobytes()
+ primitive.attributes.COLOR_0 = indx
+ accessors.append(
+ pygltflib.Accessor(
+ bufferView=indx,
+ componentType=pygltflib.FLOAT,
+ count=len(rgb),
+ type=pygltflib.VEC3,
+ max=rgb.max(axis=0).tolist(),
+ min=rgb.min(axis=0).tolist(),
+ )
+ )
+ bufferviews.append(
+ pygltflib.BufferView(
+ buffer=0,
+ byteOffset=(
+ len(faces_binary_blob)
+ + len(coords_binary_blob)
+ ),
+ byteLength=len(rgb_binary_blob),
+ target=pygltflib.ARRAY_BUFFER,
+ )
+ )
+
+ gltf = pygltflib.GLTF2(
+ scene=0,
+ scenes=[pygltflib.Scene(nodes=[0])],
+ nodes=[pygltflib.Node(mesh=0)],
+ meshes=[
+ pygltflib.Mesh(
+ primitives=[primitive]
+ )
+ ],
+ accessors=accessors,
+ bufferViews=bufferviews,
+ buffers=[
+ pygltflib.Buffer(
+ byteLength=(
+ len(faces_binary_blob)
+ + len(coords_binary_blob)
+ + len(rgb_binary_blob)
+ )
+ )
+ ],
+ )
+ gltf.set_binary_blob(
+ coords_binary_blob
+ + faces_binary_blob
+ + rgb_binary_blob
+ )
+
+ with buffered_writer(raw_f) as fio:
+ fio.write(b"".join(gltf.save_to_bytes()))

shap_e/rendering/mesh.py

@@ -5,6 +5,7 @@
 import numpy as np
 
 from .ply_util import write_ply
+from .glb_util import write_glb
 
 
 @dataclass
@@ -105,3 +106,15 @@ def write_obj(self, raw_f: BinaryIO):
  combined_data = ["v " + vertex for vertex in vertices] + faces
 
  raw_f.writelines("\n".join(combined_data))
+
+ def write_glb(self, raw_f: BinaryIO):
+ write_glb(
+ raw_f,
+ coords=self.verts,
+ rgb=(
+ np.stack([self.vertex_channels[x] for x in "RGB"], axis=1)
+ if self.has_vertex_colors()
+ else None
+ ),
+ faces=self.faces,
+ )