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
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}"
|
||||
)
|
||||
Reference in New Issue
Block a user