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

Update Gaussian Splatting's Training Schedule & Default Parameters

Browse Source 
Also update Multi-View-Images_to_3DGS_to_3DMesh(DMTet and DiffRast) workflow
This commit is contained in:
MrForExample
2024-02-20 22:46:28 +01:00
parent fcdda1792d
commit a4c20e95ec
5 changed files with 465 additions and 1691 deletions
Download Patch File Download Diff File Expand all files Collapse all files
_Example_Workflows/Multi-View-Images_to_3DGS_to_3DMesh(DMTet and DiffRast).json
+421 -418
View File
File diff suppressed because it is too large Load Diff
_Example_Workflows/Triplane_Gaussian_Transformers_to_3DGS(DMTet and DiffRast).json
-1229
View File
File diff suppressed because it is too large Load Diff
algorithms/main_3DGS.py
+23 -23
View File
@@ -15,36 +15,34 @@ from ..shared_utils.image_utils import prepare_torch_img
class GSParams:
def __init__(self,
training_iterations=3000,
training_iterations=30_000,
batch_size=1,
loss_scale=10000,
lambda_ssim=0.2,
lambda_alpha=3,
lambda_offset=0,
lambda_offset_opacity=0,
invert_bg_prob = 0.5,
feature_lr=0.01,
feature_lr=0.0025,
opacity_lr=0.05,
scaling_lr=0.005,
rotation_lr=0.005,
position_lr_init=0.001,
position_lr_final=0.00002,
position_lr_delay_mult=0.02,
position_lr_max_steps=500,
num_pts=5000,
rotation_lr=0.001,
position_lr_init=0.00016,
position_lr_final=0.0000016,
position_lr_delay_mult=0.01,
position_lr_max_steps=30_000,
num_pts=10_000,
K=3,
percent_dense=0.01,
density_start_iter=100,
density_end_iter=100000,
density_start_iter=500,
density_end_iter=15_000,
densification_interval=100,
opacity_reset_interval=700,
densify_grad_threshold=0.01,
sh_degree=0):
opacity_reset_interval=3000,
densify_grad_threshold=0.0002,
sh_degree=3):
# training params
self.training_iterations = training_iterations
self.batch_size = batch_size
self.loss_scale = loss_scale
self.lambda_ssim = lambda_ssim
self.lambda_alpha = lambda_alpha
self.lambda_offset = lambda_offset
@@ -212,25 +210,27 @@ class GaussianSplatting:
masks_batch_torch = torch.cat(masks_batch, dim=0)
ref_masks_batch_torch = torch.cat(ref_masks_batch, dim=0)
#loss_scaler = self.gs_params.loss_scale * step_ratio
# rgb loss
loss += (1 - self.gs_params.lambda_ssim) * self.gs_params.loss_scale * step_ratio * F.mse_loss(masked_rendered_img_batch_torch, masked_ref_img_batch_torch)
loss += (1 - self.gs_params.lambda_ssim) * F.l1_loss(masked_rendered_img_batch_torch, masked_ref_img_batch_torch)
# alpha loss
loss += self.gs_params.lambda_alpha * self.gs_params.loss_scale * step_ratio * F.mse_loss(masks_batch_torch, ref_masks_batch_torch)
loss += self.gs_params.lambda_alpha * F.mse_loss(masks_batch_torch, ref_masks_batch_torch)
# D-SSIM loss
# [1, 3, H, W] in [0, 1]
#loss += self.lambda_ssim * (1 - self.ssim_loss(X, Y))
loss += self.gs_params.lambda_ssim * self.gs_params.loss_scale * step_ratio * (1 - self.ms_ssim_loss(masked_ref_img_batch_torch, masked_rendered_img_batch_torch))
loss += self.gs_params.lambda_ssim * (1 - self.ms_ssim_loss(masked_ref_img_batch_torch, masked_rendered_img_batch_torch))
if self.gs_params.lambda_offset > 0:
# Reference offset loss
offset_norm = self.renderer.gaussians.get_xyz_offset.norm(dim=-1, keepdim=True)
loss += self.gs_params.lambda_offset * self.gs_params.loss_scale * step_ratio * torch.mean(offset_norm)
loss += self.gs_params.lambda_offset * torch.mean(offset_norm)
if self.gs_params.lambda_offset_opacity > 0:
# Alpha penalty loss
loss += self.gs_params.lambda_offset_opacity * self.gs_params.loss_scale * step_ratio * torch.mean(offset_norm.detach() * self.renderer.gaussians.get_opacity)
loss += self.gs_params.lambda_offset_opacity * torch.mean(offset_norm.detach() * self.renderer.gaussians.get_opacity)
# optimize step
@@ -245,12 +245,12 @@ class GaussianSplatting:
self.renderer.gaussians.add_densification_stats(viewspace_point_tensor, visibility_filter)
if step % self.gs_params.densification_interval == 0:
#self.renderer.gaussians.densify_and_prune(self.gs_params.densify_grad_threshold, min_opacity=0.01, extent=4, max_screen_size=1)
self.renderer.gaussians.densify_and_prune(self.gs_params.densify_grad_threshold, min_opacity=0.005, extent=4, max_screen_size=1)
#self.renderer.gaussians.densify_and_prune_by_compatness(self.gs_params.K, min_opacity=0.01, extent=4, max_screen_size=1)
#self.renderer.gaussians.densify_by_clone_and_split(self.gs_params.densify_grad_threshold, extent=4)
self.renderer.gaussians.densify_by_compatness(self.gs_params.K)
self.renderer.gaussians.prune(min_opacity=0.01, extent=4, max_screen_size=1, max_offset=0.1)
#self.renderer.gaussians.densify_by_compatness(self.gs_params.K)
#self.renderer.gaussians.prune(min_opacity=0.01, extent=4, max_screen_size=1, max_offset=0.1)
if step % self.gs_params.opacity_reset_interval == 0:
self.renderer.gaussians.reset_opacity()
nodes.py
+18 -19
View File
@@ -721,31 +721,30 @@ class Gaussian_Splatting:
"reference_masks": ("MASK",),
"reference_orbit_camera_poses": ("ORBIT_CAMPOSES",), # (orbit radius, elevation, azimuth, orbit center X, orbit center Y, orbit center Z)
"reference_orbit_camera_fovy": ("FLOAT", {"default": 49.1, "min": 0.0, "max": 180.0, "step": 0.1}),
"training_iterations": ("INT", {"default": 1000, "min": 1, "max": 100000}),
"batch_size": ("INT", {"default": 3, "min": 1, "max": 0xffffffffffffffff}),
"loss_value_scale": ("FLOAT", {"default": 10000.0, "min": 1.0}),
"training_iterations": ("INT", {"default": 30_000, "min": 1, "max": 0xffffffffffffffff}),
"batch_size": ("INT", {"default": 1, "min": 1, "max": 0xffffffffffffffff}),
"ms_ssim_loss_weight": ("FLOAT", {"default": 0.2, "min": 0.0, "max": 1.0, }),
"alpha_loss_weight": ("FLOAT", {"default": 3, "min": 0.0, }),
"offset_loss_weight": ("FLOAT", {"default": 0.0, "min": 0.0, }),
"offset_opacity_loss_weight": ("FLOAT", {"default": 0.0, "min": 0.0, }),
"invert_background_probability": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.1}),
"feature_learning_rate": ("FLOAT", {"default": 0.01, "min": 0.00001, "step": 0.00001}),
"opacity_learning_rate": ("FLOAT", {"default": 0.05, "min": 0.00001, "step": 0.00001}),
"scaling_learning_rate": ("FLOAT", {"default": 0.005, "min": 0.00001, "step": 0.00001}),
"rotation_learning_rate": ("FLOAT", {"default": 0.005, "min": 0.00001, "step": 0.00001}),
"position_learning_rate_init": ("FLOAT", {"default": 0.001, "min": 0.00001, "step": 0.00001}),
"position_learning_rate_final": ("FLOAT", {"default": 0.00002, "min": 0.00001, "step": 0.00001}),
"position_learning_rate_delay_mult": ("FLOAT", {"default": 0.02, "min": 0.00001, "step": 0.00001}),
"position_learning_rate_max_steps": ("INT", {"default": 500, "min": 1}),
"initial_gaussians_num": ("INT", {"default": 5000, "min": 1}),
"K_nearest_neighbors": ("INT", {"default": 3, "min": 1}),
"feature_learning_rate": ("FLOAT", {"default": 0.0025, "min": 0.000001, "step": 0.000001}),
"opacity_learning_rate": ("FLOAT", {"default": 0.05, "min": 0.000001, "step": 0.000001}),
"scaling_learning_rate": ("FLOAT", {"default": 0.005, "min": 0.000001, "step": 0.000001}),
"rotation_learning_rate": ("FLOAT", {"default": 0.001, "min": 0.000001, "step": 0.000001}),
"position_learning_rate_init": ("FLOAT", {"default": 0.00016, "min": 0.000001, "step": 0.000001}),
"position_learning_rate_final": ("FLOAT", {"default": 0.0000016, "min": 0.0000001, "step": 0.0000001}),
"position_learning_rate_delay_mult": ("FLOAT", {"default": 0.01, "min": 0.000001, "step": 0.000001}),
"position_learning_rate_max_steps": ("INT", {"default": 30_000, "min": 1, "max": 0xffffffffffffffff}),
"initial_gaussians_num": ("INT", {"default": 10_000, "min": 1, "max": 0xffffffffffffffff}),
"K_nearest_neighbors": ("INT", {"default": 3, "min": 1, "max": 0xffffffffffffffff}),
"percent_dense": ("FLOAT", {"default": 0.01, "min": 0.00001, "step": 0.00001}),
"density_start_iterations": ("INT", {"default": 100, "min": 0}),
"density_end_iterations": ("INT", {"default": 100000, "min": 0}),
"densification_interval": ("INT", {"default": 100, "min": 1}),
"opacity_reset_interval": ("INT", {"default": 700, "min": 1}),
"densify_grad_threshold": ("FLOAT", {"default": 0.01, "min": 0.00001, "step": 0.00001}),
"gaussian_sh_degree": ("INT", {"default": 0, "min": 0}),
"density_start_iterations": ("INT", {"default": 500, "min": 0, "max": 0xffffffffffffffff}),
"density_end_iterations": ("INT", {"default": 15_000, "min": 0, "max": 0xffffffffffffffff}),
"densification_interval": ("INT", {"default": 100, "min": 1, "max": 0xffffffffffffffff}),
"opacity_reset_interval": ("INT", {"default": 3000, "min": 1, "max": 0xffffffffffffffff}),
"densify_grad_threshold": ("FLOAT", {"default": 0.0002, "min": 0.00001, "step": 0.00001}),
"gaussian_sh_degree": ("INT", {"default": 3, "min": 0}),
},
"optional": {
web/js/threeVisualizer.js
+3 -2
View File
@@ -119,8 +119,9 @@ async function main(filepath="") {
loader.load( url, function ( gltf ) {
const model = gltf.scene;
model.position.set( 1, 1, 0 );
model.scale.set( 0.01, 0.01, 0.01 );
//model.position.set( 1, 1, 0 );
//model.scale.set( 0.01, 0.01, 0.01 );
model.scale.set( 3, 3, 3 );
scene.add( model );