mirror of
https://github.com/deepseek-ai/DreamCraft3D
synced 2025-06-26 18:25:49 +00:00
chores: rebase commits
This commit is contained in:
MrTornado24
8
threestudio/models/geometry/__init__.py
Normal file
8
threestudio/models/geometry/__init__.py
Normal file
@@ -0,0 +1,8 @@
|
||||
from . import (
|
||||
base,
|
||||
custom_mesh,
|
||||
implicit_sdf,
|
||||
implicit_volume,
|
||||
tetrahedra_sdf_grid,
|
||||
volume_grid,
|
||||
)
|
||||
209
threestudio/models/geometry/base.py
Normal file
209
threestudio/models/geometry/base.py
Normal file
@@ -0,0 +1,209 @@
|
||||
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.isosurface import (
|
||||
IsosurfaceHelper,
|
||||
MarchingCubeCPUHelper,
|
||||
MarchingTetrahedraHelper,
|
||||
)
|
||||
from threestudio.models.mesh import Mesh
|
||||
from threestudio.utils.base import BaseModule
|
||||
from threestudio.utils.ops import chunk_batch, scale_tensor
|
||||
from threestudio.utils.typing import *
|
||||
|
||||
|
||||
def contract_to_unisphere(
|
||||
x: Float[Tensor, "... 3"], bbox: Float[Tensor, "2 3"], unbounded: bool = False
|
||||
) -> Float[Tensor, "... 3"]:
|
||||
if unbounded:
|
||||
x = scale_tensor(x, bbox, (0, 1))
|
||||
x = x * 2 - 1 # aabb is at [-1, 1]
|
||||
mag = x.norm(dim=-1, keepdim=True)
|
||||
mask = mag.squeeze(-1) > 1
|
||||
x[mask] = (2 - 1 / mag[mask]) * (x[mask] / mag[mask])
|
||||
x = x / 4 + 0.5 # [-inf, inf] is at [0, 1]
|
||||
else:
|
||||
x = scale_tensor(x, bbox, (0, 1))
|
||||
return x
|
||||
|
||||
|
||||
class BaseGeometry(BaseModule):
|
||||
@dataclass
|
||||
class Config(BaseModule.Config):
|
||||
pass
|
||||
|
||||
cfg: Config
|
||||
|
||||
@staticmethod
|
||||
def create_from(
|
||||
other: "BaseGeometry", cfg: Optional[Union[dict, DictConfig]] = None, **kwargs
|
||||
) -> "BaseGeometry":
|
||||
raise TypeError(
|
||||
f"Cannot create {BaseGeometry.__name__} from {other.__class__.__name__}"
|
||||
)
|
||||
|
||||
def export(self, *args, **kwargs) -> Dict[str, Any]:
|
||||
return {}
|
||||
|
||||
|
||||
class BaseImplicitGeometry(BaseGeometry):
|
||||
@dataclass
|
||||
class Config(BaseGeometry.Config):
|
||||
radius: float = 1.0
|
||||
isosurface: bool = True
|
||||
isosurface_method: str = "mt"
|
||||
isosurface_resolution: int = 128
|
||||
isosurface_threshold: Union[float, str] = 0.0
|
||||
isosurface_chunk: int = 0
|
||||
isosurface_coarse_to_fine: bool = True
|
||||
isosurface_deformable_grid: bool = False
|
||||
isosurface_remove_outliers: bool = True
|
||||
isosurface_outlier_n_faces_threshold: Union[int, float] = 0.01
|
||||
|
||||
cfg: Config
|
||||
|
||||
def configure(self) -> None:
|
||||
self.bbox: Float[Tensor, "2 3"]
|
||||
self.register_buffer(
|
||||
"bbox",
|
||||
torch.as_tensor(
|
||||
[
|
||||
[-self.cfg.radius, -self.cfg.radius, -self.cfg.radius],
|
||||
[self.cfg.radius, self.cfg.radius, self.cfg.radius],
|
||||
],
|
||||
dtype=torch.float32,
|
||||
),
|
||||
)
|
||||
self.isosurface_helper: Optional[IsosurfaceHelper] = None
|
||||
self.unbounded: bool = False
|
||||
|
||||
def _initilize_isosurface_helper(self):
|
||||
if self.cfg.isosurface and self.isosurface_helper is None:
|
||||
if self.cfg.isosurface_method == "mc-cpu":
|
||||
self.isosurface_helper = MarchingCubeCPUHelper(
|
||||
self.cfg.isosurface_resolution
|
||||
).to(self.device)
|
||||
elif self.cfg.isosurface_method == "mt":
|
||||
self.isosurface_helper = MarchingTetrahedraHelper(
|
||||
self.cfg.isosurface_resolution,
|
||||
f"load/tets/{self.cfg.isosurface_resolution}_tets.npz",
|
||||
).to(self.device)
|
||||
else:
|
||||
raise AttributeError(
|
||||
"Unknown isosurface method {self.cfg.isosurface_method}"
|
||||
)
|
||||
|
||||
def forward(
|
||||
self, points: Float[Tensor, "*N Di"], output_normal: bool = False
|
||||
) -> Dict[str, Float[Tensor, "..."]]:
|
||||
raise NotImplementedError
|
||||
|
||||
def forward_field(
|
||||
self, points: Float[Tensor, "*N Di"]
|
||||
) -> Tuple[Float[Tensor, "*N 1"], Optional[Float[Tensor, "*N 3"]]]:
|
||||
# return the value of the implicit field, could be density / signed distance
|
||||
# also return a deformation field if the grid vertices can be optimized
|
||||
raise NotImplementedError
|
||||
|
||||
def forward_level(
|
||||
self, field: Float[Tensor, "*N 1"], threshold: float
|
||||
) -> Float[Tensor, "*N 1"]:
|
||||
# return the value of the implicit field, where the zero level set represents the surface
|
||||
raise NotImplementedError
|
||||
|
||||
def _isosurface(self, bbox: Float[Tensor, "2 3"], fine_stage: bool = False) -> Mesh:
|
||||
def batch_func(x):
|
||||
# scale to bbox as the input vertices are in [0, 1]
|
||||
field, deformation = self.forward_field(
|
||||
scale_tensor(
|
||||
x.to(bbox.device), self.isosurface_helper.points_range, bbox
|
||||
),
|
||||
)
|
||||
field = field.to(
|
||||
x.device
|
||||
) # move to the same device as the input (could be CPU)
|
||||
if deformation is not None:
|
||||
deformation = deformation.to(x.device)
|
||||
return field, deformation
|
||||
|
||||
assert self.isosurface_helper is not None
|
||||
|
||||
field, deformation = chunk_batch(
|
||||
batch_func,
|
||||
self.cfg.isosurface_chunk,
|
||||
self.isosurface_helper.grid_vertices,
|
||||
)
|
||||
|
||||
threshold: float
|
||||
|
||||
if isinstance(self.cfg.isosurface_threshold, float):
|
||||
threshold = self.cfg.isosurface_threshold
|
||||
elif self.cfg.isosurface_threshold == "auto":
|
||||
eps = 1.0e-5
|
||||
threshold = field[field > eps].mean().item()
|
||||
threestudio.info(
|
||||
f"Automatically determined isosurface threshold: {threshold}"
|
||||
)
|
||||
else:
|
||||
raise TypeError(
|
||||
f"Unknown isosurface_threshold {self.cfg.isosurface_threshold}"
|
||||
)
|
||||
|
||||
level = self.forward_level(field, threshold)
|
||||
mesh: Mesh = self.isosurface_helper(level, deformation=deformation)
|
||||
mesh.v_pos = scale_tensor(
|
||||
mesh.v_pos, self.isosurface_helper.points_range, bbox
|
||||
) # scale to bbox as the grid vertices are in [0, 1]
|
||||
mesh.add_extra("bbox", bbox)
|
||||
|
||||
if self.cfg.isosurface_remove_outliers:
|
||||
# remove outliers components with small number of faces
|
||||
# only enabled when the mesh is not differentiable
|
||||
mesh = mesh.remove_outlier(self.cfg.isosurface_outlier_n_faces_threshold)
|
||||
|
||||
return mesh
|
||||
|
||||
def isosurface(self) -> Mesh:
|
||||
if not self.cfg.isosurface:
|
||||
raise NotImplementedError(
|
||||
"Isosurface is not enabled in the current configuration"
|
||||
)
|
||||
self._initilize_isosurface_helper()
|
||||
if self.cfg.isosurface_coarse_to_fine:
|
||||
threestudio.debug("First run isosurface to get a tight bounding box ...")
|
||||
with torch.no_grad():
|
||||
mesh_coarse = self._isosurface(self.bbox)
|
||||
vmin, vmax = mesh_coarse.v_pos.amin(dim=0), mesh_coarse.v_pos.amax(dim=0)
|
||||
vmin_ = (vmin - (vmax - vmin) * 0.1).max(self.bbox[0])
|
||||
vmax_ = (vmax + (vmax - vmin) * 0.1).min(self.bbox[1])
|
||||
threestudio.debug("Run isosurface again with the tight bounding box ...")
|
||||
mesh = self._isosurface(torch.stack([vmin_, vmax_], dim=0), fine_stage=True)
|
||||
else:
|
||||
mesh = self._isosurface(self.bbox)
|
||||
return mesh
|
||||
|
||||
|
||||
class BaseExplicitGeometry(BaseGeometry):
|
||||
@dataclass
|
||||
class Config(BaseGeometry.Config):
|
||||
radius: float = 1.0
|
||||
|
||||
cfg: Config
|
||||
|
||||
def configure(self) -> None:
|
||||
self.bbox: Float[Tensor, "2 3"]
|
||||
self.register_buffer(
|
||||
"bbox",
|
||||
torch.as_tensor(
|
||||
[
|
||||
[-self.cfg.radius, -self.cfg.radius, -self.cfg.radius],
|
||||
[self.cfg.radius, self.cfg.radius, self.cfg.radius],
|
||||
],
|
||||
dtype=torch.float32,
|
||||
),
|
||||
)
|
||||
178
threestudio/models/geometry/custom_mesh.py
Normal file
178
threestudio/models/geometry/custom_mesh.py
Normal file
@@ -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
|
||||
413
threestudio/models/geometry/implicit_sdf.py
Normal file
413
threestudio/models/geometry/implicit_sdf.py
Normal file
@@ -0,0 +1,413 @@
|
||||
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 BaseImplicitGeometry, contract_to_unisphere
|
||||
from threestudio.models.mesh import Mesh
|
||||
from threestudio.models.networks import get_encoding, get_mlp
|
||||
from threestudio.utils.misc import broadcast, get_rank
|
||||
from threestudio.utils.typing import *
|
||||
|
||||
|
||||
@threestudio.register("implicit-sdf")
|
||||
class ImplicitSDF(BaseImplicitGeometry):
|
||||
@dataclass
|
||||
class Config(BaseImplicitGeometry.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,
|
||||
}
|
||||
)
|
||||
normal_type: Optional[
|
||||
str
|
||||
] = "finite_difference" # in ['pred', 'finite_difference', 'finite_difference_laplacian']
|
||||
finite_difference_normal_eps: Union[
|
||||
float, str
|
||||
] = 0.01 # in [float, "progressive"]
|
||||
shape_init: Optional[str] = None
|
||||
shape_init_params: Optional[Any] = None
|
||||
shape_init_mesh_up: str = "+z"
|
||||
shape_init_mesh_front: str = "+x"
|
||||
force_shape_init: bool = False
|
||||
sdf_bias: Union[float, str] = 0.0
|
||||
sdf_bias_params: Optional[Any] = None
|
||||
|
||||
# no need to removal outlier for SDF
|
||||
isosurface_remove_outliers: bool = False
|
||||
|
||||
cfg: Config
|
||||
|
||||
def configure(self) -> None:
|
||||
super().configure()
|
||||
self.encoding = get_encoding(
|
||||
self.cfg.n_input_dims, self.cfg.pos_encoding_config
|
||||
)
|
||||
self.sdf_network = get_mlp(
|
||||
self.encoding.n_output_dims, 1, self.cfg.mlp_network_config
|
||||
)
|
||||
|
||||
if self.cfg.n_feature_dims > 0:
|
||||
self.feature_network = get_mlp(
|
||||
self.encoding.n_output_dims,
|
||||
self.cfg.n_feature_dims,
|
||||
self.cfg.mlp_network_config,
|
||||
)
|
||||
|
||||
if self.cfg.normal_type == "pred":
|
||||
self.normal_network = get_mlp(
|
||||
self.encoding.n_output_dims, 3, self.cfg.mlp_network_config
|
||||
)
|
||||
if self.cfg.isosurface_deformable_grid:
|
||||
assert (
|
||||
self.cfg.isosurface_method == "mt"
|
||||
), "isosurface_deformable_grid only works with mt"
|
||||
self.deformation_network = get_mlp(
|
||||
self.encoding.n_output_dims, 3, self.cfg.mlp_network_config
|
||||
)
|
||||
|
||||
self.finite_difference_normal_eps: Optional[float] = None
|
||||
|
||||
def initialize_shape(self) -> None:
|
||||
if self.cfg.shape_init is None and not self.cfg.force_shape_init:
|
||||
return
|
||||
|
||||
# do not initialize shape if weights are provided
|
||||
if self.cfg.weights is not None and not self.cfg.force_shape_init:
|
||||
return
|
||||
|
||||
if self.cfg.sdf_bias != 0.0:
|
||||
threestudio.warn(
|
||||
"shape_init and sdf_bias are both specified, which may lead to unexpected results."
|
||||
)
|
||||
|
||||
get_gt_sdf: Callable[[Float[Tensor, "N 3"]], Float[Tensor, "N 1"]]
|
||||
assert isinstance(self.cfg.shape_init, str)
|
||||
if self.cfg.shape_init == "ellipsoid":
|
||||
assert (
|
||||
isinstance(self.cfg.shape_init_params, Sized)
|
||||
and len(self.cfg.shape_init_params) == 3
|
||||
)
|
||||
size = torch.as_tensor(self.cfg.shape_init_params).to(self.device)
|
||||
|
||||
def func(points_rand: Float[Tensor, "N 3"]) -> Float[Tensor, "N 1"]:
|
||||
return ((points_rand / size) ** 2).sum(
|
||||
dim=-1, keepdim=True
|
||||
).sqrt() - 1.0 # pseudo signed distance of an ellipsoid
|
||||
|
||||
get_gt_sdf = func
|
||||
elif self.cfg.shape_init == "sphere":
|
||||
assert isinstance(self.cfg.shape_init_params, float)
|
||||
radius = self.cfg.shape_init_params
|
||||
|
||||
def func(points_rand: Float[Tensor, "N 3"]) -> Float[Tensor, "N 1"]:
|
||||
return (points_rand**2).sum(dim=-1, keepdim=True).sqrt() - radius
|
||||
|
||||
get_gt_sdf = func
|
||||
elif 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
|
||||
|
||||
from pysdf import SDF
|
||||
|
||||
sdf = SDF(mesh.vertices, mesh.faces)
|
||||
|
||||
def func(points_rand: Float[Tensor, "N 3"]) -> Float[Tensor, "N 1"]:
|
||||
# add a negative signed here
|
||||
# as in pysdf the inside of the shape has positive signed distance
|
||||
return torch.from_numpy(-sdf(points_rand.cpu().numpy())).to(
|
||||
points_rand
|
||||
)[..., None]
|
||||
|
||||
get_gt_sdf = func
|
||||
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Unknown shape initialization type: {self.cfg.shape_init}"
|
||||
)
|
||||
|
||||
# Initialize SDF to a given shape when no weights are provided or force_shape_init is True
|
||||
optim = torch.optim.Adam(self.parameters(), lr=1e-3)
|
||||
from tqdm import tqdm
|
||||
|
||||
for _ in tqdm(
|
||||
range(1000),
|
||||
desc=f"Initializing SDF to a(n) {self.cfg.shape_init}:",
|
||||
disable=get_rank() != 0,
|
||||
):
|
||||
points_rand = (
|
||||
torch.rand((10000, 3), dtype=torch.float32).to(self.device) * 2.0 - 1.0
|
||||
)
|
||||
sdf_gt = get_gt_sdf(points_rand)
|
||||
sdf_pred = self.forward_sdf(points_rand)
|
||||
loss = F.mse_loss(sdf_pred, sdf_gt)
|
||||
optim.zero_grad()
|
||||
loss.backward()
|
||||
optim.step()
|
||||
|
||||
# explicit broadcast to ensure param consistency across ranks
|
||||
for param in self.parameters():
|
||||
broadcast(param, src=0)
|
||||
|
||||
def get_shifted_sdf(
|
||||
self, points: Float[Tensor, "*N Di"], sdf: Float[Tensor, "*N 1"]
|
||||
) -> Float[Tensor, "*N 1"]:
|
||||
sdf_bias: Union[float, Float[Tensor, "*N 1"]]
|
||||
if self.cfg.sdf_bias == "ellipsoid":
|
||||
assert (
|
||||
isinstance(self.cfg.sdf_bias_params, Sized)
|
||||
and len(self.cfg.sdf_bias_params) == 3
|
||||
)
|
||||
size = torch.as_tensor(self.cfg.sdf_bias_params).to(points)
|
||||
sdf_bias = ((points / size) ** 2).sum(
|
||||
dim=-1, keepdim=True
|
||||
).sqrt() - 1.0 # pseudo signed distance of an ellipsoid
|
||||
elif self.cfg.sdf_bias == "sphere":
|
||||
assert isinstance(self.cfg.sdf_bias_params, float)
|
||||
radius = self.cfg.sdf_bias_params
|
||||
sdf_bias = (points**2).sum(dim=-1, keepdim=True).sqrt() - radius
|
||||
elif isinstance(self.cfg.sdf_bias, float):
|
||||
sdf_bias = self.cfg.sdf_bias
|
||||
else:
|
||||
raise ValueError(f"Unknown sdf bias {self.cfg.sdf_bias}")
|
||||
return sdf + sdf_bias
|
||||
|
||||
def forward(
|
||||
self, points: Float[Tensor, "*N Di"], output_normal: bool = False
|
||||
) -> Dict[str, Float[Tensor, "..."]]:
|
||||
grad_enabled = torch.is_grad_enabled()
|
||||
|
||||
if output_normal and self.cfg.normal_type == "analytic":
|
||||
torch.set_grad_enabled(True)
|
||||
points.requires_grad_(True)
|
||||
|
||||
points_unscaled = points # points in the original scale
|
||||
points = contract_to_unisphere(
|
||||
points, self.bbox, self.unbounded
|
||||
) # points normalized to (0, 1)
|
||||
|
||||
enc = self.encoding(points.view(-1, self.cfg.n_input_dims))
|
||||
sdf = self.sdf_network(enc).view(*points.shape[:-1], 1)
|
||||
sdf = self.get_shifted_sdf(points_unscaled, sdf)
|
||||
output = {"sdf": sdf}
|
||||
|
||||
if self.cfg.n_feature_dims > 0:
|
||||
features = self.feature_network(enc).view(
|
||||
*points.shape[:-1], self.cfg.n_feature_dims
|
||||
)
|
||||
output.update({"features": features})
|
||||
|
||||
if output_normal:
|
||||
if (
|
||||
self.cfg.normal_type == "finite_difference"
|
||||
or self.cfg.normal_type == "finite_difference_laplacian"
|
||||
):
|
||||
assert self.finite_difference_normal_eps is not None
|
||||
eps: float = self.finite_difference_normal_eps
|
||||
if self.cfg.normal_type == "finite_difference_laplacian":
|
||||
offsets: Float[Tensor, "6 3"] = torch.as_tensor(
|
||||
[
|
||||
[eps, 0.0, 0.0],
|
||||
[-eps, 0.0, 0.0],
|
||||
[0.0, eps, 0.0],
|
||||
[0.0, -eps, 0.0],
|
||||
[0.0, 0.0, eps],
|
||||
[0.0, 0.0, -eps],
|
||||
]
|
||||
).to(points_unscaled)
|
||||
points_offset: Float[Tensor, "... 6 3"] = (
|
||||
points_unscaled[..., None, :] + offsets
|
||||
).clamp(-self.cfg.radius, self.cfg.radius)
|
||||
sdf_offset: Float[Tensor, "... 6 1"] = self.forward_sdf(
|
||||
points_offset
|
||||
)
|
||||
sdf_grad = (
|
||||
0.5
|
||||
* (sdf_offset[..., 0::2, 0] - sdf_offset[..., 1::2, 0])
|
||||
/ eps
|
||||
)
|
||||
else:
|
||||
offsets: Float[Tensor, "3 3"] = torch.as_tensor(
|
||||
[[eps, 0.0, 0.0], [0.0, eps, 0.0], [0.0, 0.0, eps]]
|
||||
).to(points_unscaled)
|
||||
points_offset: Float[Tensor, "... 3 3"] = (
|
||||
points_unscaled[..., None, :] + offsets
|
||||
).clamp(-self.cfg.radius, self.cfg.radius)
|
||||
sdf_offset: Float[Tensor, "... 3 1"] = self.forward_sdf(
|
||||
points_offset
|
||||
)
|
||||
sdf_grad = (sdf_offset[..., 0::1, 0] - sdf) / eps
|
||||
normal = F.normalize(sdf_grad, dim=-1)
|
||||
elif self.cfg.normal_type == "pred":
|
||||
normal = self.normal_network(enc).view(*points.shape[:-1], 3)
|
||||
normal = F.normalize(normal, dim=-1)
|
||||
sdf_grad = normal
|
||||
elif self.cfg.normal_type == "analytic":
|
||||
sdf_grad = -torch.autograd.grad(
|
||||
sdf,
|
||||
points_unscaled,
|
||||
grad_outputs=torch.ones_like(sdf),
|
||||
create_graph=True,
|
||||
)[0]
|
||||
normal = F.normalize(sdf_grad, dim=-1)
|
||||
if not grad_enabled:
|
||||
sdf_grad = sdf_grad.detach()
|
||||
normal = normal.detach()
|
||||
else:
|
||||
raise AttributeError(f"Unknown normal type {self.cfg.normal_type}")
|
||||
output.update(
|
||||
{"normal": normal, "shading_normal": normal, "sdf_grad": sdf_grad}
|
||||
)
|
||||
return output
|
||||
|
||||
def forward_sdf(self, points: Float[Tensor, "*N Di"]) -> Float[Tensor, "*N 1"]:
|
||||
points_unscaled = points
|
||||
points = contract_to_unisphere(points_unscaled, self.bbox, self.unbounded)
|
||||
|
||||
sdf = self.sdf_network(
|
||||
self.encoding(points.reshape(-1, self.cfg.n_input_dims))
|
||||
).reshape(*points.shape[:-1], 1)
|
||||
sdf = self.get_shifted_sdf(points_unscaled, sdf)
|
||||
return sdf
|
||||
|
||||
def forward_field(
|
||||
self, points: Float[Tensor, "*N Di"]
|
||||
) -> Tuple[Float[Tensor, "*N 1"], Optional[Float[Tensor, "*N 3"]]]:
|
||||
points_unscaled = points
|
||||
points = contract_to_unisphere(points_unscaled, self.bbox, self.unbounded)
|
||||
enc = self.encoding(points.reshape(-1, self.cfg.n_input_dims))
|
||||
sdf = self.sdf_network(enc).reshape(*points.shape[:-1], 1)
|
||||
sdf = self.get_shifted_sdf(points_unscaled, sdf)
|
||||
deformation: Optional[Float[Tensor, "*N 3"]] = None
|
||||
if self.cfg.isosurface_deformable_grid:
|
||||
deformation = self.deformation_network(enc).reshape(*points.shape[:-1], 3)
|
||||
return sdf, deformation
|
||||
|
||||
def forward_level(
|
||||
self, field: Float[Tensor, "*N 1"], threshold: float
|
||||
) -> Float[Tensor, "*N 1"]:
|
||||
return field - threshold
|
||||
|
||||
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, self.unbounded)
|
||||
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
|
||||
|
||||
def update_step(self, epoch: int, global_step: int, on_load_weights: bool = False):
|
||||
if (
|
||||
self.cfg.normal_type == "finite_difference"
|
||||
or self.cfg.normal_type == "finite_difference_laplacian"
|
||||
):
|
||||
if isinstance(self.cfg.finite_difference_normal_eps, float):
|
||||
self.finite_difference_normal_eps = (
|
||||
self.cfg.finite_difference_normal_eps
|
||||
)
|
||||
elif self.cfg.finite_difference_normal_eps == "progressive":
|
||||
# progressive finite difference eps from Neuralangelo
|
||||
# https://arxiv.org/abs/2306.03092
|
||||
hg_conf: Any = self.cfg.pos_encoding_config
|
||||
assert (
|
||||
hg_conf.otype == "ProgressiveBandHashGrid"
|
||||
), "finite_difference_normal_eps=progressive only works with ProgressiveBandHashGrid"
|
||||
current_level = min(
|
||||
hg_conf.start_level
|
||||
+ max(global_step - hg_conf.start_step, 0) // hg_conf.update_steps,
|
||||
hg_conf.n_levels,
|
||||
)
|
||||
grid_res = hg_conf.base_resolution * hg_conf.per_level_scale ** (
|
||||
current_level - 1
|
||||
)
|
||||
grid_size = 2 * self.cfg.radius / grid_res
|
||||
if grid_size != self.finite_difference_normal_eps:
|
||||
threestudio.info(
|
||||
f"Update finite_difference_normal_eps to {grid_size}"
|
||||
)
|
||||
self.finite_difference_normal_eps = grid_size
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Unknown finite_difference_normal_eps={self.cfg.finite_difference_normal_eps}"
|
||||
)
|
||||
325
threestudio/models/geometry/implicit_volume.py
Normal file
325
threestudio/models/geometry/implicit_volume.py
Normal file
@@ -0,0 +1,325 @@
|
||||
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 (
|
||||
BaseGeometry,
|
||||
BaseImplicitGeometry,
|
||||
contract_to_unisphere,
|
||||
)
|
||||
from threestudio.models.networks import get_encoding, get_mlp
|
||||
from threestudio.utils.ops import get_activation
|
||||
from threestudio.utils.typing import *
|
||||
|
||||
|
||||
@threestudio.register("implicit-volume")
|
||||
class ImplicitVolume(BaseImplicitGeometry):
|
||||
@dataclass
|
||||
class Config(BaseImplicitGeometry.Config):
|
||||
n_input_dims: int = 3
|
||||
n_feature_dims: int = 3
|
||||
density_activation: Optional[str] = "softplus"
|
||||
density_bias: Union[float, str] = "blob_magic3d"
|
||||
density_blob_scale: float = 10.0
|
||||
density_blob_std: float = 0.5
|
||||
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,
|
||||
}
|
||||
)
|
||||
normal_type: Optional[
|
||||
str
|
||||
] = "finite_difference" # in ['pred', 'finite_difference', 'finite_difference_laplacian']
|
||||
finite_difference_normal_eps: Union[
|
||||
float, str
|
||||
] = 0.01 # in [float, "progressive"]
|
||||
|
||||
# automatically determine the threshold
|
||||
isosurface_threshold: Union[float, str] = 25.0
|
||||
|
||||
# 4D Gaussian Annealing
|
||||
anneal_density_blob_std_config: Optional[dict] = None
|
||||
|
||||
cfg: Config
|
||||
|
||||
def configure(self) -> None:
|
||||
super().configure()
|
||||
self.encoding = get_encoding(
|
||||
self.cfg.n_input_dims, self.cfg.pos_encoding_config
|
||||
)
|
||||
self.density_network = get_mlp(
|
||||
self.encoding.n_output_dims, 1, self.cfg.mlp_network_config
|
||||
)
|
||||
if self.cfg.n_feature_dims > 0:
|
||||
self.feature_network = get_mlp(
|
||||
self.encoding.n_output_dims,
|
||||
self.cfg.n_feature_dims,
|
||||
self.cfg.mlp_network_config,
|
||||
)
|
||||
if self.cfg.normal_type == "pred":
|
||||
self.normal_network = get_mlp(
|
||||
self.encoding.n_output_dims, 3, self.cfg.mlp_network_config
|
||||
)
|
||||
|
||||
self.finite_difference_normal_eps: Optional[float] = None
|
||||
|
||||
def get_activated_density(
|
||||
self, points: Float[Tensor, "*N Di"], density: Float[Tensor, "*N 1"]
|
||||
) -> Tuple[Float[Tensor, "*N 1"], Float[Tensor, "*N 1"]]:
|
||||
density_bias: Union[float, Float[Tensor, "*N 1"]]
|
||||
if self.cfg.density_bias == "blob_dreamfusion":
|
||||
# pre-activation density bias
|
||||
density_bias = (
|
||||
self.cfg.density_blob_scale
|
||||
* torch.exp(
|
||||
-0.5 * (points**2).sum(dim=-1) / self.cfg.density_blob_std**2
|
||||
)[..., None]
|
||||
)
|
||||
elif self.cfg.density_bias == "blob_magic3d":
|
||||
# pre-activation density bias
|
||||
density_bias = (
|
||||
self.cfg.density_blob_scale
|
||||
* (
|
||||
1
|
||||
- torch.sqrt((points**2).sum(dim=-1)) / self.cfg.density_blob_std
|
||||
)[..., None]
|
||||
)
|
||||
elif isinstance(self.cfg.density_bias, float):
|
||||
density_bias = self.cfg.density_bias
|
||||
else:
|
||||
raise ValueError(f"Unknown density bias {self.cfg.density_bias}")
|
||||
raw_density: Float[Tensor, "*N 1"] = density + density_bias
|
||||
density = get_activation(self.cfg.density_activation)(raw_density)
|
||||
return raw_density, density
|
||||
|
||||
def forward(
|
||||
self, points: Float[Tensor, "*N Di"], output_normal: bool = False
|
||||
) -> Dict[str, Float[Tensor, "..."]]:
|
||||
grad_enabled = torch.is_grad_enabled()
|
||||
|
||||
if output_normal and self.cfg.normal_type == "analytic":
|
||||
torch.set_grad_enabled(True)
|
||||
points.requires_grad_(True)
|
||||
|
||||
points_unscaled = points # points in the original scale
|
||||
points = contract_to_unisphere(
|
||||
points, self.bbox, self.unbounded
|
||||
) # points normalized to (0, 1)
|
||||
|
||||
enc = self.encoding(points.view(-1, self.cfg.n_input_dims))
|
||||
density = self.density_network(enc).view(*points.shape[:-1], 1)
|
||||
raw_density, density = self.get_activated_density(points_unscaled, density)
|
||||
|
||||
output = {
|
||||
"density": density,
|
||||
}
|
||||
|
||||
if self.cfg.n_feature_dims > 0:
|
||||
features = self.feature_network(enc).view(
|
||||
*points.shape[:-1], self.cfg.n_feature_dims
|
||||
)
|
||||
output.update({"features": features})
|
||||
|
||||
if output_normal:
|
||||
if (
|
||||
self.cfg.normal_type == "finite_difference"
|
||||
or self.cfg.normal_type == "finite_difference_laplacian"
|
||||
):
|
||||
# TODO: use raw density
|
||||
assert self.finite_difference_normal_eps is not None
|
||||
eps: float = self.finite_difference_normal_eps
|
||||
if self.cfg.normal_type == "finite_difference_laplacian":
|
||||
offsets: Float[Tensor, "6 3"] = torch.as_tensor(
|
||||
[
|
||||
[eps, 0.0, 0.0],
|
||||
[-eps, 0.0, 0.0],
|
||||
[0.0, eps, 0.0],
|
||||
[0.0, -eps, 0.0],
|
||||
[0.0, 0.0, eps],
|
||||
[0.0, 0.0, -eps],
|
||||
]
|
||||
).to(points_unscaled)
|
||||
points_offset: Float[Tensor, "... 6 3"] = (
|
||||
points_unscaled[..., None, :] + offsets
|
||||
).clamp(-self.cfg.radius, self.cfg.radius)
|
||||
density_offset: Float[Tensor, "... 6 1"] = self.forward_density(
|
||||
points_offset
|
||||
)
|
||||
normal = (
|
||||
-0.5
|
||||
* (density_offset[..., 0::2, 0] - density_offset[..., 1::2, 0])
|
||||
/ eps
|
||||
)
|
||||
else:
|
||||
offsets: Float[Tensor, "3 3"] = torch.as_tensor(
|
||||
[[eps, 0.0, 0.0], [0.0, eps, 0.0], [0.0, 0.0, eps]]
|
||||
).to(points_unscaled)
|
||||
points_offset: Float[Tensor, "... 3 3"] = (
|
||||
points_unscaled[..., None, :] + offsets
|
||||
).clamp(-self.cfg.radius, self.cfg.radius)
|
||||
density_offset: Float[Tensor, "... 3 1"] = self.forward_density(
|
||||
points_offset
|
||||
)
|
||||
normal = -(density_offset[..., 0::1, 0] - density) / eps
|
||||
normal = F.normalize(normal, dim=-1)
|
||||
elif self.cfg.normal_type == "pred":
|
||||
normal = self.normal_network(enc).view(*points.shape[:-1], 3)
|
||||
normal = F.normalize(normal, dim=-1)
|
||||
elif self.cfg.normal_type == "analytic":
|
||||
normal = -torch.autograd.grad(
|
||||
density,
|
||||
points_unscaled,
|
||||
grad_outputs=torch.ones_like(density),
|
||||
create_graph=True,
|
||||
)[0]
|
||||
normal = F.normalize(normal, dim=-1)
|
||||
if not grad_enabled:
|
||||
normal = normal.detach()
|
||||
else:
|
||||
raise AttributeError(f"Unknown normal type {self.cfg.normal_type}")
|
||||
output.update({"normal": normal, "shading_normal": normal})
|
||||
|
||||
torch.set_grad_enabled(grad_enabled)
|
||||
return output
|
||||
|
||||
def forward_density(self, points: Float[Tensor, "*N Di"]) -> Float[Tensor, "*N 1"]:
|
||||
points_unscaled = points
|
||||
points = contract_to_unisphere(points_unscaled, self.bbox, self.unbounded)
|
||||
|
||||
density = self.density_network(
|
||||
self.encoding(points.reshape(-1, self.cfg.n_input_dims))
|
||||
).reshape(*points.shape[:-1], 1)
|
||||
|
||||
_, density = self.get_activated_density(points_unscaled, density)
|
||||
return density
|
||||
|
||||
def forward_field(
|
||||
self, points: Float[Tensor, "*N Di"]
|
||||
) -> Tuple[Float[Tensor, "*N 1"], Optional[Float[Tensor, "*N 3"]]]:
|
||||
if self.cfg.isosurface_deformable_grid:
|
||||
threestudio.warn(
|
||||
f"{self.__class__.__name__} does not support isosurface_deformable_grid. Ignoring."
|
||||
)
|
||||
density = self.forward_density(points)
|
||||
return density, None
|
||||
|
||||
def forward_level(
|
||||
self, field: Float[Tensor, "*N 1"], threshold: float
|
||||
) -> Float[Tensor, "*N 1"]:
|
||||
return -(field - threshold)
|
||||
|
||||
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, self.unbounded)
|
||||
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
|
||||
|
||||
@staticmethod
|
||||
@torch.no_grad()
|
||||
def create_from(
|
||||
other: BaseGeometry,
|
||||
cfg: Optional[Union[dict, DictConfig]] = None,
|
||||
copy_net: bool = True,
|
||||
**kwargs,
|
||||
) -> "ImplicitVolume":
|
||||
if isinstance(other, ImplicitVolume):
|
||||
instance = ImplicitVolume(cfg, **kwargs)
|
||||
instance.encoding.load_state_dict(other.encoding.state_dict())
|
||||
instance.density_network.load_state_dict(other.density_network.state_dict())
|
||||
if copy_net:
|
||||
if (
|
||||
instance.cfg.n_feature_dims > 0
|
||||
and other.cfg.n_feature_dims == instance.cfg.n_feature_dims
|
||||
):
|
||||
instance.feature_network.load_state_dict(
|
||||
other.feature_network.state_dict()
|
||||
)
|
||||
if (
|
||||
instance.cfg.normal_type == "pred"
|
||||
and other.cfg.normal_type == "pred"
|
||||
):
|
||||
instance.normal_network.load_state_dict(
|
||||
other.normal_network.state_dict()
|
||||
)
|
||||
return instance
|
||||
else:
|
||||
raise TypeError(
|
||||
f"Cannot create {ImplicitVolume.__name__} from {other.__class__.__name__}"
|
||||
)
|
||||
|
||||
# FIXME: use progressive normal eps
|
||||
def update_step(
|
||||
self, epoch: int, global_step: int, on_load_weights: bool = False
|
||||
) -> None:
|
||||
if self.cfg.anneal_density_blob_std_config is not None:
|
||||
min_step = self.cfg.anneal_density_blob_std_config.min_anneal_step
|
||||
max_step = self.cfg.anneal_density_blob_std_config.max_anneal_step
|
||||
if global_step >= min_step and global_step <= max_step:
|
||||
end_val = self.cfg.anneal_density_blob_std_config.end_val
|
||||
start_val = self.cfg.anneal_density_blob_std_config.start_val
|
||||
self.density_blob_std = start_val + (global_step - min_step) * (
|
||||
end_val - start_val
|
||||
) / (max_step - min_step)
|
||||
|
||||
if (
|
||||
self.cfg.normal_type == "finite_difference"
|
||||
or self.cfg.normal_type == "finite_difference_laplacian"
|
||||
):
|
||||
if isinstance(self.cfg.finite_difference_normal_eps, float):
|
||||
self.finite_difference_normal_eps = (
|
||||
self.cfg.finite_difference_normal_eps
|
||||
)
|
||||
elif self.cfg.finite_difference_normal_eps == "progressive":
|
||||
# progressive finite difference eps from Neuralangelo
|
||||
# https://arxiv.org/abs/2306.03092
|
||||
hg_conf: Any = self.cfg.pos_encoding_config
|
||||
assert (
|
||||
hg_conf.otype == "ProgressiveBandHashGrid"
|
||||
), "finite_difference_normal_eps=progressive only works with ProgressiveBandHashGrid"
|
||||
current_level = min(
|
||||
hg_conf.start_level
|
||||
+ max(global_step - hg_conf.start_step, 0) // hg_conf.update_steps,
|
||||
hg_conf.n_levels,
|
||||
)
|
||||
grid_res = hg_conf.base_resolution * hg_conf.per_level_scale ** (
|
||||
current_level - 1
|
||||
)
|
||||
grid_size = 2 * self.cfg.radius / grid_res
|
||||
if grid_size != self.finite_difference_normal_eps:
|
||||
threestudio.info(
|
||||
f"Update finite_difference_normal_eps to {grid_size}"
|
||||
)
|
||||
self.finite_difference_normal_eps = grid_size
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Unknown finite_difference_normal_eps={self.cfg.finite_difference_normal_eps}"
|
||||
)
|
||||
369
threestudio/models/geometry/tetrahedra_sdf_grid.py
Normal file
369
threestudio/models/geometry/tetrahedra_sdf_grid.py
Normal file
@@ -0,0 +1,369 @@
|
||||
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.geometry.implicit_sdf import ImplicitSDF
|
||||
from threestudio.models.geometry.implicit_volume import ImplicitVolume
|
||||
from threestudio.models.isosurface import MarchingTetrahedraHelper
|
||||
from threestudio.models.mesh import Mesh
|
||||
from threestudio.models.networks import get_encoding, get_mlp
|
||||
from threestudio.utils.misc import broadcast
|
||||
from threestudio.utils.ops import scale_tensor
|
||||
from threestudio.utils.typing import *
|
||||
|
||||
|
||||
@threestudio.register("tetrahedra-sdf-grid")
|
||||
class TetrahedraSDFGrid(BaseExplicitGeometry):
|
||||
@dataclass
|
||||
class Config(BaseExplicitGeometry.Config):
|
||||
isosurface_resolution: int = 128
|
||||
isosurface_deformable_grid: bool = True
|
||||
isosurface_remove_outliers: bool = False
|
||||
isosurface_outlier_n_faces_threshold: Union[int, float] = 0.01
|
||||
|
||||
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: Optional[str] = None
|
||||
shape_init_params: Optional[Any] = None
|
||||
shape_init_mesh_up: str = "+z"
|
||||
shape_init_mesh_front: str = "+x"
|
||||
force_shape_init: bool = False
|
||||
geometry_only: bool = False
|
||||
fix_geometry: bool = False
|
||||
|
||||
cfg: Config
|
||||
|
||||
def configure(self) -> None:
|
||||
super().configure()
|
||||
|
||||
# this should be saved to state_dict, register as buffer
|
||||
self.isosurface_bbox: Float[Tensor, "2 3"]
|
||||
self.register_buffer("isosurface_bbox", self.bbox.clone())
|
||||
|
||||
self.isosurface_helper = MarchingTetrahedraHelper(
|
||||
self.cfg.isosurface_resolution,
|
||||
f"load/tets/{self.cfg.isosurface_resolution}_tets.npz",
|
||||
)
|
||||
|
||||
self.sdf: Float[Tensor, "Nv 1"]
|
||||
self.deformation: Optional[Float[Tensor, "Nv 3"]]
|
||||
|
||||
if not self.cfg.fix_geometry:
|
||||
self.register_parameter(
|
||||
"sdf",
|
||||
nn.Parameter(
|
||||
torch.zeros(
|
||||
(self.isosurface_helper.grid_vertices.shape[0], 1),
|
||||
dtype=torch.float32,
|
||||
)
|
||||
),
|
||||
)
|
||||
if self.cfg.isosurface_deformable_grid:
|
||||
self.register_parameter(
|
||||
"deformation",
|
||||
nn.Parameter(
|
||||
torch.zeros_like(self.isosurface_helper.grid_vertices)
|
||||
),
|
||||
)
|
||||
else:
|
||||
self.deformation = None
|
||||
else:
|
||||
self.register_buffer(
|
||||
"sdf",
|
||||
torch.zeros(
|
||||
(self.isosurface_helper.grid_vertices.shape[0], 1),
|
||||
dtype=torch.float32,
|
||||
),
|
||||
)
|
||||
if self.cfg.isosurface_deformable_grid:
|
||||
self.register_buffer(
|
||||
"deformation",
|
||||
torch.zeros_like(self.isosurface_helper.grid_vertices),
|
||||
)
|
||||
else:
|
||||
self.deformation = None
|
||||
|
||||
if not self.cfg.geometry_only:
|
||||
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,
|
||||
)
|
||||
|
||||
self.mesh: Optional[Mesh] = None
|
||||
|
||||
def initialize_shape(self) -> None:
|
||||
if self.cfg.shape_init is None and not self.cfg.force_shape_init:
|
||||
return
|
||||
|
||||
# do not initialize shape if weights are provided
|
||||
if self.cfg.weights is not None and not self.cfg.force_shape_init:
|
||||
return
|
||||
|
||||
get_gt_sdf: Callable[[Float[Tensor, "N 3"]], Float[Tensor, "N 1"]]
|
||||
assert isinstance(self.cfg.shape_init, str)
|
||||
if self.cfg.shape_init == "ellipsoid":
|
||||
assert (
|
||||
isinstance(self.cfg.shape_init_params, Sized)
|
||||
and len(self.cfg.shape_init_params) == 3
|
||||
)
|
||||
size = torch.as_tensor(self.cfg.shape_init_params).to(self.device)
|
||||
|
||||
def func(points_rand: Float[Tensor, "N 3"]) -> Float[Tensor, "N 1"]:
|
||||
return ((points_rand / size) ** 2).sum(
|
||||
dim=-1, keepdim=True
|
||||
).sqrt() - 1.0 # pseudo signed distance of an ellipsoid
|
||||
|
||||
get_gt_sdf = func
|
||||
elif self.cfg.shape_init == "sphere":
|
||||
assert isinstance(self.cfg.shape_init_params, float)
|
||||
radius = self.cfg.shape_init_params
|
||||
|
||||
def func(points_rand: Float[Tensor, "N 3"]) -> Float[Tensor, "N 1"]:
|
||||
return (points_rand**2).sum(dim=-1, keepdim=True).sqrt() - radius
|
||||
|
||||
get_gt_sdf = func
|
||||
elif 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
|
||||
|
||||
mesh = trimesh.load(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
|
||||
|
||||
from pysdf import SDF
|
||||
|
||||
sdf = SDF(mesh.vertices, mesh.faces)
|
||||
|
||||
def func(points_rand: Float[Tensor, "N 3"]) -> Float[Tensor, "N 1"]:
|
||||
# add a negative signed here
|
||||
# as in pysdf the inside of the shape has positive signed distance
|
||||
return torch.from_numpy(-sdf(points_rand.cpu().numpy())).to(
|
||||
points_rand
|
||||
)[..., None]
|
||||
|
||||
get_gt_sdf = func
|
||||
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Unknown shape initialization type: {self.cfg.shape_init}"
|
||||
)
|
||||
|
||||
sdf_gt = get_gt_sdf(
|
||||
scale_tensor(
|
||||
self.isosurface_helper.grid_vertices,
|
||||
self.isosurface_helper.points_range,
|
||||
self.isosurface_bbox,
|
||||
)
|
||||
)
|
||||
self.sdf.data = sdf_gt
|
||||
|
||||
# explicit broadcast to ensure param consistency across ranks
|
||||
for param in self.parameters():
|
||||
broadcast(param, src=0)
|
||||
|
||||
def isosurface(self) -> Mesh:
|
||||
# return cached mesh if fix_geometry is True to save computation
|
||||
if self.cfg.fix_geometry and self.mesh is not None:
|
||||
return self.mesh
|
||||
mesh = self.isosurface_helper(self.sdf, self.deformation)
|
||||
mesh.v_pos = scale_tensor(
|
||||
mesh.v_pos, self.isosurface_helper.points_range, self.isosurface_bbox
|
||||
)
|
||||
if self.cfg.isosurface_remove_outliers:
|
||||
mesh = mesh.remove_outlier(self.cfg.isosurface_outlier_n_faces_threshold)
|
||||
self.mesh = mesh
|
||||
return mesh
|
||||
|
||||
def forward(
|
||||
self, points: Float[Tensor, "*N Di"], output_normal: bool = False
|
||||
) -> Dict[str, Float[Tensor, "..."]]:
|
||||
if self.cfg.geometry_only:
|
||||
return {}
|
||||
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}
|
||||
|
||||
@staticmethod
|
||||
@torch.no_grad()
|
||||
def create_from(
|
||||
other: BaseGeometry,
|
||||
cfg: Optional[Union[dict, DictConfig]] = None,
|
||||
copy_net: bool = True,
|
||||
**kwargs,
|
||||
) -> "TetrahedraSDFGrid":
|
||||
if isinstance(other, TetrahedraSDFGrid):
|
||||
instance = TetrahedraSDFGrid(cfg, **kwargs)
|
||||
assert instance.cfg.isosurface_resolution == other.cfg.isosurface_resolution
|
||||
instance.isosurface_bbox = other.isosurface_bbox.clone()
|
||||
instance.sdf.data = other.sdf.data.clone()
|
||||
if (
|
||||
instance.cfg.isosurface_deformable_grid
|
||||
and other.cfg.isosurface_deformable_grid
|
||||
):
|
||||
assert (
|
||||
instance.deformation is not None and other.deformation is not None
|
||||
)
|
||||
instance.deformation.data = other.deformation.data.clone()
|
||||
if (
|
||||
not instance.cfg.geometry_only
|
||||
and not other.cfg.geometry_only
|
||||
and copy_net
|
||||
):
|
||||
instance.encoding.load_state_dict(other.encoding.state_dict())
|
||||
instance.feature_network.load_state_dict(
|
||||
other.feature_network.state_dict()
|
||||
)
|
||||
return instance
|
||||
elif isinstance(other, ImplicitVolume):
|
||||
instance = TetrahedraSDFGrid(cfg, **kwargs)
|
||||
if other.cfg.isosurface_method != "mt":
|
||||
other.cfg.isosurface_method = "mt"
|
||||
threestudio.warn(
|
||||
f"Override isosurface_method of the source geometry to 'mt'"
|
||||
)
|
||||
if other.cfg.isosurface_resolution != instance.cfg.isosurface_resolution:
|
||||
other.cfg.isosurface_resolution = instance.cfg.isosurface_resolution
|
||||
threestudio.warn(
|
||||
f"Override isosurface_resolution of the source geometry to {instance.cfg.isosurface_resolution}"
|
||||
)
|
||||
mesh = other.isosurface()
|
||||
instance.isosurface_bbox = mesh.extras["bbox"]
|
||||
instance.sdf.data = (
|
||||
mesh.extras["grid_level"].to(instance.sdf.data).clamp(-1, 1)
|
||||
)
|
||||
if not instance.cfg.geometry_only and copy_net:
|
||||
instance.encoding.load_state_dict(other.encoding.state_dict())
|
||||
instance.feature_network.load_state_dict(
|
||||
other.feature_network.state_dict()
|
||||
)
|
||||
return instance
|
||||
elif isinstance(other, ImplicitSDF):
|
||||
instance = TetrahedraSDFGrid(cfg, **kwargs)
|
||||
if other.cfg.isosurface_method != "mt":
|
||||
other.cfg.isosurface_method = "mt"
|
||||
threestudio.warn(
|
||||
f"Override isosurface_method of the source geometry to 'mt'"
|
||||
)
|
||||
if other.cfg.isosurface_resolution != instance.cfg.isosurface_resolution:
|
||||
other.cfg.isosurface_resolution = instance.cfg.isosurface_resolution
|
||||
threestudio.warn(
|
||||
f"Override isosurface_resolution of the source geometry to {instance.cfg.isosurface_resolution}"
|
||||
)
|
||||
mesh = other.isosurface()
|
||||
instance.isosurface_bbox = mesh.extras["bbox"]
|
||||
instance.sdf.data = mesh.extras["grid_level"].to(instance.sdf.data)
|
||||
if (
|
||||
instance.cfg.isosurface_deformable_grid
|
||||
and other.cfg.isosurface_deformable_grid
|
||||
):
|
||||
assert instance.deformation is not None
|
||||
instance.deformation.data = mesh.extras["grid_deformation"].to(
|
||||
instance.deformation.data
|
||||
)
|
||||
if not instance.cfg.geometry_only and copy_net:
|
||||
instance.encoding.load_state_dict(other.encoding.state_dict())
|
||||
instance.feature_network.load_state_dict(
|
||||
other.feature_network.state_dict()
|
||||
)
|
||||
return instance
|
||||
else:
|
||||
raise TypeError(
|
||||
f"Cannot create {TetrahedraSDFGrid.__name__} from {other.__class__.__name__}"
|
||||
)
|
||||
|
||||
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
|
||||
190
threestudio/models/geometry/volume_grid.py
Normal file
190
threestudio/models/geometry/volume_grid.py
Normal file
@@ -0,0 +1,190 @@
|
||||
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 BaseImplicitGeometry, contract_to_unisphere
|
||||
from threestudio.utils.ops import get_activation
|
||||
from threestudio.utils.typing import *
|
||||
|
||||
|
||||
@threestudio.register("volume-grid")
|
||||
class VolumeGrid(BaseImplicitGeometry):
|
||||
@dataclass
|
||||
class Config(BaseImplicitGeometry.Config):
|
||||
grid_size: Tuple[int, int, int] = field(default_factory=lambda: (100, 100, 100))
|
||||
n_feature_dims: int = 3
|
||||
density_activation: Optional[str] = "softplus"
|
||||
density_bias: Union[float, str] = "blob"
|
||||
density_blob_scale: float = 5.0
|
||||
density_blob_std: float = 0.5
|
||||
normal_type: Optional[
|
||||
str
|
||||
] = "finite_difference" # in ['pred', 'finite_difference', 'finite_difference_laplacian']
|
||||
|
||||
# automatically determine the threshold
|
||||
isosurface_threshold: Union[float, str] = "auto"
|
||||
|
||||
cfg: Config
|
||||
|
||||
def configure(self) -> None:
|
||||
super().configure()
|
||||
self.grid_size = self.cfg.grid_size
|
||||
|
||||
self.grid = nn.Parameter(
|
||||
torch.zeros(1, self.cfg.n_feature_dims + 1, *self.grid_size)
|
||||
)
|
||||
if self.cfg.density_bias == "blob":
|
||||
self.register_buffer("density_scale", torch.tensor(0.0))
|
||||
else:
|
||||
self.density_scale = nn.Parameter(torch.tensor(0.0))
|
||||
|
||||
if self.cfg.normal_type == "pred":
|
||||
self.normal_grid = nn.Parameter(torch.zeros(1, 3, *self.grid_size))
|
||||
|
||||
def get_density_bias(self, points: Float[Tensor, "*N Di"]):
|
||||
if self.cfg.density_bias == "blob":
|
||||
# density_bias: Float[Tensor, "*N 1"] = self.cfg.density_blob_scale * torch.exp(-0.5 * (points ** 2).sum(dim=-1) / self.cfg.density_blob_std ** 2)[...,None]
|
||||
density_bias: Float[Tensor, "*N 1"] = (
|
||||
self.cfg.density_blob_scale
|
||||
* (
|
||||
1
|
||||
- torch.sqrt((points.detach() ** 2).sum(dim=-1))
|
||||
/ self.cfg.density_blob_std
|
||||
)[..., None]
|
||||
)
|
||||
return density_bias
|
||||
elif isinstance(self.cfg.density_bias, float):
|
||||
return self.cfg.density_bias
|
||||
else:
|
||||
raise AttributeError(f"Unknown density bias {self.cfg.density_bias}")
|
||||
|
||||
def get_trilinear_feature(
|
||||
self, points: Float[Tensor, "*N Di"], grid: Float[Tensor, "1 Df G1 G2 G3"]
|
||||
) -> Float[Tensor, "*N Df"]:
|
||||
points_shape = points.shape[:-1]
|
||||
df = grid.shape[1]
|
||||
di = points.shape[-1]
|
||||
out = F.grid_sample(
|
||||
grid, points.view(1, 1, 1, -1, di), align_corners=False, mode="bilinear"
|
||||
)
|
||||
out = out.reshape(df, -1).T.reshape(*points_shape, df)
|
||||
return out
|
||||
|
||||
def forward(
|
||||
self, points: Float[Tensor, "*N Di"], output_normal: bool = False
|
||||
) -> Dict[str, Float[Tensor, "..."]]:
|
||||
points_unscaled = points # points in the original scale
|
||||
points = contract_to_unisphere(
|
||||
points, self.bbox, self.unbounded
|
||||
) # points normalized to (0, 1)
|
||||
points = points * 2 - 1 # convert to [-1, 1] for grid sample
|
||||
|
||||
out = self.get_trilinear_feature(points, self.grid)
|
||||
density, features = out[..., 0:1], out[..., 1:]
|
||||
density = density * torch.exp(self.density_scale) # exp scaling in DreamFusion
|
||||
|
||||
# breakpoint()
|
||||
density = get_activation(self.cfg.density_activation)(
|
||||
density + self.get_density_bias(points_unscaled)
|
||||
)
|
||||
|
||||
output = {
|
||||
"density": density,
|
||||
"features": features,
|
||||
}
|
||||
|
||||
if output_normal:
|
||||
if (
|
||||
self.cfg.normal_type == "finite_difference"
|
||||
or self.cfg.normal_type == "finite_difference_laplacian"
|
||||
):
|
||||
eps = 1.0e-3
|
||||
if self.cfg.normal_type == "finite_difference_laplacian":
|
||||
offsets: Float[Tensor, "6 3"] = torch.as_tensor(
|
||||
[
|
||||
[eps, 0.0, 0.0],
|
||||
[-eps, 0.0, 0.0],
|
||||
[0.0, eps, 0.0],
|
||||
[0.0, -eps, 0.0],
|
||||
[0.0, 0.0, eps],
|
||||
[0.0, 0.0, -eps],
|
||||
]
|
||||
).to(points_unscaled)
|
||||
points_offset: Float[Tensor, "... 6 3"] = (
|
||||
points_unscaled[..., None, :] + offsets
|
||||
).clamp(-self.cfg.radius, self.cfg.radius)
|
||||
density_offset: Float[Tensor, "... 6 1"] = self.forward_density(
|
||||
points_offset
|
||||
)
|
||||
normal = (
|
||||
-0.5
|
||||
* (density_offset[..., 0::2, 0] - density_offset[..., 1::2, 0])
|
||||
/ eps
|
||||
)
|
||||
else:
|
||||
offsets: Float[Tensor, "3 3"] = torch.as_tensor(
|
||||
[[eps, 0.0, 0.0], [0.0, eps, 0.0], [0.0, 0.0, eps]]
|
||||
).to(points_unscaled)
|
||||
points_offset: Float[Tensor, "... 3 3"] = (
|
||||
points_unscaled[..., None, :] + offsets
|
||||
).clamp(-self.cfg.radius, self.cfg.radius)
|
||||
density_offset: Float[Tensor, "... 3 1"] = self.forward_density(
|
||||
points_offset
|
||||
)
|
||||
normal = -(density_offset[..., 0::1, 0] - density) / eps
|
||||
normal = F.normalize(normal, dim=-1)
|
||||
elif self.cfg.normal_type == "pred":
|
||||
normal = self.get_trilinear_feature(points, self.normal_grid)
|
||||
normal = F.normalize(normal, dim=-1)
|
||||
else:
|
||||
raise AttributeError(f"Unknown normal type {self.cfg.normal_type}")
|
||||
output.update({"normal": normal, "shading_normal": normal})
|
||||
return output
|
||||
|
||||
def forward_density(self, points: Float[Tensor, "*N Di"]) -> Float[Tensor, "*N 1"]:
|
||||
points_unscaled = points
|
||||
points = contract_to_unisphere(points_unscaled, self.bbox, self.unbounded)
|
||||
points = points * 2 - 1 # convert to [-1, 1] for grid sample
|
||||
|
||||
out = self.get_trilinear_feature(points, self.grid)
|
||||
density = out[..., 0:1]
|
||||
density = density * torch.exp(self.density_scale)
|
||||
|
||||
density = get_activation(self.cfg.density_activation)(
|
||||
density + self.get_density_bias(points_unscaled)
|
||||
)
|
||||
return density
|
||||
|
||||
def forward_field(
|
||||
self, points: Float[Tensor, "*N Di"]
|
||||
) -> Tuple[Float[Tensor, "*N 1"], Optional[Float[Tensor, "*N 3"]]]:
|
||||
if self.cfg.isosurface_deformable_grid:
|
||||
threestudio.warn(
|
||||
f"{self.__class__.__name__} does not support isosurface_deformable_grid. Ignoring."
|
||||
)
|
||||
density = self.forward_density(points)
|
||||
return density, None
|
||||
|
||||
def forward_level(
|
||||
self, field: Float[Tensor, "*N 1"], threshold: float
|
||||
) -> Float[Tensor, "*N 1"]:
|
||||
return -(field - threshold)
|
||||
|
||||
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, self.bbox, self.unbounded)
|
||||
points = points * 2 - 1 # convert to [-1, 1] for grid sample
|
||||
features = self.get_trilinear_feature(points, self.grid)[..., 1:]
|
||||
out.update(
|
||||
{
|
||||
"features": features,
|
||||
}
|
||||
)
|
||||
return out
|
||||
Reference in New Issue
Block a user