From a6e1891c06b5446b82baf5bfe58aa7132701a09a Mon Sep 17 00:00:00 2001 From: yankeesong Date: Mon, 10 Jul 2023 13:48:54 -0400 Subject: [PATCH 1/4] custom_mesh --- DOCUMENTATION.md | 13 +- configs/fantasia3d-texture.yaml | 20 ++- threestudio/models/geometry/__init__.py | 9 +- threestudio/models/geometry/custom_mesh.py | 178 +++++++++++++++++++++ 4 files changed, 211 insertions(+), 9 deletions(-) create mode 100644 threestudio/models/geometry/custom_mesh.py diff --git a/DOCUMENTATION.md b/DOCUMENTATION.md index 70545f47..d2b73374 100644 --- a/DOCUMENTATION.md +++ b/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. diff --git a/configs/fantasia3d-texture.yaml b/configs/fantasia3d-texture.yaml index e232d557..ddf8d54c 100644 --- a/configs/fantasia3d-texture.yaml +++ b/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 @@ -85,10 +101,10 @@ system: eps: 1.e-15 trainer: - max_steps: 5000 + max_steps: 600 log_every_n_steps: 1 num_sanity_val_steps: 1 - val_check_interval: 500 + val_check_interval: 200 enable_progress_bar: true precision: 16-mixed diff --git a/threestudio/models/geometry/__init__.py b/threestudio/models/geometry/__init__.py index f88c4b8c..19185d26 100644 --- a/threestudio/models/geometry/__init__.py +++ b/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, +) diff --git a/threestudio/models/geometry/custom_mesh.py b/threestudio/models/geometry/custom_mesh.py new file mode 100644 index 00000000..8fe08a87 --- /dev/null +++ b/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.geometry_only or 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 From 93314af99c1b957d6127b858d94821be06cf634d Mon Sep 17 00:00:00 2001 From: yankeesong Date: Mon, 10 Jul 2023 13:54:48 -0400 Subject: [PATCH 2/4] typecheck for custom mesh --- threestudio/models/geometry/implicit_sdf.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/threestudio/models/geometry/implicit_sdf.py b/threestudio/models/geometry/implicit_sdf.py index 713e100b..7330df73 100644 --- a/threestudio/models/geometry/implicit_sdf.py +++ b/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) From 06df7ec732a0e0e1af64477d103f4dcf371b93a8 Mon Sep 17 00:00:00 2001 From: yankeesong Date: Mon, 10 Jul 2023 14:01:09 -0400 Subject: [PATCH 3/4] fix export --- threestudio/models/geometry/custom_mesh.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/threestudio/models/geometry/custom_mesh.py b/threestudio/models/geometry/custom_mesh.py index 8fe08a87..08ebc630 100644 --- a/threestudio/models/geometry/custom_mesh.py +++ b/threestudio/models/geometry/custom_mesh.py @@ -162,7 +162,7 @@ def forward( def export(self, points: Float[Tensor, "*N Di"], **kwargs) -> Dict[str, Any]: out: Dict[str, Any] = {} - if self.cfg.geometry_only or self.cfg.n_feature_dims == 0: + if self.cfg.n_feature_dims == 0: return out points_unscaled = points points = contract_to_unisphere(points_unscaled, self.bbox) From 48009b7ab598ed989939a7c095be730d4bb4a402 Mon Sep 17 00:00:00 2001 From: yankeesong Date: Tue, 18 Jul 2023 09:25:58 -0400 Subject: [PATCH 4/4] minor change revert --- configs/fantasia3d-texture.yaml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/configs/fantasia3d-texture.yaml b/configs/fantasia3d-texture.yaml index ddf8d54c..9d6fc3a3 100644 --- a/configs/fantasia3d-texture.yaml +++ b/configs/fantasia3d-texture.yaml @@ -101,10 +101,10 @@ system: eps: 1.e-15 trainer: - max_steps: 600 + max_steps: 5000 log_every_n_steps: 1 num_sanity_val_steps: 1 - val_check_interval: 200 + val_check_interval: 500 enable_progress_bar: true precision: 16-mixed