nodes/MrForExample-ComfyUI-3D-Pack
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Integrated Triplane Gaussian Transformers & Better 3DGS Import & Export

Browse Source
This commit is contained in:
MrForExample
2024-01-25 15:09:07 +01:00
parent 553416ac15
commit b95a0e475a
67 changed files with 9212 additions and 1008 deletions
Download Patch File Download Diff File Expand all files Collapse all files
algorithms/diff_texturing.py
+1
View File
@@ -29,6 +29,7 @@ class DiffTextureBaker:
self.batch_size = batch_size
def prepare_img(self, img):
# (H, W, C) -> (C, H, W)
img_new = img.permute(2, 0, 1).unsqueeze(0).to(self.device)
img_new = F.interpolate(img_new, (self.ref_size_H, self.ref_size_W), mode="bilinear", align_corners=False).contiguous()
return img_new
algorithms/main_3DGS_renderer.py
+15 -45
View File
@@ -31,7 +31,7 @@ from simple_knn._C import distCUDA2
from ..shared_utils.sh_utils import eval_sh, SH2RGB, RGB2SH
from ..mesh_processer.mesh import Mesh
from ..mesh_processer.mesh_utils import decimate_mesh, clean_mesh
from ..mesh_processer.mesh_utils import read_gs_ply
def inverse_sigmoid(x):
return torch.log(x/(1-x))
@@ -500,6 +500,11 @@ class GaussianModel:
lr = self.xyz_scheduler_args(iteration)
param_group['lr'] = lr
return lr
def reset_opacity(self):
opacities_new = inverse_sigmoid(torch.min(self.get_opacity, torch.ones_like(self.get_opacity)*0.01))
optimizable_tensors = self.replace_tensor_to_optimizer(opacities_new, "opacity")
self._opacity = optimizable_tensors["opacity"]
def construct_list_of_attributes(self):
l = ['x', 'y', 'z', 'nx', 'ny', 'nz']
@@ -525,9 +530,7 @@ class GaussianModel:
def get_mesh(self):
xyz = self._xyz.detach().cpu().numpy()
def save_ply(self, path):
os.makedirs(os.path.dirname(path), exist_ok=True)
def to_ply(self):
xyz = self._xyz.detach().cpu().numpy()
normals = np.zeros_like(xyz)
f_dc = self._features_dc.detach().transpose(1, 2).flatten(start_dim=1).contiguous().cpu().numpy()
@@ -542,45 +545,10 @@ class GaussianModel:
attributes = np.concatenate((xyz, normals, f_dc, f_rest, opacities, scale, rotation), axis=1)
elements[:] = list(map(tuple, attributes))
el = PlyElement.describe(elements, 'vertex')
PlyData([el]).write(path)
return PlyData([el])
def reset_opacity(self):
opacities_new = inverse_sigmoid(torch.min(self.get_opacity, torch.ones_like(self.get_opacity)*0.01))
optimizable_tensors = self.replace_tensor_to_optimizer(opacities_new, "opacity")
self._opacity = optimizable_tensors["opacity"]
def load_ply(self, path):
plydata = PlyData.read(path)
xyz = np.stack((np.asarray(plydata.elements[0]["x"]),
np.asarray(plydata.elements[0]["y"]),
np.asarray(plydata.elements[0]["z"])), axis=1)
opacities = np.asarray(plydata.elements[0]["opacity"])[..., np.newaxis]
print("Number of points at loading : ", xyz.shape[0])
features_dc = np.zeros((xyz.shape[0], 3, 1))
features_dc[:, 0, 0] = np.asarray(plydata.elements[0]["f_dc_0"])
features_dc[:, 1, 0] = np.asarray(plydata.elements[0]["f_dc_1"])
features_dc[:, 2, 0] = np.asarray(plydata.elements[0]["f_dc_2"])
extra_f_names = [p.name for p in plydata.elements[0].properties if p.name.startswith("f_rest_")]
assert len(extra_f_names)==3*(self.max_sh_degree + 1) ** 2 - 3
features_extra = np.zeros((xyz.shape[0], len(extra_f_names)))
for idx, attr_name in enumerate(extra_f_names):
features_extra[:, idx] = np.asarray(plydata.elements[0][attr_name])
# Reshape (P,F*SH_coeffs) to (P, F, SH_coeffs except DC)
features_extra = features_extra.reshape((features_extra.shape[0], 3, (self.max_sh_degree + 1) ** 2 - 1))
scale_names = [p.name for p in plydata.elements[0].properties if p.name.startswith("scale_")]
scales = np.zeros((xyz.shape[0], len(scale_names)))
for idx, attr_name in enumerate(scale_names):
scales[:, idx] = np.asarray(plydata.elements[0][attr_name])
rot_names = [p.name for p in plydata.elements[0].properties if p.name.startswith("rot")]
rots = np.zeros((xyz.shape[0], len(rot_names)))
for idx, attr_name in enumerate(rot_names):
rots[:, idx] = np.asarray(plydata.elements[0][attr_name])
def create_from_ply(self, plydata):
xyz, features_dc, features_extra, opacities, scales, rots = read_gs_ply(plydata)
self._xyz = nn.Parameter(torch.tensor(xyz, dtype=torch.float, device="cuda").requires_grad_(True))
self._features_dc = nn.Parameter(torch.tensor(features_dc, dtype=torch.float, device="cuda").transpose(1, 2).contiguous().requires_grad_(True))
@@ -894,6 +862,11 @@ class Renderer:
if isinstance(input, Mesh):
# load from 3D mesh
self.gaussians.create_from_mesh(input, num_pts)
elif isinstance(input, PlyData):
self.gaussians.create_from_ply(input)
elif isinstance(input, PointCloud):
# load from a provided pcd
self.gaussians.create_from_pcd(input, 1)
elif input is None:
# init from random point cloud
@@ -913,9 +886,6 @@ class Renderer:
points=xyz, colors=SH2RGB(shs), normals=np.zeros((num_pts, 3))
)
self.gaussians.create_from_pcd(pcd, 10)
elif isinstance(input, PointCloud):
# load from a provided pcd
self.gaussians.create_from_pcd(input, 1)
else:
self.gaussians.create_from_uv_data(input)
algorithms/triplane_gaussian_transformers.py
+167
View File
@@ -0,0 +1,167 @@
import torch
from dataclasses import dataclass, field
from einops import rearrange
from ..tgs import find
from ..tgs.models.image_feature import ImageFeature
from ..tgs.utils.saving import SaverMixin
from ..tgs.utils.config import parse_structured
from ..tgs.utils.ops import points_projection
from ..tgs.utils.misc import load_module_weights
from ..tgs.utils.typing import *
class TGS(torch.nn.Module, SaverMixin):
@dataclass
class Config:
weights: Optional[str] = None
weights_ignore_modules: Optional[List[str]] = None
camera_embedder_cls: str = ""
camera_embedder: dict = field(default_factory=dict)
image_feature: dict = field(default_factory=dict)
image_tokenizer_cls: str = ""
image_tokenizer: dict = field(default_factory=dict)
tokenizer_cls: str = ""
tokenizer: dict = field(default_factory=dict)
backbone_cls: str = ""
backbone: dict = field(default_factory=dict)
post_processor_cls: str = ""
post_processor: dict = field(default_factory=dict)
renderer_cls: str = ""
renderer: dict = field(default_factory=dict)
pointcloud_generator_cls: str = ""
pointcloud_generator: dict = field(default_factory=dict)
pointcloud_encoder_cls: str = ""
pointcloud_encoder: dict = field(default_factory=dict)
cfg: Config
def load_weights(self, weights: str, ignore_modules: Optional[List[str]] = None):
state_dict = load_module_weights(
weights, ignore_modules=ignore_modules, map_location="cpu"
)
self.load_state_dict(state_dict, strict=False)
def __init__(self, cfg):
super().__init__()
self.cfg = parse_structured(self.Config, cfg)
self.image_tokenizer = find(self.cfg.image_tokenizer_cls)(
self.cfg.image_tokenizer
)
assert self.cfg.camera_embedder_cls == 'tgs.models.networks.MLP'
weights = self.cfg.camera_embedder.pop("weights") if "weights" in self.cfg.camera_embedder else None
self.camera_embedder = find(self.cfg.camera_embedder_cls)(**self.cfg.camera_embedder)
if weights:
from ..tgs.utils.misc import load_module_weights
weights_path, module_name = weights.split(":")
state_dict = load_module_weights(
weights_path, module_name=module_name, map_location="cpu"
)
self.camera_embedder.load_state_dict(state_dict)
self.image_feature = ImageFeature(self.cfg.image_feature)
self.tokenizer = find(self.cfg.tokenizer_cls)(self.cfg.tokenizer)
self.backbone = find(self.cfg.backbone_cls)(self.cfg.backbone)
self.post_processor = find(self.cfg.post_processor_cls)(
self.cfg.post_processor
)
self.renderer = find(self.cfg.renderer_cls)(self.cfg.renderer)
# pointcloud generator
self.pointcloud_generator = find(self.cfg.pointcloud_generator_cls)(self.cfg.pointcloud_generator)
self.point_encoder = find(self.cfg.pointcloud_encoder_cls)(self.cfg.pointcloud_encoder)
# load checkpoint
if self.cfg.weights is not None:
self.load_weights(self.cfg.weights, self.cfg.weights_ignore_modules)
def _forward(self, batch: Dict[str, Any]) -> Dict[str, Any]:
# generate point cloud
out = self.pointcloud_generator(batch)
pointclouds = out["points"]
batch_size, n_input_views = batch["rgb_cond"].shape[:2]
# Camera modulation
camera_extri = batch["c2w_cond"].view(*batch["c2w_cond"].shape[:-2], -1)
camera_intri = batch["intrinsic_normed_cond"].view(*batch["intrinsic_normed_cond"].shape[:-2], -1)
camera_feats = torch.cat([camera_intri, camera_extri], dim=-1)
camera_feats = self.camera_embedder(camera_feats)
input_image_tokens: Float[Tensor, "B Cit Nit"] = self.image_tokenizer(
rearrange(batch["rgb_cond"], 'B Nv H W C -> B Nv C H W'),
modulation_cond=camera_feats,
)
input_image_tokens = rearrange(input_image_tokens, 'B Nv C Nt -> B (Nv Nt) C', Nv=n_input_views)
# get image features for projection
image_features = self.image_feature(
rgb = batch["rgb_cond"],
mask = batch.get("mask_cond", None),
feature = input_image_tokens
)
# only support number of input view is one
c2w_cond = batch["c2w_cond"].squeeze(1)
intrinsic_cond = batch["intrinsic_cond"].squeeze(1)
proj_feats = points_projection(pointclouds, c2w_cond, intrinsic_cond, image_features)
point_cond_embeddings = self.point_encoder(torch.cat([pointclouds, proj_feats], dim=-1))
tokens: Float[Tensor, "B Ct Nt"] = self.tokenizer(batch_size, cond_embeddings=point_cond_embeddings)
tokens = self.backbone(
tokens,
encoder_hidden_states=input_image_tokens,
modulation_cond=None,
)
scene_codes = self.post_processor(self.tokenizer.detokenize(tokens))
rend_out = self.renderer(scene_codes,
query_points=pointclouds,
additional_features=proj_feats,
**batch)
return {**out, **rend_out}
def forward(self, batch):
plys = []
out = self._forward(batch)
batch_size = batch["index"].shape[0]
for b in range(batch_size):
if batch["view_index"][b, 0] == 0:
plys.append(out["3dgs"][b].to_ply())
#self.render_test_video(out, batch, b)
return plys
def render_test_video(self, out, batch, b):
for index, render_image in enumerate(out["comp_rgb"][b]):
view_index = batch["view_index"][b, index]
self.save_image_grid(
f"video/{batch['instance_id'][b]}/{view_index}.png",
[
{
"type": "rgb",
"img": render_image,
"kwargs": {"data_format": "HWC"},
}
]
)