Refactor gen_image function in app.py to return base64 encoded GLB data instead of a file path. Update mesh export in inference.py to return .obj file extension. Modify Mesh class to streamline GLB export and enhance texture handling, ensuring proper integration of vertex attributes.

app.py

@@ -96,6 +96,7 @@ def preprocess_image(image, background_choice, foreground_ratio, backgroud_color
 
 
 @spaces.GPU
+
 def gen_image(input_image, seed, scale, step):
     global pipeline, model, args
     pipeline.set_seed(seed)
@@ -107,19 +108,13 @@ def gen_image(input_image, seed, scale, step):
     
     glb_path = generate3d(model, np_imgs, np_xyzs, args.device)
     
-    # Create a temporary file with a proper name for the GLB data
-    import tempfile
-    import shutil
-    
-    # Create a temporary file with a proper extension
-    temp_glb = tempfile.NamedTemporaryFile(suffix='.glb', delete=False)
-    temp_glb.close()
-    
-    # Copy the generated GLB file to our temporary file
-    shutil.copy2(glb_path, temp_glb.name)
+    # Read the GLB file and encode it in base64
+    with open(glb_path, 'rb') as f:
+        glb_bytes = f.read()
+    encoded_glb = 'data:model/gltf-binary;base64,' + base64.b64encode(glb_bytes).decode('utf-8')
     
-    # Return images and the path to the temporary GLB file
-    return Image.fromarray(np_imgs), Image.fromarray(np_xyzs), temp_glb.name
+    # Return images and the encoded GLB data
+    return Image.fromarray(np_imgs), Image.fromarray(np_xyzs), encoded_glb
 
 
 parser = argparse.ArgumentParser()
@@ -220,7 +215,7 @@ with gr.Blocks() as demo:
             xyz_ouput = gr.Image(interactive=False, label="Output CCM image")
 
             output_model = gr.Model3D(
-                label="Output 3D Model (GLB)",
+                label="Output OBJ",
                 interactive=False,
             )
             gr.Markdown("Note: Ensure that the input image is correctly pre-processed into a grey background, otherwise the results will be unpredictable.")

inference.py

@@ -73,7 +73,7 @@ def generate3d(model, rgb, ccm, device):
 
     start_time = time.time()
     with torch.no_grad():
-        mesh_path_glb = tempfile.NamedTemporaryFile(suffix=f".glb", delete=False).name
+        mesh_path_glb = tempfile.NamedTemporaryFile(suffix=f"", delete=False).name
         model.export_mesh(data_config, mesh_path_glb, tri_fea_2 = triplane_feature2)
 
         # glctx = dr.RasterizeGLContext()#dr.RasterizeCudaContext()
@@ -96,4 +96,4 @@ def generate3d(model, rgb, ccm, device):
     end_time = time.time()
     elapsed_time = end_time - start_time
     print(f"uv takes {elapsed_time}s")
-    return mesh_path_glb
+    return mesh_path_glb+".obj"

mesh.py

@@ -10,7 +10,6 @@ from kiui.typing import *
 class Mesh:
     """
     A torch-native trimesh class, with support for ``ply/obj/glb`` formats.
-
     Note:
         This class only supports one mesh with a single texture image (an albedo texture and a metallic-roughness texture).
     """
@@ -28,7 +27,6 @@ class Mesh:
         device: Optional[torch.device] = None,
     ):
         """Init a mesh directly using all attributes.
-
         Args:
             v (Optional[Tensor]): vertices, float [N, 3]. Defaults to None.
             f (Optional[Tensor]): faces, int [M, 3]. Defaults to None.
@@ -62,7 +60,6 @@ class Mesh:
     @classmethod
     def load(cls, path, resize=True, clean=False, renormal=True, retex=False, bound=0.9, front_dir='+z', **kwargs):
         """load mesh from path.
-
         Args:
             path (str): path to mesh file, supports ply, obj, glb.
             clean (bool, optional): perform mesh cleaning at load (e.g., merge close vertices). Defaults to False.
@@ -76,7 +73,6 @@ class Mesh:
         Note:
             a ``device`` keyword argument can be provided to specify the torch device. 
             If it's not provided, we will try to use ``'cuda'`` as the device if it's available.
-
         Returns:
             Mesh: the loaded Mesh object.
         """
@@ -140,7 +136,6 @@ class Mesh:
     @classmethod
     def load_obj(cls, path, albedo_path=None, device=None):
         """load an ``obj`` mesh.
-
         Args:
             path (str): path to mesh.
             albedo_path (str, optional): path to the albedo texture image, will overwrite the existing texture path if specified in mtl. Defaults to None.
@@ -149,7 +144,6 @@ class Mesh:
         Note: 
             We will try to read `mtl` path from `obj`, else we assume the file name is the same as `obj` but with `mtl` extension.
             The `usemtl` statement is ignored, and we only use the last material path in `mtl` file.
-
         Returns:
             Mesh: the loaded Mesh object.
         """
@@ -313,17 +307,13 @@ class Mesh:
     @classmethod
     def load_trimesh(cls, path, device=None):
         """load a mesh using ``trimesh.load()``.
-
         Can load various formats like ``glb`` and serves as a fallback.
-
         Note:
             We will try to merge all meshes if the glb contains more than one, 
             but **this may cause the texture to lose**, since we only support one texture image!
-
         Args:
             path (str): path to the mesh file.
             device (torch.device, optional): torch device. Defaults to None.
-
         Returns:
             Mesh: the loaded Mesh object.
         """
@@ -423,10 +413,8 @@ class Mesh:
     # sample surface (using trimesh)
     def sample_surface(self, count: int):
         """sample points on the surface of the mesh.
-
         Args:
             count (int): number of points to sample.
-
         Returns:
             torch.Tensor: the sampled points, float [count, 3].
         """
@@ -438,7 +426,6 @@ class Mesh:
     # aabb
     def aabb(self):
         """get the axis-aligned bounding box of the mesh.
-
         Returns:
             Tuple[torch.Tensor]: the min xyz and max xyz of the mesh.
         """
@@ -448,7 +435,6 @@ class Mesh:
     @torch.no_grad()
     def auto_size(self, bound=0.9):
         """auto resize the mesh.
-
         Args:
             bound (float, optional): resizing into ``[-bound, bound]^3``. Defaults to 0.9.
         """
@@ -484,7 +470,6 @@ class Mesh:
 
     def auto_uv(self, cache_path=None, vmap=True):
         """auto calculate the uv coordinates.
-
         Args:
             cache_path (str, optional): path to save/load the uv cache as a npz file, this can avoid calculating uv every time when loading the same mesh, which is time-consuming. Defaults to None.
             vmap (bool, optional): remap vertices based on uv coordinates, so each v correspond to a unique vt (necessary for formats like gltf). 
@@ -523,7 +508,6 @@ class Mesh:
     
     def align_v_to_vt(self, vmapping=None):
         """ remap v/f and vn/fn to vt/ft.
-
         Args:
             vmapping (np.ndarray, optional): the mapping relationship from f to ft. Defaults to None.
         """
@@ -542,10 +526,8 @@ class Mesh:
 
     def to(self, device):
         """move all tensor attributes to device.
-
         Args:
             device (torch.device): target device.
-
         Returns:
             Mesh: self.
         """
@@ -558,7 +540,6 @@ class Mesh:
     
     def write(self, path):
         """write the mesh to a path.
-
         Args:
             path (str): path to write, supports ply, obj and glb.
         """
@@ -573,7 +554,6 @@ class Mesh:
     
     def write_ply(self, path):
         """write the mesh in ply format. Only for geometry!
-
         Args:
             path (str): path to write.
         """
@@ -591,16 +571,16 @@ class Mesh:
     def write_glb(self, path):
         """write the mesh in glb/gltf format.
           This will create a scene with a single mesh.
-
         Args:
             path (str): path to write.
         """
-        import pygltflib
 
         # assert self.v.shape[0] == self.vn.shape[0] and self.v.shape[0] == self.vt.shape[0]
         if self.vt is not None and self.v.shape[0] != self.vt.shape[0]:
             self.align_v_to_vt()
 
+        import pygltflib
+
         f_np = self.f.detach().cpu().numpy().astype(np.uint32)
         f_np_blob = f_np.flatten().tobytes()
 
@@ -610,104 +590,79 @@ class Mesh:
         blob = f_np_blob + v_np_blob
         byteOffset = len(blob)
 
-        # Create attributes dictionary
-        attributes = pygltflib.Attributes(POSITION=1)
-        accessor_count = 2  # Start after position (0) and indices (1)
-
-        # Initialize buffer views list
-        buffer_views = [
-            # triangles; as flatten (element) array
-            pygltflib.BufferView(
-                buffer=0,
-                byteLength=len(f_np_blob),
-                target=pygltflib.ELEMENT_ARRAY_BUFFER,
-            ),
-            # positions; as vec3 array
-            pygltflib.BufferView(
-                buffer=0,
-                byteOffset=len(f_np_blob),
-                byteLength=len(v_np_blob),
-                byteStride=12,  # vec3
-                target=pygltflib.ARRAY_BUFFER,
-            ),
-        ]
-
-        # Initialize accessors list
-        accessors = [
-            # 0 = triangles
-            pygltflib.Accessor(
-                bufferView=0,
-                componentType=pygltflib.UNSIGNED_INT,
-                count=f_np.size,
-                type=pygltflib.SCALAR,
-                max=[int(f_np.max())],
-                min=[int(f_np.min())],
-            ),
-            # 1 = positions
-            pygltflib.Accessor(
-                bufferView=1,
-                componentType=pygltflib.FLOAT,
-                count=len(v_np),
-                type=pygltflib.VEC3,
-                max=v_np.max(axis=0).tolist(),
-                min=v_np.min(axis=0).tolist(),
-            ),
-        ]
-
-        # Add vertex colors if they exist
-        if self.vc is not None:
-            vc_np = self.vc.detach().cpu().numpy().astype(np.float32)
-            vc_np_blob = vc_np.tobytes()
-            
-            # Add vertex color buffer view
-            buffer_views.append(
-                pygltflib.BufferView(
-                    buffer=0,
-                    byteOffset=byteOffset,
-                    byteLength=len(vc_np_blob),
-                    byteStride=12,  # vec3
-                    target=pygltflib.ARRAY_BUFFER,
-                )
-            )
-            
-            # Add vertex color accessor
-            accessors.append(
-                pygltflib.Accessor(
-                    bufferView=accessor_count,
-                    componentType=pygltflib.FLOAT,
-                    count=len(vc_np),
-                    type=pygltflib.VEC3,
-                    max=vc_np.max(axis=0).tolist(),
-                    min=vc_np.min(axis=0).tolist(),
-                )
-            )
-            
-            # Add to attributes
-            attributes.COLOR_0 = accessor_count
-            accessor_count += 1
-            
-            blob += vc_np_blob
-            byteOffset += len(vc_np_blob)
-
-        # Create the GLTF object with all components
+        # base mesh
         gltf = pygltflib.GLTF2(
             scene=0,
             scenes=[pygltflib.Scene(nodes=[0])],
             nodes=[pygltflib.Node(mesh=0)],
             meshes=[pygltflib.Mesh(primitives=[pygltflib.Primitive(
-                attributes=attributes,
+                # indices to accessors (0 is triangles)
+                attributes=pygltflib.Attributes(
+                    POSITION=1,
+                ),
                 indices=0,
             )])],
-            buffers=[pygltflib.Buffer(byteLength=byteOffset)],
-            bufferViews=buffer_views,
-            accessors=accessors,
+            buffers=[
+                pygltflib.Buffer(byteLength=len(f_np_blob) + len(v_np_blob))
+            ],
+            # buffer view (based on dtype)
+            bufferViews=[
+                # triangles; as flatten (element) array
+                pygltflib.BufferView(
+                    buffer=0,
+                    byteLength=len(f_np_blob),
+                    target=pygltflib.ELEMENT_ARRAY_BUFFER, # GL_ELEMENT_ARRAY_BUFFER (34963)
+                ),
+                # positions; as vec3 array
+                pygltflib.BufferView(
+                    buffer=0,
+                    byteOffset=len(f_np_blob),
+                    byteLength=len(v_np_blob),
+                    byteStride=12, # vec3
+                    target=pygltflib.ARRAY_BUFFER, # GL_ARRAY_BUFFER (34962)
+                ),
+            ],
+            accessors=[
+                # 0 = triangles
+                pygltflib.Accessor(
+                    bufferView=0,
+                    componentType=pygltflib.UNSIGNED_INT, # GL_UNSIGNED_INT (5125)
+                    count=f_np.size,
+                    type=pygltflib.SCALAR,
+                    max=[int(f_np.max())],
+                    min=[int(f_np.min())],
+                ),
+                # 1 = positions
+                pygltflib.Accessor(
+                    bufferView=1,
+                    componentType=pygltflib.FLOAT, # GL_FLOAT (5126)
+                    count=len(v_np),
+                    type=pygltflib.VEC3,
+                    max=v_np.max(axis=0).tolist(),
+                    min=v_np.min(axis=0).tolist(),
+                ),
+            ],
         )
 
-        # Add material for vertex colors
-        if self.vc is not None:
+        # append texture info
+        if self.vt is not None:
+
+            vt_np = self.vt.detach().cpu().numpy().astype(np.float32)
+            vt_np_blob = vt_np.tobytes()
+
+            albedo = self.albedo.detach().cpu().numpy()
+            albedo = (albedo * 255).astype(np.uint8)
+            albedo = cv2.cvtColor(albedo, cv2.COLOR_RGB2BGR)
+            albedo_blob = cv2.imencode('.png', albedo)[1].tobytes()
+
+            # update primitive
+            gltf.meshes[0].primitives[0].attributes.TEXCOORD_0 = 2
+            gltf.meshes[0].primitives[0].material = 0
+
+            # update materials
             gltf.materials.append(pygltflib.Material(
                 pbrMetallicRoughness=pygltflib.PbrMetallicRoughness(
-                    baseColorFactor=[1.0, 1.0, 1.0, 1.0],
+                    baseColorTexture=pygltflib.TextureInfo(index=0, texCoord=0),
                     metallicFactor=0.0,
                     roughnessFactor=1.0,
                 ),
@@ -715,28 +670,27 @@ class Mesh:
                 alphaCutoff=None,
                 doubleSided=True,
             ))
-            gltf.meshes[0].primitives[0].material = 0
 
-        # Handle textures if they exist
-        if self.vt is not None:
-            vt_np = self.vt.detach().cpu().numpy().astype(np.float32)
-            vt_np_blob = vt_np.tobytes()
+            gltf.textures.append(pygltflib.Texture(sampler=0, source=0))
+            gltf.samplers.append(pygltflib.Sampler(magFilter=pygltflib.LINEAR, minFilter=pygltflib.LINEAR_MIPMAP_LINEAR, wrapS=pygltflib.REPEAT, wrapT=pygltflib.REPEAT))
+            gltf.images.append(pygltflib.Image(bufferView=3, mimeType="image/png"))
 
-            # Add texture coordinates buffer view
+            # update buffers
             gltf.bufferViews.append(
+                # index = 2, texcoords; as vec2 array
                 pygltflib.BufferView(
                     buffer=0,
                     byteOffset=byteOffset,
                     byteLength=len(vt_np_blob),
-                    byteStride=8,  # vec2
+                    byteStride=8, # vec2
                     target=pygltflib.ARRAY_BUFFER,
                 )
             )
 
-            # Add texture coordinates accessor
             gltf.accessors.append(
+                # 2 = texcoords
                 pygltflib.Accessor(
-                    bufferView=len(gltf.bufferViews) - 1,
+                    bufferView=2,
                     componentType=pygltflib.FLOAT,
                     count=len(vt_np),
                     type=pygltflib.VEC2,
@@ -745,25 +699,64 @@ class Mesh:
                 )
             )
 
-            # Add texture coordinates to attributes
-            gltf.meshes[0].primitives[0].attributes.TEXCOORD_0 = len(gltf.accessors) - 1
-
-            blob += vt_np_blob
+            blob += vt_np_blob 
             byteOffset += len(vt_np_blob)
 
-            # Update buffer size
+            gltf.bufferViews.append(
+                # index = 3, albedo texture; as none target
+                pygltflib.BufferView(
+                    buffer=0,
+                    byteOffset=byteOffset,
+                    byteLength=len(albedo_blob),
+                )
+            )
+
+            blob += albedo_blob
+            byteOffset += len(albedo_blob)
+
             gltf.buffers[0].byteLength = byteOffset
 
-        # Set the binary blob
+            # append metllic roughness
+            if self.metallicRoughness is not None:
+                metallicRoughness = self.metallicRoughness.detach().cpu().numpy()
+                metallicRoughness = (metallicRoughness * 255).astype(np.uint8)
+                metallicRoughness = cv2.cvtColor(metallicRoughness, cv2.COLOR_RGB2BGR)
+                metallicRoughness_blob = cv2.imencode('.png', metallicRoughness)[1].tobytes()
+
+                # update texture definition
+                gltf.materials[0].pbrMetallicRoughness.metallicFactor = 1.0
+                gltf.materials[0].pbrMetallicRoughness.roughnessFactor = 1.0
+                gltf.materials[0].pbrMetallicRoughness.metallicRoughnessTexture = pygltflib.TextureInfo(index=1, texCoord=0)
+
+                gltf.textures.append(pygltflib.Texture(sampler=1, source=1))
+                gltf.samplers.append(pygltflib.Sampler(magFilter=pygltflib.LINEAR, minFilter=pygltflib.LINEAR_MIPMAP_LINEAR, wrapS=pygltflib.REPEAT, wrapT=pygltflib.REPEAT))
+                gltf.images.append(pygltflib.Image(bufferView=4, mimeType="image/png"))
+
+                # update buffers
+                gltf.bufferViews.append(
+                    # index = 4, metallicRoughness texture; as none target
+                    pygltflib.BufferView(
+                        buffer=0,
+                        byteOffset=byteOffset,
+                        byteLength=len(metallicRoughness_blob),
+                    )
+                )
+
+                blob += metallicRoughness_blob
+                byteOffset += len(metallicRoughness_blob)
+
+                gltf.buffers[0].byteLength = byteOffset
+
+            
+        # set actual data
         gltf.set_binary_blob(blob)
 
-        # Save the GLB file
+        # glb = b"".join(gltf.save_to_bytes())
         gltf.save(path)
 
 
     def write_obj(self, path):
         """write the mesh in obj format. Will also write the texture and mtl files.
-
         Args:
             path (str): path to write.
         """
@@ -826,5 +819,4 @@ class Mesh:
             metallicRoughness = self.metallicRoughness.detach().cpu().numpy()
             metallicRoughness = (metallicRoughness * 255).astype(np.uint8)
             cv2.imwrite(metallic_path, metallicRoughness[..., 2])
-            cv2.imwrite(roughness_path, metallicRoughness[..., 1])
-
+            cv2.imwrite(roughness_path, metallicRoughness[..., 1])

model/crm/model.py

@@ -98,28 +98,17 @@ class CRM(nn.Module):
         # Expect predicted colors value range from [-1, 1]
         colors = (colors * 0.5 + 0.5).clip(0, 1)
 
-        # Transform vertices to match GLB coordinate system
-        # GLB uses right-handed coordinate system with Y up
-        verts = verts[..., [0, 2, 1]]  # Swap Y and Z to get Y up
-        verts[..., 0] *= -1  # Flip X to get right-handed
+        verts = verts[..., [0, 2, 1]]
+        verts[..., 0]*= -1
+        verts[..., 2]*= -1
         verts = verts.squeeze().cpu().numpy()
-        
-        # Transform faces to maintain correct winding order
-        faces = faces[..., [2, 1, 0]]  # Reverse winding order
-        faces = faces.squeeze().cpu().numpy()
+        faces = faces[..., [2, 1, 0]][..., [0, 2, 1]]#[..., [1, 0, 2]]
+        faces = faces.squeeze().cpu().numpy()#faces[..., [2, 1, 0]].squeeze().cpu().numpy()
 
         # export the final mesh
         with torch.no_grad():
-            # Create a Mesh object with the data
-            from mesh import Mesh
-            mesh = Mesh(
-                v=torch.from_numpy(verts).float(),
-                f=torch.from_numpy(faces).int(),
-                vc=torch.from_numpy(colors).float(),
-                device='cpu'
-            )
-            # Write as GLB
-            mesh.write(out_dir)
+            mesh = trimesh.Trimesh(verts, faces, vertex_colors=colors, process=False) # important, process=True leads to seg fault...
+            mesh.export(f'{out_dir}.obj')
 
     def export_mesh_wt_uv(self, ctx, data, out_dir, ind, device, res, tri_fea_2=None):