Load fixed custom mesh and train texture only.

DOCUMENTATION.md

@@ -203,7 +203,8 @@ An explicit geometry parameterized with a feature volume. The feature volume has
 
 | name | type | description |
 | ---- | ---- | ----------- |
-
+| pos_encoding_config                  | dict              | Configurations for the positional encoding. See https://github.com/NVlabs/tiny-cuda-nn/blob/master/DOCUMENTATION.md#encodings for supported arguments. Default: {}                                                                                                               |
+| mlp_network_config                   | dict              | Configurations for the MLP head for feature prediction. See https://github.com/NVlabs/tiny-cuda-nn/blob/master/DOCUMENTATION.md#networks for supported arguments. Default: {}                                                                                                                                                                                                                                                     |
 ### tetrahedra-sdf-grid
 
 | name                                 | type              | description                                                                                                                                                                                                                                                                      |
@@ -212,14 +213,14 @@ An explicit geometry parameterized with a feature volume. The feature volume has
 | isosurface_deformable_grid           | bool              | Whether to optimize positions of tetrahedra grid vertices for surface extraction. Default: True                                                                                                                                                                                  |
 | isosurface_remove_outliers           | bool              | Whether to remove outlier components according to the number of faces. Only remove if the isosurface process does not require gradient. Default: False                                                                                                                           |
 | isosurface_outlier_n_faces_threshold | Union[int, float] | Extracted mesh components with number of faces less than this threshold will be removed if `isosurface_remove_outliers=True`. If `int`, direcly used as the threshold number of faces; if `float`, used as the ratio of all face numbers to compute the threshold. Default: 0.01 |
-| pos_encoding_config                  | dict              | Configurations for the positional encoding. See https://github.com/NVlabs/tiny-cuda-nn/blob/master/DOCUMENTATION.md#encodings for supported arguments. Default: {}                                                                                                               |
-| mlp_network_config                   | dict              | Configurations for the MLP head for feature prediction. See https://github.com/NVlabs/tiny-cuda-nn/blob/master/DOCUMENTATION.md#networks for supported arguments. Default: {}                                                                                                    |
-| shape_init                           | Optional[str]     | The shape to initializa the SDF as, in [None, "sphere", "ellipsoid"]. If None, does not initialize; if "sphere", initialized as a sphere; if "ellipsoid", initialized as an ellipsoid. Default: None                                                                             |
-| shape_init_params                    | Optional[Any]     | Parameters to specify the SDF initialization. If `shape_init="sphere"`, a float is used for the sphere radius; if `shape_init="ellipsoid"`, a tuple of three floats is used for the radius along x/y/z axis. Default: None                                                       |
-| force_shape_init                     | bool              | Whether to force initialization of the SDf even if weights are provided. Default:False                                                                                                                                                                                           |
 | geometry_only                        | bool              | Whether to only model the SDF. If True, the feature prediction is ommited. Default:False                                                                                                                                                                                         |
 | fix_geometry                         | bool              | Whether to optimize the geometry. If True, the SDF (and grid vertices if `isosurface_deformable_grid=True`) is fixed. Default: False                                                                                                                                             |
 
+### Custom mesh
+
+|  shape_init                   | str        | The shape to initializa the SDF as. Should be formatted as "mesh:path", where `path` points to the custom mesh. Default: ""                            |
+| shape_init_params            | Optional[Any]       | Parameters to specify the SDF initialization. A single float is used for uniform scaling; a tuple of three floats is used for scalings along x/y/z axis. Default: None   |
+
 ## Material
 
 The material module outputs colors or color latents conditioned on the sampled positions, view directions, and sometimes light directions and normals.

configs/fantasia3d-texture.yaml

@@ -23,6 +23,8 @@ system_type: "fantasia3d-system"
 system:
   # do texture training
   texture: true
+
+  # If using geometry from previous training
   geometry_convert_from: ???
   geometry_convert_inherit_texture: false
   geometry_type: "tetrahedra-sdf-grid"
@@ -40,6 +42,20 @@ system:
     n_feature_dims: 8 # albedo3 + roughness1 + metallic1 + bump3
     fix_geometry: true
 
+  # If using custom mesh
+  # geometry_type: "custom-mesh"
+  # geometry:
+  #   shape_init: ???
+  #   radius: 1.0 # consistent with coarse
+  #   pos_encoding_config:
+  #     otype: HashGrid
+  #     n_levels: 16
+  #     n_features_per_level: 2
+  #     log2_hashmap_size: 19
+  #     base_resolution: 16
+  #     per_level_scale: 1.4472692374403782 # max resolution 4096
+  #   n_feature_dims: 8 # albedo3 + roughness1 + metallic1 + bump3
+
   material_type: "pbr-material"
   material:
     material_activation: sigmoid

threestudio/models/geometry/__init__.py

@@ -1 +1,8 @@
-from . import base, implicit_sdf, implicit_volume, tetrahedra_sdf_grid, volume_grid
+from . import (
+    base,
+    custom_mesh,
+    implicit_sdf,
+    implicit_volume,
+    tetrahedra_sdf_grid,
+    volume_grid,
+)

threestudio/models/geometry/custom_mesh.py

@@ -0,0 +1,178 @@
+import os
+from dataclasses import dataclass, field
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import threestudio
+from threestudio.models.geometry.base import (
+    BaseExplicitGeometry,
+    BaseGeometry,
+    contract_to_unisphere,
+)
+from threestudio.models.mesh import Mesh
+from threestudio.models.networks import get_encoding, get_mlp
+from threestudio.utils.ops import scale_tensor
+from threestudio.utils.typing import *
+
+
+@threestudio.register("custom-mesh")
+class CustomMesh(BaseExplicitGeometry):
+    @dataclass
+    class Config(BaseExplicitGeometry.Config):
+        n_input_dims: int = 3
+        n_feature_dims: int = 3
+        pos_encoding_config: dict = field(
+            default_factory=lambda: {
+                "otype": "HashGrid",
+                "n_levels": 16,
+                "n_features_per_level": 2,
+                "log2_hashmap_size": 19,
+                "base_resolution": 16,
+                "per_level_scale": 1.447269237440378,
+            }
+        )
+        mlp_network_config: dict = field(
+            default_factory=lambda: {
+                "otype": "VanillaMLP",
+                "activation": "ReLU",
+                "output_activation": "none",
+                "n_neurons": 64,
+                "n_hidden_layers": 1,
+            }
+        )
+        shape_init: str = ""
+        shape_init_params: Optional[Any] = None
+        shape_init_mesh_up: str = "+z"
+        shape_init_mesh_front: str = "+x"
+
+    cfg: Config
+
+    def configure(self) -> None:
+        super().configure()
+
+        self.encoding = get_encoding(
+            self.cfg.n_input_dims, self.cfg.pos_encoding_config
+        )
+        self.feature_network = get_mlp(
+            self.encoding.n_output_dims,
+            self.cfg.n_feature_dims,
+            self.cfg.mlp_network_config,
+        )
+
+        # Initialize custom mesh
+        if self.cfg.shape_init.startswith("mesh:"):
+            assert isinstance(self.cfg.shape_init_params, float)
+            mesh_path = self.cfg.shape_init[5:]
+            if not os.path.exists(mesh_path):
+                raise ValueError(f"Mesh file {mesh_path} does not exist.")
+
+            import trimesh
+
+            scene = trimesh.load(mesh_path)
+            if isinstance(scene, trimesh.Trimesh):
+                mesh = scene
+            elif isinstance(scene, trimesh.scene.Scene):
+                mesh = trimesh.Trimesh()
+                for obj in scene.geometry.values():
+                    mesh = trimesh.util.concatenate([mesh, obj])
+            else:
+                raise ValueError(f"Unknown mesh type at {mesh_path}.")
+
+            # move to center
+            centroid = mesh.vertices.mean(0)
+            mesh.vertices = mesh.vertices - centroid
+
+            # align to up-z and front-x
+            dirs = ["+x", "+y", "+z", "-x", "-y", "-z"]
+            dir2vec = {
+                "+x": np.array([1, 0, 0]),
+                "+y": np.array([0, 1, 0]),
+                "+z": np.array([0, 0, 1]),
+                "-x": np.array([-1, 0, 0]),
+                "-y": np.array([0, -1, 0]),
+                "-z": np.array([0, 0, -1]),
+            }
+            if (
+                self.cfg.shape_init_mesh_up not in dirs
+                or self.cfg.shape_init_mesh_front not in dirs
+            ):
+                raise ValueError(
+                    f"shape_init_mesh_up and shape_init_mesh_front must be one of {dirs}."
+                )
+            if self.cfg.shape_init_mesh_up[1] == self.cfg.shape_init_mesh_front[1]:
+                raise ValueError(
+                    "shape_init_mesh_up and shape_init_mesh_front must be orthogonal."
+                )
+            z_, x_ = (
+                dir2vec[self.cfg.shape_init_mesh_up],
+                dir2vec[self.cfg.shape_init_mesh_front],
+            )
+            y_ = np.cross(z_, x_)
+            std2mesh = np.stack([x_, y_, z_], axis=0).T
+            mesh2std = np.linalg.inv(std2mesh)
+
+            # scaling
+            scale = np.abs(mesh.vertices).max()
+            mesh.vertices = mesh.vertices / scale * self.cfg.shape_init_params
+            mesh.vertices = np.dot(mesh2std, mesh.vertices.T).T
+
+            v_pos = torch.tensor(mesh.vertices, dtype=torch.float32).to(self.device)
+            t_pos_idx = torch.tensor(mesh.faces, dtype=torch.int64).to(self.device)
+            self.mesh = Mesh(v_pos=v_pos, t_pos_idx=t_pos_idx)
+            self.register_buffer(
+                "v_buffer",
+                v_pos,
+            )
+            self.register_buffer(
+                "t_buffer",
+                t_pos_idx,
+            )
+
+        else:
+            raise ValueError(
+                f"Unknown shape initialization type: {self.cfg.shape_init}"
+            )
+        print(self.mesh.v_pos.device)
+
+    def isosurface(self) -> Mesh:
+        if hasattr(self, "mesh"):
+            return self.mesh
+        elif hasattr(self, "v_buffer"):
+            self.mesh = Mesh(v_pos=self.v_buffer, t_pos_idx=self.t_buffer)
+            return self.mesh
+        else:
+            raise ValueError(f"custom mesh is not initialized")
+
+    def forward(
+        self, points: Float[Tensor, "*N Di"], output_normal: bool = False
+    ) -> Dict[str, Float[Tensor, "..."]]:
+        assert (
+            output_normal == False
+        ), f"Normal output is not supported for {self.__class__.__name__}"
+        points_unscaled = points  # points in the original scale
+        points = contract_to_unisphere(points, self.bbox)  # points normalized to (0, 1)
+        enc = self.encoding(points.view(-1, self.cfg.n_input_dims))
+        features = self.feature_network(enc).view(
+            *points.shape[:-1], self.cfg.n_feature_dims
+        )
+        return {"features": features}
+
+    def export(self, points: Float[Tensor, "*N Di"], **kwargs) -> Dict[str, Any]:
+        out: Dict[str, Any] = {}
+        if self.cfg.n_feature_dims == 0:
+            return out
+        points_unscaled = points
+        points = contract_to_unisphere(points_unscaled, self.bbox)
+        enc = self.encoding(points.reshape(-1, self.cfg.n_input_dims))
+        features = self.feature_network(enc).view(
+            *points.shape[:-1], self.cfg.n_feature_dims
+        )
+        out.update(
+            {
+                "features": features,
+            }
+        )
+        return out

threestudio/models/geometry/implicit_sdf.py

@@ -132,7 +132,15 @@ def func(points_rand: Float[Tensor, "N 3"]) -> Float[Tensor, "N 1"]:
 
             import trimesh
 
-            mesh = trimesh.load(mesh_path)
+            scene = trimesh.load(mesh_path)
+            if isinstance(scene, trimesh.Trimesh):
+                mesh = scene
+            elif isinstance(scene, trimesh.scene.Scene):
+                mesh = trimesh.Trimesh()
+                for obj in scene.geometry.values():
+                    mesh = trimesh.util.concatenate([mesh, obj])
+            else:
+                raise ValueError(f"Unknown mesh type at {mesh_path}.")
 
             # move to center
             centroid = mesh.vertices.mean(0)