Add Instant_NGP

algorithms/Instant_NGP.py

@@ -0,0 +1,206 @@
+import tqdm
+import random
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import nerfacc
+
+import comfy.utils
+
+from pytorch_msssim import SSIM, MS_SSIM
+
+from kiui.op import safe_normalize
+from kiui.cam import orbit_camera
+from kiui.nn import MLP, trunc_exp
+
+from ..shared_utils.image_utils import prepare_torch_img
+
+
+class InstantNGP(nn.Module):
+    def __init__(self, resolution=128, device="cuda"):
+        super().__init__()
+        from kiui.gridencoder import GridEncoder
+        
+        self.device = torch.device(device)
+        self.ref_size_H = resolution
+        self.ref_size_W = resolution
+        
+        self.render_step_size = 5e-3
+        self.aabb = torch.tensor([-1.0, -1.0, -1.0, 1.0, 1.0, 1.0], device=self.device)
+        self.estimator = nerfacc.OccGridEstimator(roi_aabb=self.aabb, resolution=64, levels=1)
+
+        self.encoder_density = GridEncoder(num_levels=12) # VMEncoder(output_dim=16, mode='sum')
+        self.encoder = GridEncoder(num_levels=12)
+        self.mlp_density = MLP(self.encoder_density.output_dim, 1, 32, 2, bias=False)
+        self.mlp = MLP(self.encoder.output_dim, 3, 32, 2, bias=False)
+        
+    def get_rays(self, pose, h, w, fovy, opengl=True):
+
+        x, y = torch.meshgrid(
+            torch.arange(w, device=pose.device),
+            torch.arange(h, device=pose.device),
+            indexing="xy",
+        )
+        x = x.flatten()
+        y = y.flatten()
+
+        cx = w * 0.5
+        cy = h * 0.5
+
+        focal = h * 0.5 / np.tan(0.5 * np.deg2rad(fovy))
+
+        camera_dirs = F.pad(
+            torch.stack(
+                [
+                    (x - cx + 0.5) / focal,
+                    (y - cy + 0.5) / focal * (-1.0 if opengl else 1.0),
+                ],
+                dim=-1,
+            ),
+            (0, 1),
+            value=(-1.0 if opengl else 1.0),
+        )  # [hw, 3]
+
+        rays_d = camera_dirs @ pose[:3, :3].transpose(0, 1)  # [hw, 3]
+        rays_o = pose[:3, 3].unsqueeze(0).expand_as(rays_d) # [hw, 3]
+
+        rays_o = rays_o.view(h, w, 3)
+        rays_d = safe_normalize(rays_d).view(h, w, 3)
+
+        return rays_o, rays_d
+    
+    def get_color(self, xs):
+        return torch.sigmoid(self.mlp(self.encoder(xs.to(self.device))))
+    
+    def get_density(self, xs):
+        # xs: [..., 3]
+        xs = xs.to(self.device)
+        prefix = xs.shape[:-1]
+        xs = xs.view(-1, 3)
+        feats = self.encoder_density(xs)
+        density = trunc_exp(self.mlp_density(feats))
+        density = density.view(*prefix, 1)
+        return density
+    
+    def prepare_training(self, reference_images, reference_masks, reference_orbit_camera_poses, reference_orbit_camera_fovy):
+        self.ref_imgs_num = len(reference_images)
+
+        self.all_ref_cam_poses = reference_orbit_camera_poses
+        self.ref_cam_fovy = reference_orbit_camera_fovy
+        
+        # prepare reference images and masks
+        ref_imgs_torch_list = []
+        ref_masks_torch_list = []
+        for i in range(self.ref_imgs_num):
+            ref_imgs_torch_list.append(prepare_torch_img(reference_images[i].unsqueeze(0), self.ref_size_H, self.ref_size_W, self.device))
+            ref_masks_torch_list.append(prepare_torch_img(reference_masks[i].unsqueeze(2).unsqueeze(0), self.ref_size_H, self.ref_size_W, self.device))
+            
+        self.ref_imgs_torch = torch.cat(ref_imgs_torch_list, dim=0) # [N, 3, H, W]
+        self.ref_masks_torch = torch.cat(ref_masks_torch_list, dim=0).squeeze(1) # [N, H, W]
+        
+    def render_nerf(self, pose):
+        
+        pose = torch.from_numpy(pose.astype(np.float32)).to(self.device)
+        
+        # get rays
+        rays_o, rays_d = self.get_rays(pose, self.ref_size_H, self.ref_size_W, self.ref_cam_fovy)
+        hw = rays_o.shape[0] * rays_o.shape[1]
+        rays_o = rays_o.view(hw, 3)
+        rays_d = rays_d.view(hw, 3)
+        
+        # update occ grid
+        if self.training:
+            def occ_eval_fn(xs):
+                sigmas = self.get_density(xs)
+                return self.render_step_size * sigmas
+            
+            self.estimator.update_every_n_steps(self.render_step, occ_eval_fn=occ_eval_fn, occ_thre=0.01, n=8)
+            self.render_step += 1
+
+        # render
+        def sigma_fn(t_starts, t_ends, ray_indices):
+            t_origins = rays_o[ray_indices]
+            t_dirs = rays_d[ray_indices]
+            xs = t_origins + t_dirs * (t_starts + t_ends)[:, None] / 2.0
+            sigmas = self.get_density(xs)
+            return sigmas.squeeze(-1)
+
+        with torch.no_grad():
+            ray_indices, t_starts, t_ends = self.estimator.sampling(
+                rays_o,
+                rays_d,
+                sigma_fn=sigma_fn,
+                near_plane=0.01,
+                far_plane=100,
+                render_step_size=self.render_step_size,
+                stratified=self.training,
+                cone_angle=0,
+            )
+
+        t_origins = rays_o[ray_indices]
+        t_dirs = rays_d[ray_indices]
+        xs = t_origins + t_dirs * (t_starts + t_ends)[:, None] / 2.0
+        sigmas = self.get_density(xs).squeeze(-1)
+        rgbs = torch.sigmoid(self.mlp(self.encoder(xs)))
+
+        n_rays=rays_o.shape[0]
+        weights, t, alphas = nerfacc.render_weight_from_density(t_starts, t_ends, sigmas, ray_indices=ray_indices, n_rays=n_rays)
+        color = nerfacc.accumulate_along_rays(weights, values=rgbs, ray_indices=ray_indices, n_rays=n_rays)
+        alpha = nerfacc.accumulate_along_rays(weights, values=None, ray_indices=ray_indices, n_rays=n_rays)
+
+        color = color + 1 * (1.0 - alpha)
+
+        color = color.view(self.ref_size_H, self.ref_size_W, 3).clamp(0, 1).permute(2, 0, 1).contiguous()
+        alpha = alpha.view(self.ref_size_H, self.ref_size_W).clamp(0, 1).contiguous()
+        
+        return color, alpha
+
+    def fit_nerf(self, iters=512):
+
+        optimizer = torch.optim.Adam([
+            {'params': self.encoder_density.parameters(), 'lr': 1e-2},
+            {'params': self.encoder.parameters(), 'lr': 1e-2},
+            {'params': self.mlp_density.parameters(), 'lr': 1e-3},
+            {'params': self.mlp.parameters(), 'lr': 1e-3},
+        ])
+
+        print(f"[INFO] fitting nerf...")
+        self.render_step = 0
+        
+        ref_imgs_num_minus_1 = self.ref_imgs_num-1
+        
+        comfy_pbar = comfy.utils.ProgressBar(iters)
+        pbar = tqdm.trange(iters)
+        for step in pbar:
+            
+            i = random.randint(0, ref_imgs_num_minus_1)
+
+            radius, elevation, azimuth, center_X, center_Y, center_Z = self.all_ref_cam_poses[i]
+            
+            orbit_target = np.array([center_X, center_Y, center_Z], dtype=np.float32)
+            pose = orbit_camera(elevation, azimuth, radius, target=orbit_target)
+            
+            image_gt = self.ref_imgs_torch[i]   # [3, H, W]
+            alpha_gt = self.ref_masks_torch[i]  # [H, W]
+            image_pred, alpha_pred = self.render_nerf(pose)
+
+            # if i % 200 == 0:
+            #     kiui.vis.plot_image(image_gt, alpha_gt, image_pred, alpha_pred)
+            
+            loss_mse = F.mse_loss(image_pred, image_gt) + 0.1 * F.mse_loss(alpha_pred, alpha_gt)
+            loss = loss_mse #+ 0.1 * self.encoder_density.tv_loss() #+ 0.0001 * self.encoder_density.density_loss()
+            #loss += self.lambda_ssim * (1 - self.ms_ssim_loss(image_gt, image_pred))
+
+            loss.backward()
+            self.encoder_density.grad_total_variation(1e-8)
+        
+            optimizer.step()
+            optimizer.zero_grad()
+
+            pbar.set_description(f"NeRF Fitting Loss = {loss_mse.item():.6f}")
+            comfy_pbar.update_absolute(step + 1)
+            
+        torch.cuda.synchronize()
+        
+        print(f"[INFO] finished fitting nerf!")

algorithms/diff_mesh.py

@@ -2,9 +2,12 @@ import random
 import tqdm
 
 import torch
+import torch.nn as nn
 import torch.nn.functional as F
 import numpy as np
+from kiui.mesh_utils import clean_mesh, decimate_mesh
 from kiui.cam import orbit_camera
+from kiui.mesh_utils import laplacian_smooth_loss, normal_consistency
 from pytorch_msssim import SSIM, MS_SSIM
 
 import comfy.utils
@@ -13,11 +16,14 @@ from .diff_mesh_renderer import DiffRastRenderer
 from ..shared_utils.camera_utils import OrbitCamera
 from ..shared_utils.image_utils import prepare_torch_img
 
-class DiffTextureBaker:
+class DiffMesh:
     
-    def __init__(self, mesh, training_iterations, batch_size, texture_learning_rate, train_mesh_geometry, geometry_learning_rate, ms_ssim_loss_weight, force_cuda_rasterize):
+    def __init__(self, mesh, training_iterations, batch_size, texture_learning_rate, train_mesh_geometry, geometry_learning_rate, ms_ssim_loss_weight, remesh_after_n_iteration, force_cuda_rasterize):
         self.device = torch.device("cuda")
         
+        self.train_mesh_geometry = train_mesh_geometry
+        self.remesh_after_n_iteration = remesh_after_n_iteration
+        
         # prepare main components for optimization
         self.renderer = DiffRastRenderer(mesh, force_cuda_rasterize).to(self.device)
 
@@ -51,7 +57,7 @@ class DiffTextureBaker:
         self.ref_imgs_torch = torch.cat(ref_imgs_torch_list, dim=0)
         self.ref_masks_torch = torch.cat(ref_masks_torch_list, dim=0)
     
-    def training(self):
+    def training(self, decimate_target=5e4):
         starter = torch.cuda.Event(enable_timing=True)
         ender = torch.cuda.Event(enable_timing=True)
         starter.record()
@@ -100,6 +106,24 @@ class DiffTextureBaker:
             # [1, 3, H, W] in [0, 1]
             #loss += self.lambda_ssim * (1 - self.ssim_loss(X, Y))
             loss += self.lambda_ssim * (1 - self.ms_ssim_loss(masked_ref_img_batch_torch, masked_rendered_img_batch_torch))
+            
+            # Regularization loss
+            if self.train_mesh_geometry:
+                current_v = self.renderer.mesh.v + self.renderer.v_offsets
+                loss += 0.01 * laplacian_smooth_loss(current_v, self.renderer.mesh.f)
+                loss += 0.001 * normal_consistency(current_v, self.renderer.mesh.f)
+                loss += 0.1 * (self.renderer.v_offsets ** 2).sum(-1).mean()
+                
+                # remesh periodically
+                if step > 0 and step % self.remesh_after_n_iteration == 0:
+                    vertices = (self.renderer.mesh.v + self.renderer.v_offsets).detach().cpu().numpy()
+                    triangles = self.renderer.mesh.f.detach().cpu().numpy()
+                    vertices, triangles = clean_mesh(vertices, triangles, remesh=True, remesh_size=0.01)
+                    if triangles.shape[0] > decimate_target:
+                        vertices, triangles = decimate_mesh(vertices, triangles, decimate_target, optimalplacement=False)
+                    self.renderer.mesh.v = torch.from_numpy(vertices).contiguous().float().to(self.device)
+                    self.renderer.mesh.f = torch.from_numpy(triangles).contiguous().int().to(self.device)
+                    self.renderer.v_offsets = nn.Parameter(torch.zeros_like(self.renderer.mesh.v)).to(self.device)
 
             # optimize step
             loss.backward()

algorithms/diff_mesh_renderer.py

@@ -7,6 +7,8 @@ import torch.nn as nn
 import torch.nn.functional as F
 import nvdiffrast.torch as dr
 
+from kiui.op import inverse_sigmoid
+
 from ..mesh_processer.mesh import safe_normalize
 
 def scale_img_nhwc(x, size, mag='bilinear', min='bilinear'):
@@ -30,10 +32,6 @@ def scale_img_nhw(x, size, mag='bilinear', min='bilinear'):
 def scale_img_hw(x, size, mag='bilinear', min='bilinear'):
     return scale_img_nhwc(x[None, ..., None], size, mag, min)[0, ..., 0]
 
-def trunc_rev_sigmoid(x, eps=1e-6):
-    x = x.clamp(eps, 1 - eps)
-    return torch.log(x / (1 - x))
-
 def make_divisible(x, m=8):
     return int(math.ceil(x / m) * m)
 
@@ -51,7 +49,7 @@ class DiffRastRenderer(nn.Module):
         
         # extract trainable parameters
         self.v_offsets = nn.Parameter(torch.zeros_like(self.mesh.v), requires_grad=True)
-        self.raw_albedo = nn.Parameter(trunc_rev_sigmoid(self.mesh.albedo), requires_grad=True)
+        self.raw_albedo = nn.Parameter(inverse_sigmoid(self.mesh.albedo), requires_grad=True)
 
 
     def get_params(self, texture_lr, train_geo, geom_lr):

algorithms/gaussian_model.txt

@@ -0,0 +1,397 @@
+import torch
+import numpy as np
+from utils.general_utils import inverse_sigmoid, get_expon_lr_func, build_rotation
+from torch import nn
+import os
+from utils.system_utils import mkdir_p
+from plyfile import PlyData, PlyElement
+from utils.sh_utils import RGB2SH
+from simple_knn._C import distCUDA2
+from utils.graphics_utils import BasicPointCloud
+from utils.general_utils import strip_symmetric, build_scaling_rotation
+
+class GaussianModel:
+
+    def setup_functions(self):
+        def build_covariance_from_scaling_rotation(scaling, scaling_modifier, rotation):
+            L = build_scaling_rotation(scaling_modifier * scaling, rotation)
+            actual_covariance = L @ L.transpose(1, 2)
+            symm = strip_symmetric(actual_covariance)
+            return symm
+        
+        self.scaling_activation = torch.exp
+        self.scaling_inverse_activation = torch.log
+
+        self.covariance_activation = build_covariance_from_scaling_rotation
+
+        self.opacity_activation = torch.sigmoid
+        self.inverse_opacity_activation = inverse_sigmoid
+
+        self.rotation_activation = torch.nn.functional.normalize
+
+
+    def __init__(self, sh_degree : int, divide_ratio : float):
+        self.active_sh_degree = 0
+        self.max_sh_degree = sh_degree  
+        self._xyz = torch.empty(0)
+        self._features_dc = torch.empty(0)
+        self._features_rest = torch.empty(0)
+        self._scaling = torch.empty(0)
+        self._rotation = torch.empty(0)
+        self._opacity = torch.empty(0)
+        self.max_radii2D = torch.empty(0)
+        self.xyz_gradient_accum = torch.empty(0)
+        self.denom = torch.empty(0)
+        self.optimizer = None
+        self.percent_dense = 0
+        self.spatial_lr_scale = 0
+        self.setup_functions()
+        self.divide_ratio = divide_ratio
+
+    def capture(self):
+        return (
+            self.active_sh_degree,
+            self._xyz,
+            self._features_dc,
+            self._features_rest,
+            self._scaling,
+            self._rotation,
+            self._opacity,
+            self.max_radii2D,
+            self.xyz_gradient_accum,
+            self.denom,
+            self.optimizer.state_dict(),
+            self.spatial_lr_scale,
+        )
+    
+    def restore(self, model_args, training_args):
+        (self.active_sh_degree, 
+        self._xyz, 
+        self._features_dc, 
+        self._features_rest,
+        self._scaling, 
+        self._rotation, 
+        self._opacity,
+        self.max_radii2D, 
+        xyz_gradient_accum, 
+        denom,
+        opt_dict, 
+        self.spatial_lr_scale) = model_args
+        self.training_setup(training_args)
+        self.xyz_gradient_accum = xyz_gradient_accum
+        self.denom = denom
+        self.optimizer.load_state_dict(opt_dict)
+
+    @property
+    def get_scaling(self):
+        return self.scaling_activation(self._scaling)
+    
+    @property
+    def get_rotation(self):
+        return self.rotation_activation(self._rotation)
+    
+    @property
+    def get_xyz(self):
+        return self._xyz
+    
+    @property
+    def get_features(self):
+        features_dc = self._features_dc
+        features_rest = self._features_rest
+        return torch.cat((features_dc, features_rest), dim=1)
+    
+    @property
+    def get_opacity(self):
+        return self.opacity_activation(self._opacity)
+    
+    def get_covariance(self, scaling_modifier = 1):
+        return self.covariance_activation(self.get_scaling, scaling_modifier, self._rotation)
+
+    def oneupSHdegree(self):
+        if self.active_sh_degree < self.max_sh_degree:
+            self.active_sh_degree += 1
+
+    def create_from_pcd(self, pcd : BasicPointCloud, spatial_lr_scale : float):
+        self.spatial_lr_scale = spatial_lr_scale
+        fused_point_cloud = torch.tensor(np.asarray(pcd.points)).float().cuda()
+        fused_color = RGB2SH(torch.tensor(np.asarray(pcd.colors)).float().cuda())
+        features = torch.zeros((fused_color.shape[0], 3, (self.max_sh_degree + 1) ** 2)).float().cuda()
+        features[:, :3, 0 ] = fused_color
+        features[:, 3:, 1:] = 0.0
+
+        print("Number of points at initialisation : ", fused_point_cloud.shape[0])
+
+        dist2 = torch.clamp_min(distCUDA2(torch.from_numpy(np.asarray(pcd.points)).float().cuda()), 0.0000001)
+        scales = torch.log(torch.sqrt(dist2))[...,None].repeat(1, 3)
+        rots = torch.zeros((fused_point_cloud.shape[0], 4), device="cuda")
+        rots[:, 0] = 1
+
+        opacities = inverse_sigmoid(0.1 * torch.ones((fused_point_cloud.shape[0], 1), dtype=torch.float, device="cuda"))
+
+        self._xyz = nn.Parameter(fused_point_cloud.requires_grad_(True))
+        self._features_dc = nn.Parameter(features[:,:,0:1].transpose(1, 2).contiguous().requires_grad_(True))
+        self._features_rest = nn.Parameter(features[:,:,1:].transpose(1, 2).contiguous().requires_grad_(True))
+        self._scaling = nn.Parameter(scales.requires_grad_(True)) 
+        self._rotation = nn.Parameter(rots.requires_grad_(True))
+        self._opacity = nn.Parameter(opacities.requires_grad_(True))
+        self.max_radii2D = torch.zeros((self.get_xyz.shape[0]), device="cuda")
+
+    def training_setup(self, training_args):
+        self.percent_dense = training_args.percent_dense
+        self.xyz_gradient_accum = torch.zeros((self.get_xyz.shape[0], 1), device="cuda")
+        self.denom = torch.zeros((self.get_xyz.shape[0], 1), device="cuda")
+
+        l = [
+            {'params': [self._xyz], 'lr': training_args.position_lr_init * self.spatial_lr_scale, "name": "xyz"},
+            {'params': [self._features_dc], 'lr': training_args.feature_lr, "name": "f_dc"},
+            {'params': [self._features_rest], 'lr': training_args.feature_lr / 20.0, "name": "f_rest"},
+            {'params': [self._opacity], 'lr': training_args.opacity_lr, "name": "opacity"},
+            {'params': [self._scaling], 'lr': training_args.scaling_lr, "name": "scaling"},
+            {'params': [self._rotation], 'lr': training_args.rotation_lr, "name": "rotation"}
+        ]
+
+        self.optimizer = torch.optim.Adam(l, lr=0.0, eps=1e-15)
+        self.xyz_scheduler_args = get_expon_lr_func(lr_init=training_args.position_lr_init*self.spatial_lr_scale,
+                                                    lr_final=training_args.position_lr_final*self.spatial_lr_scale,
+                                                    lr_delay_mult=training_args.position_lr_delay_mult,
+                                                    max_steps=training_args.position_lr_max_steps)
+
+    def update_learning_rate(self, iteration):
+        ''' Learning rate scheduling per step '''
+        for param_group in self.optimizer.param_groups:
+            if param_group["name"] == "xyz":
+                lr = self.xyz_scheduler_args(iteration)
+                param_group['lr'] = lr
+                return lr
+
+    def construct_list_of_attributes(self):
+        l = ['x', 'y', 'z', 'nx', 'ny', 'nz']
+        
+        for i in range(self._features_dc.shape[1]*self._features_dc.shape[2]):
+            l.append('f_dc_{}'.format(i))
+        for i in range(self._features_rest.shape[1]*self._features_rest.shape[2]):
+            l.append('f_rest_{}'.format(i))
+        l.append('opacity')
+        for i in range(self._scaling.shape[1]):
+            l.append('scale_{}'.format(i))
+        for i in range(self._rotation.shape[1]):
+            l.append('rot_{}'.format(i))
+        return l
+
+    def save_ply(self, path):
+        mkdir_p(os.path.dirname(path))
+
+        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()
+        f_rest = self._features_rest.detach().transpose(1, 2).flatten(start_dim=1).contiguous().cpu().numpy()
+        opacities = self._opacity.detach().cpu().numpy()
+        scale = self._scaling.detach().cpu().numpy()
+        rotation = self._rotation.detach().cpu().numpy()
+
+        dtype_full = [(attribute, 'f4') for attribute in self.construct_list_of_attributes()]
+
+        elements = np.empty(xyz.shape[0], dtype=dtype_full)
+        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)
+
+    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]
+
+        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_")]
+        extra_f_names = sorted(extra_f_names, key = lambda x: int(x.split('_')[-1]))
+        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])
+        
+        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_")]
+        scale_names = sorted(scale_names, key = lambda x: int(x.split('_')[-1]))
+        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")]
+        rot_names = sorted(rot_names, key = lambda x: int(x.split('_')[-1]))
+        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])
+
+        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))
+        self._features_rest = nn.Parameter(torch.tensor(features_extra, dtype=torch.float, device="cuda").transpose(1, 2).contiguous().requires_grad_(True))
+        self._opacity = nn.Parameter(torch.tensor(opacities, dtype=torch.float, device="cuda").requires_grad_(True))
+        self._scaling = nn.Parameter(torch.tensor(scales, dtype=torch.float, device="cuda").requires_grad_(True))
+        self._rotation = nn.Parameter(torch.tensor(rots, dtype=torch.float, device="cuda").requires_grad_(True))
+
+        self.active_sh_degree = self.max_sh_degree
+
+    def replace_tensor_to_optimizer(self, tensor, name):
+        optimizable_tensors = {}
+        for group in self.optimizer.param_groups:
+            if group["name"] == name:
+                stored_state = self.optimizer.state.get(group['params'][0], None)
+                stored_state["exp_avg"] = torch.zeros_like(tensor)
+                stored_state["exp_avg_sq"] = torch.zeros_like(tensor)
+
+                del self.optimizer.state[group['params'][0]]
+                group["params"][0] = nn.Parameter(tensor.requires_grad_(True))
+                self.optimizer.state[group['params'][0]] = stored_state
+
+                optimizable_tensors[group["name"]] = group["params"][0]
+        return optimizable_tensors
+
+    def _prune_optimizer(self, mask):
+        optimizable_tensors = {}
+        for group in self.optimizer.param_groups:
+            stored_state = self.optimizer.state.get(group['params'][0], None)
+            if stored_state is not None:
+                stored_state["exp_avg"] = stored_state["exp_avg"][mask]
+                stored_state["exp_avg_sq"] = stored_state["exp_avg_sq"][mask]
+
+                del self.optimizer.state[group['params'][0]]
+                group["params"][0] = nn.Parameter((group["params"][0][mask].requires_grad_(True)))
+                self.optimizer.state[group['params'][0]] = stored_state
+
+                optimizable_tensors[group["name"]] = group["params"][0]
+            else:
+                group["params"][0] = nn.Parameter(group["params"][0][mask].requires_grad_(True))
+                optimizable_tensors[group["name"]] = group["params"][0]
+        return optimizable_tensors
+
+    def prune_points(self, mask):
+        valid_points_mask = ~mask
+        optimizable_tensors = self._prune_optimizer(valid_points_mask)
+
+        self._xyz = optimizable_tensors["xyz"]
+        self._features_dc = optimizable_tensors["f_dc"]
+        self._features_rest = optimizable_tensors["f_rest"]
+        self._opacity = optimizable_tensors["opacity"]
+        self._scaling = optimizable_tensors["scaling"]
+        self._rotation = optimizable_tensors["rotation"]
+
+        self.xyz_gradient_accum = self.xyz_gradient_accum[valid_points_mask]
+
+        self.denom = self.denom[valid_points_mask]
+        self.max_radii2D = self.max_radii2D[valid_points_mask]
+
+    def cat_tensors_to_optimizer(self, tensors_dict):
+        optimizable_tensors = {}
+        for group in self.optimizer.param_groups:
+            assert len(group["params"]) == 1
+            extension_tensor = tensors_dict[group["name"]]
+            stored_state = self.optimizer.state.get(group['params'][0], None)
+            if stored_state is not None:
+
+                stored_state["exp_avg"] = torch.cat((stored_state["exp_avg"], torch.zeros_like(extension_tensor)), dim=0)
+                stored_state["exp_avg_sq"] = torch.cat((stored_state["exp_avg_sq"], torch.zeros_like(extension_tensor)), dim=0)
+
+                del self.optimizer.state[group['params'][0]]
+                group["params"][0] = nn.Parameter(torch.cat((group["params"][0], extension_tensor), dim=0).requires_grad_(True))
+                self.optimizer.state[group['params'][0]] = stored_state
+
+                optimizable_tensors[group["name"]] = group["params"][0]
+            else:
+                group["params"][0] = nn.Parameter(torch.cat((group["params"][0], extension_tensor), dim=0).requires_grad_(True))
+                optimizable_tensors[group["name"]] = group["params"][0]
+
+        return optimizable_tensors
+
+    def densification_postfix(self, new_xyz, new_features_dc, new_features_rest, new_opacities, new_scaling, new_rotation):
+        d = {"xyz": new_xyz,
+        "f_dc": new_features_dc,
+        "f_rest": new_features_rest,
+        "opacity": new_opacities,
+        "scaling" : new_scaling,
+        "rotation" : new_rotation}
+
+        optimizable_tensors = self.cat_tensors_to_optimizer(d)
+        self._xyz = optimizable_tensors["xyz"]
+        self._features_dc = optimizable_tensors["f_dc"]
+        self._features_rest = optimizable_tensors["f_rest"]
+        self._opacity = optimizable_tensors["opacity"]
+        self._scaling = optimizable_tensors["scaling"]
+        self._rotation = optimizable_tensors["rotation"]
+
+        self.xyz_gradient_accum = torch.zeros((self.get_xyz.shape[0], 1), device="cuda")
+        self.denom = torch.zeros((self.get_xyz.shape[0], 1), device="cuda")
+        self.max_radii2D = torch.zeros((self.get_xyz.shape[0]), device="cuda")
+
+    def densify_and_split(self, grads, grad_threshold, scene_extent, N=2):
+        n_init_points = self.get_xyz.shape[0]
+        
+        padded_grad = torch.zeros((n_init_points), device="cuda")
+        padded_grad[:grads.shape[0]] = grads.squeeze()
+        selected_pts_mask = torch.where(padded_grad >= grad_threshold, True, False)
+        selected_pts_mask = torch.logical_and(selected_pts_mask,
+                                              torch.max(self.get_scaling, dim=1).values > self.percent_dense*scene_extent)
+
+        stds = self.get_scaling[selected_pts_mask].repeat(N,1)
+        means =torch.zeros((stds.size(0), 3),device="cuda")
+        samples = torch.normal(mean=means, std=stds)
+        rots = build_rotation(self._rotation[selected_pts_mask]).repeat(N,1,1)
+        new_xyz = torch.bmm(rots, samples.unsqueeze(-1)).squeeze(-1) + self.get_xyz[selected_pts_mask].repeat(N, 1)
+        new_scaling = self.scaling_inverse_activation(self.get_scaling[selected_pts_mask].repeat(N,1) / (self.divide_ratio*N))
+        new_rotation = self._rotation[selected_pts_mask].repeat(N,1)
+        new_features_dc = self._features_dc[selected_pts_mask].repeat(N,1,1)
+        new_features_rest = self._features_rest[selected_pts_mask].repeat(N,1,1)
+        new_opacity = self._opacity[selected_pts_mask].repeat(N,1)
+
+        self.densification_postfix(new_xyz, new_features_dc, new_features_rest, new_opacity, new_scaling, new_rotation)
+
+        prune_filter = torch.cat((selected_pts_mask, torch.zeros(N * selected_pts_mask.sum(), device="cuda", dtype=bool)))
+        self.prune_points(prune_filter)
+
+    def densify_and_clone(self, grads, grad_threshold, scene_extent):
+        
+        selected_pts_mask = torch.where(torch.norm(grads, dim=-1) >= grad_threshold, True, False)
+        selected_pts_mask = torch.logical_and(selected_pts_mask,
+                                              torch.max(self.get_scaling, dim=1).values <= self.percent_dense*scene_extent)
+        
+        new_xyz = self._xyz[selected_pts_mask]
+        new_features_dc = self._features_dc[selected_pts_mask]
+        new_features_rest = self._features_rest[selected_pts_mask]
+        new_opacities = self._opacity[selected_pts_mask]
+        new_scaling = self._scaling[selected_pts_mask]
+        new_rotation = self._rotation[selected_pts_mask]
+
+        self.densification_postfix(new_xyz, new_features_dc, new_features_rest, new_opacities, new_scaling, new_rotation)
+
+    def densify_and_prune(self, max_grad, min_opacity, extent, max_screen_size):
+        grads = self.xyz_gradient_accum / self.denom
+        grads[grads.isnan()] = 0.0
+
+        self.densify_and_clone(grads, max_grad, extent)
+        self.densify_and_split(grads, max_grad, extent)
+
+        prune_mask = (self.get_opacity < min_opacity).squeeze()
+        if max_screen_size:
+            big_points_vs = self.max_radii2D > max_screen_size
+            big_points_ws = self.get_scaling.max(dim=1).values > 0.1 * extent
+            prune_mask = torch.logical_or(torch.logical_or(prune_mask, big_points_vs), big_points_ws)
+        self.prune_points(prune_mask)
+
+        torch.cuda.empty_cache()
+
+    def add_densification_stats(self, viewspace_point_tensor, update_filter):
+        self.xyz_gradient_accum[update_filter] += torch.norm(viewspace_point_tensor.grad[update_filter,:2], dim=-1, keepdim=True)
+        self.denom[update_filter] += 1

algorithms/main_3DGS_renderer.py

@@ -27,13 +27,12 @@ from diff_gaussian_rasterization import (
 )
 from simple_knn._C import distCUDA2
 
+from kiui.op import inverse_sigmoid
+
 from ..shared_utils.sh_utils import eval_sh, SH2RGB, RGB2SH
 from ..mesh_processer.mesh import Mesh, PointCloud
 from ..mesh_processer.mesh_utils import construct_list_of_gs_attributes, write_gs_ply, read_gs_ply
 
-def inverse_sigmoid(x):
-    return torch.log(x/(1-x))
-
 def get_expon_lr_func(
     lr_init, lr_final, lr_delay_steps=0, lr_delay_mult=1.0, max_steps=1000000
 ):

algorithms/nerf_marching_cubes_converter.py

@@ -162,6 +162,7 @@ class GSConverterNeRFMarchingCubes(nn.Module):
         ])
 
         print(f"[INFO] fitting nerf...")
+        imgs, alphas = [], []
         pbar = tqdm.trange(iters)
         for i in pbar:
 
@@ -172,6 +173,8 @@ class GSConverterNeRFMarchingCubes(nn.Module):
             pose = orbit_camera(ver, hor, rad)
             
             image_gt, alpha_gt = self.render_gs(pose)
+            imgs.append(image_gt.permute(1, 2, 0).unsqueeze(0))
+            alphas.append(alpha_gt.unsqueeze(0))
             image_pred, alpha_pred = self.render_nerf(pose)
 
             # if i % 200 == 0:
@@ -189,6 +192,8 @@ class GSConverterNeRFMarchingCubes(nn.Module):
             pbar.set_description(f"MSE = {loss_mse.item():.6f}")
         
         print(f"[INFO] finished fitting nerf!")
+        
+        return torch.cat(imgs, dim=0), torch.cat(alphas, dim=0)
     
     def render_mesh(self, pose):
 
@@ -265,6 +270,11 @@ class GSConverterNeRFMarchingCubes(nn.Module):
         self.v = torch.from_numpy(vertices).contiguous().float().to(self.device)
         self.f = torch.from_numpy(triangles).contiguous().int().to(self.device)
         self.deform = nn.Parameter(torch.zeros_like(self.v)).to(self.device)
+        
+        # Here coarse mesh already resemble the target shape closely
+        #mcube_mesh = Mesh(v=self.v, f=self.f, albedo=None, device=self.device)
+        #mcube_mesh.auto_normal()
+        #mcube_mesh.auto_uv()
 
         # fit mesh from gs
         lr_factor = 1