232 lines
9.9 KiB
Python
232 lines
9.9 KiB
Python
import random
|
|
import tqdm
|
|
|
|
import torch
|
|
import torch.nn.functional as F
|
|
from pytorch_msssim import SSIM, MS_SSIM
|
|
|
|
import comfy.utils
|
|
|
|
from .main_3DGS_renderer import GaussianSplattingRenderer
|
|
|
|
from shared_utils.camera_utils import BaseCameraController, MiniCam, calculate_fovX, get_projection_matrix
|
|
from shared_utils.image_utils import prepare_torch_img
|
|
|
|
class GSParams:
|
|
def __init__(
|
|
self,
|
|
training_iterations=30_000,
|
|
batch_size=1,
|
|
lambda_ssim=0.2,
|
|
lambda_alpha=3,
|
|
lambda_offset=0,
|
|
lambda_offset_opacity=0,
|
|
invert_bg_prob = 0.5,
|
|
feature_lr=0.0025,
|
|
opacity_lr=0.05,
|
|
scaling_lr=0.005,
|
|
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=500,
|
|
density_end_iter=15_000,
|
|
densification_interval=100,
|
|
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.lambda_ssim = lambda_ssim
|
|
self.lambda_alpha = lambda_alpha
|
|
self.lambda_offset = lambda_offset
|
|
self.lambda_offset_opacity = lambda_offset_opacity
|
|
self.invert_bg_prob = invert_bg_prob
|
|
|
|
# learning params
|
|
self.feature_lr = feature_lr
|
|
self.opacity_lr = opacity_lr
|
|
self.scaling_lr = scaling_lr
|
|
self.rotation_lr = rotation_lr
|
|
self.position_lr_init = position_lr_init
|
|
self.position_lr_final = position_lr_final
|
|
self.position_lr_delay_mult = position_lr_delay_mult
|
|
self.position_lr_max_steps = position_lr_max_steps
|
|
|
|
# densify and prune params
|
|
self.num_pts = num_pts
|
|
self.K = K
|
|
self.percent_dense = percent_dense
|
|
self.density_start_iter=density_start_iter
|
|
self.density_end_iter=density_end_iter
|
|
self.densification_interval=densification_interval
|
|
self.opacity_reset_interval=opacity_reset_interval
|
|
self.densify_grad_threshold=densify_grad_threshold
|
|
|
|
# other gaussian params
|
|
self.sh_degree = sh_degree
|
|
|
|
class GaussianSplattingCameraController(BaseCameraController):
|
|
def get_render_result(self, render_pose, bg_color, **kwargs):
|
|
render_cam = MiniCam(render_pose, self.cam.W, self.cam.H, self.cam.fovy, self.cam.fovx, self.cam.near, self.cam.far, self.projection_matrix)
|
|
return self.renderer.render(render_cam, bg_color=bg_color, **kwargs)
|
|
|
|
def post_init(self):
|
|
self.projection_matrix = get_projection_matrix(self.cam.near, self.cam.far, self.cam.fovx, self.cam.fovy).transpose(0, 1).cuda()
|
|
|
|
class GaussianSplatting3D:
|
|
|
|
def __init__(self, gs_params=None, init_input=None, device='cuda'):
|
|
self.device = torch.device(device)
|
|
|
|
# prepare renderer for optimization
|
|
if gs_params is None:
|
|
gs_params = GSParams()
|
|
|
|
self.renderer = GaussianSplattingRenderer(sh_degree=gs_params.sh_degree)
|
|
self.renderer.initialize(init_input, num_pts=gs_params.num_pts)
|
|
|
|
# setup training
|
|
self.renderer.gaussians.training_setup(gs_params)
|
|
# do not do progressive sh-level
|
|
self.renderer.gaussians.active_sh_degree = self.renderer.gaussians.max_sh_degree
|
|
self.optimizer = self.renderer.gaussians.optimizer
|
|
#self.ssim_loss = SSIM(data_range=1, size_average=True, channel=3)
|
|
self.ms_ssim_loss = MS_SSIM(data_range=1, size_average=True, channel=3)
|
|
|
|
self.gs_params = gs_params
|
|
|
|
def prepare_training(self, reference_images, reference_masks, reference_orbit_camera_poses, reference_orbit_camera_fovy):
|
|
self.ref_imgs_num = len(reference_images)
|
|
|
|
self.ref_size_H = reference_images[0].shape[0]
|
|
self.ref_size_W = reference_images[0].shape[1]
|
|
|
|
# default camera settings
|
|
self.cam_controller = GaussianSplattingCameraController(
|
|
self.renderer, self.ref_size_W, self.ref_size_H, reference_orbit_camera_fovy, self.gs_params.invert_bg_prob, None, self.device
|
|
)
|
|
|
|
self.all_ref_cam_poses = reference_orbit_camera_poses
|
|
|
|
# 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)
|
|
self.ref_masks_torch = torch.cat(ref_masks_torch_list, dim=0)
|
|
|
|
def training(self):
|
|
starter = torch.cuda.Event(enable_timing=True)
|
|
ender = torch.cuda.Event(enable_timing=True)
|
|
starter.record()
|
|
|
|
ref_imgs_masked = []
|
|
for i in range(self.ref_imgs_num):
|
|
ref_imgs_masked.append((self.ref_imgs_torch[i] * self.ref_masks_torch[i]).unsqueeze(0))
|
|
|
|
ref_imgs_num_minus_1 = self.ref_imgs_num-1
|
|
|
|
comfy_pbar = comfy.utils.ProgressBar(self.gs_params.training_iterations)
|
|
|
|
for step in tqdm.trange(self.gs_params.training_iterations):
|
|
|
|
### calculate loss between reference and rendered image from known view
|
|
loss = 0
|
|
masked_rendered_img_batch = []
|
|
masked_ref_img_batch = []
|
|
masks_batch = []
|
|
ref_masks_batch = []
|
|
|
|
step_ratio = min(1, step / self.gs_params.training_iterations)
|
|
|
|
# update lr
|
|
self.renderer.gaussians.update_learning_rate(step)
|
|
|
|
loss = 0
|
|
|
|
for _ in range(self.gs_params.batch_size):
|
|
### calculate loss between reference and rendered image from known view
|
|
|
|
i = random.randint(0, ref_imgs_num_minus_1)
|
|
|
|
out = self.cam_controller.render_at_pose(self.all_ref_cam_poses[i])
|
|
|
|
image = out["image"] # [3, H, W] in [0, 1]
|
|
mask = out["alpha"] # [1, H, W] in [0, 1]
|
|
|
|
ref_mask = self.ref_masks_torch[i]
|
|
image_masked = (image * ref_mask).unsqueeze(0)
|
|
|
|
masked_rendered_img_batch.append(image_masked)
|
|
masked_ref_img_batch.append(ref_imgs_masked[i])
|
|
masks_batch.append(mask.unsqueeze(0))
|
|
ref_masks_batch.append(ref_mask.unsqueeze(0))
|
|
|
|
masked_rendered_img_batch_torch = torch.cat(masked_rendered_img_batch, dim=0)
|
|
masked_ref_img_batch_torch = torch.cat(masked_ref_img_batch, dim=0)
|
|
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) * F.l1_loss(masked_rendered_img_batch_torch, masked_ref_img_batch_torch)
|
|
|
|
# alpha loss
|
|
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 * (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 * torch.mean(offset_norm)
|
|
|
|
if self.gs_params.lambda_offset_opacity > 0:
|
|
# Alpha penalty loss
|
|
loss += self.gs_params.lambda_offset_opacity * torch.mean(offset_norm.detach() * self.renderer.gaussians.get_opacity)
|
|
|
|
|
|
# optimize step
|
|
loss.backward()
|
|
self.optimizer.step()
|
|
self.optimizer.zero_grad()
|
|
|
|
# densify and prune
|
|
if step >= self.gs_params.density_start_iter and step <= self.gs_params.density_end_iter:
|
|
viewspace_point_tensor, visibility_filter, radii = out["viewspace_points"], out["visibility_filter"], out["radii"]
|
|
self.renderer.gaussians.max_radii2D[visibility_filter] = torch.max(self.renderer.gaussians.max_radii2D[visibility_filter], radii[visibility_filter])
|
|
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.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)
|
|
|
|
if step % self.gs_params.opacity_reset_interval == 0:
|
|
self.renderer.gaussians.reset_opacity()
|
|
|
|
comfy_pbar.update_absolute(step + 1)
|
|
|
|
ender.record()
|
|
torch.cuda.synchronize()
|
|
#t = starter.elapsed_time(ender)
|
|
|
|
self.need_update = True |