Spaces:

topabaem
/

CC-4DGS

Running on Zero

CC-4DGS / utils /compute_regularization.py

Deploy CC-4DGS viewer (part 3)

ea66a99 verified 1 day ago

10.7 kB

import torch
from utils.graphics_utils import apply_rotation, batch_quaternion_multiply

def _compute_smooth_regularization_71(self,
                                    xyz, 
                                    scales_emb, 
                                    rotations_emb, 
                                    opacity_emb, 
                                    shs_emb,
                                    time,
                                    out_pts, 
                                    out_scales, 
                                    out_rot, 
                                    out_opacity, 
                                    out_shs,
                                    smooth_cfg,
                                    orig_cache=None):
    """
    Compute smooth regularization loss on the deformation outputs.
    The idea is: for a small perturbation in the input (both spatial and temporal),
    the deformation network should produce similar outputs.
    This function re-computes the hidden representation and each branch's deformation
    for the original and perturbed inputs, then calculates an L2 loss for each branch,
    weighted by cfg, and returns the averaged loss.
    """
    # Retrieve perturbation magnitudes from config (or use defaults)
    spatial_eps = smooth_cfg.get("spatial_eps", 5e-3)
    temporal_eps = smooth_cfg.get("temporal_eps", 5e-3)
    # temporal_eps = 0.0
    
    choice = None
    if smooth_cfg['downsample_ratio'] < 1.0:
        # == Choose Random Points == 
        device = xyz.device
        N_all = xyz.shape[0]
        subN = int(max(1, N_all * smooth_cfg['downsample_ratio']))
        choice = torch.randperm(N_all, device=device)[:subN]

        # (a) Subset + detach
        xyz   = xyz[choice,:3].clone().detach()
        scales_emb    = scales_emb[choice, :3].clone().detach()
        rotations_emb   = rotations_emb[choice, :4].clone().detach()
        time     = time[choice,:1].clone().detach()   # [subN,1]
    
    # === Compute original deformation outputs ===
    if orig_cache is None:
        hidden_orig = self.query_time(xyz, scales_emb, rotations_emb, None, time)
        dx_orig = self.pos_deform(hidden_orig) if not self.args.no_dx else None
        ds_orig = self.scales_deform(hidden_orig) if not self.args.no_ds else None
        dr_orig = self.rotations_deform(hidden_orig) if not self.args.no_dr else None
        do_orig = self.opacity_deform(hidden_orig) if not self.args.no_do else None
        dshs_orig = (
            self.shs_deform(hidden_orig).reshape(xyz.shape[0], 16, 3)
            if not self.args.no_dshs else None
        )
    else:
        dx_orig = orig_cache.get("dx") if orig_cache else None
        ds_orig = orig_cache.get("ds") if orig_cache else None
        dr_orig = orig_cache.get("dr") if orig_cache else None
        do_orig = orig_cache.get("do") if orig_cache else None
        dshs_orig = orig_cache.get("dshs") if orig_cache else None
        if not self.args.no_dx and dx_orig is None:
            raise ValueError("orig_cache missing 'dx' for smooth regularization.")
        if not self.args.no_ds and ds_orig is None:
            raise ValueError("orig_cache missing 'ds' for smooth regularization.")
        if not self.args.no_dr and dr_orig is None:
            raise ValueError("orig_cache missing 'dr' for smooth regularization.")
        if not self.args.no_do and do_orig is None:
            raise ValueError("orig_cache missing 'do' for smooth regularization.")
        if not self.args.no_dshs and dshs_orig is None:
            raise ValueError("orig_cache missing 'dshs' for smooth regularization.")
        if choice is not None:
            if dx_orig is not None:
                dx_orig = dx_orig.index_select(0, choice)
            if ds_orig is not None:
                ds_orig = ds_orig.index_select(0, choice)
            if dr_orig is not None:
                dr_orig = dr_orig.index_select(0, choice)
            if do_orig is not None:
                do_orig = do_orig.index_select(0, choice)
            if dshs_orig is not None:
                dshs_orig = dshs_orig.index_select(0, choice)

    # === Generate perturbed inputs ===
    # Perturb the spatial (rays_pts_emb) and temporal (time_emb) inputs.
    xyz = self.normalize_aabb(xyz, self.grid.get_aabb)
    perturbed_rays = xyz + (torch.rand_like(xyz) * 2 - 1) * spatial_eps
    perturbed_time = time + (torch.rand_like(time) * 2 - 1) * temporal_eps

    # Compute the hidden representation for the perturbed inputs.
    hidden_pert = self.query_time(perturbed_rays, scales_emb, rotations_emb, None, perturbed_time, need_normalize=False)
    # hidden_pert = self.feature_out(hidden_pert)
    
    dx_pert = self.pos_deform(hidden_pert) if dx_orig is not None else None
    ds_pert = self.scales_deform(hidden_pert) if ds_orig is not None else None
    dr_pert = self.rotations_deform(hidden_pert) if dr_orig is not None else None
    do_pert = self.opacity_deform(hidden_pert) if do_orig is not None else None
    dshs_pert = (self.shs_deform(hidden_pert)
                 .reshape(xyz.shape[0], 16, 3)) if dshs_orig is not None else None

    # === Compute L2 differences per branch ===
    loss = 0.0
    count = 0
    if dx_orig is not None:
        loss_dx = torch.mean((dx_orig - dx_pert)**2)
        loss += smooth_cfg.get("w_dx", 1.0) * loss_dx
        count += 1
    if ds_orig is not None:
        loss_ds = torch.mean((ds_orig - ds_pert)**2)
        loss += smooth_cfg.get("w_ds", 1.0) * loss_ds
        count += 1
    if dr_orig is not None:
        loss_dr = torch.mean((dr_orig - dr_pert)**2)
        loss += smooth_cfg.get("w_dr", 1.0) * loss_dr
        count += 1
    if do_orig is not None:
        loss_do = torch.mean((do_orig - do_pert)**2)
        loss += smooth_cfg.get("w_do", 1.0) * loss_do
        count += 1
    if dshs_orig is not None:
        loss_dshs = torch.mean((dshs_orig - dshs_pert)**2)
        loss += smooth_cfg.get("w_dshs", 1.0) * loss_dshs
        count += 1
    assert count > 0, "No deformation branches enabled for smooth regularization."
    # Return the averaged loss across the branches used.
    return loss / (count if count > 0 else 1.0)

def _compute_smooth_regularization_72(self,
                                    xyz, 
                                    scales_emb, 
                                    rotations_emb, 
                                    opacity_emb, 
                                    shs_emb,
                                    time,
                                    out_pts, 
                                    out_scales, 
                                    out_rot, 
                                    out_opacity, 
                                    out_shs,
                                    smooth_cfg):
    """
    Compute smooth regularization loss on the deformation outputs.
    The idea is: for a small perturbation in the input (both spatial and temporal),
    the deformation network should produce similar outputs.
    This function re-computes the hidden representation and each branch's deformation
    for the original and perturbed inputs, then calculates an L2 loss for each branch,
    weighted by cfg, and returns the averaged loss.
    """
    # Retrieve perturbation magnitudes from config (or use defaults)
    spatial_eps = smooth_cfg.get("spatial_eps", 5e-3)
    temporal_eps = smooth_cfg.get("temporal_eps", 5e-3)
    # temporal_eps = 0.0
    
    if smooth_cfg['downsample_ratio'] < 1.0:
        # == Choose Random Points == 
        device = xyz.device
        N_all = xyz.shape[0]
        subN = int(max(1, N_all * smooth_cfg['downsample_ratio']))
        choice = torch.randperm(N_all, device=device)[:subN]

        # (a) Subset + detach
        xyz   = xyz[choice,:3].clone().detach()
        scales_emb    = scales_emb[choice, :3].clone().detach()
        rotations_emb   = rotations_emb[choice, :4].clone().detach()
        time     = time[choice,:1].clone().detach()   # [subN,1]
    
    # === Compute original deformation outputs ===
    # For consistency, re-run query_time and feature_out on the original input.
    hidden_orig = self.query_time(xyz, scales_emb, rotations_emb, None, time)
    
    dx_orig = self.pos_deform(hidden_orig) if not self.args.no_dx else None
    ds_orig = self.scales_deform(hidden_orig) if not self.args.no_ds else None
    dr_orig = self.rotations_deform(hidden_orig) if not self.args.no_dr else None
    do_orig = self.opacity_deform(hidden_orig) if not self.args.no_do else None
    dshs_orig = (self.shs_deform(hidden_orig)
                 .reshape(xyz.shape[0], 16, 3)) if not self.args.no_dshs else None

    # === Generate perturbed inputs ===
    # Perturb the spatial (rays_pts_emb) and temporal (time_emb) inputs.
    xyz = self.normalize_aabb(xyz, self.grid.get_aabb)
    perturbed_rays = xyz + (torch.rand_like(xyz) * 2 - 1) * spatial_eps
    perturbed_time = time + (torch.rand_like(time) * 2 - 1) * temporal_eps

    # Compute the hidden representation for the perturbed inputs.
    hidden_pert = self.query_time(perturbed_rays, scales_emb, rotations_emb, None, perturbed_time, need_normalize=False)
    # hidden_pert = self.feature_out(hidden_pert)
    
    dx_pert = self.pos_deform(hidden_pert) if dx_orig is not None else None
    ds_pert = self.scales_deform(hidden_pert) if ds_orig is not None else None
    dr_pert = self.rotations_deform(hidden_pert) if dr_orig is not None else None
    do_pert = self.opacity_deform(hidden_pert) if do_orig is not None else None
    dshs_pert = (self.shs_deform(hidden_pert)
                 .reshape(xyz.shape[0], 16, 3)) if dshs_orig is not None else None

    # === Compute L2 differences per branch ===
    loss = 0.0
    count = 0
    if dx_orig is not None:
        loss_dx = torch.mean((dx_orig - dx_pert)**2)
        loss += smooth_cfg.get("w_dx", 1.0) * loss_dx
        count += 1
    if ds_orig is not None:
        loss_ds = torch.mean((ds_orig - ds_pert)**2)
        loss += smooth_cfg.get("w_ds", 0.3) * loss_ds
        count += 1
    if dr_orig is not None:
        loss_dr = torch.mean((dr_orig - dr_pert)**2)
        loss += smooth_cfg.get("w_dr", 0.3) * loss_dr
        count += 1
    if do_orig is not None:
        loss_do = torch.mean((do_orig - do_pert)**2)
        loss += smooth_cfg.get("w_do", 0.1) * loss_do
        count += 1
    if dshs_orig is not None:
        loss_dshs = torch.mean((dshs_orig - dshs_pert)**2)
        loss += smooth_cfg.get("w_dshs", 0.1) * loss_dshs
        count += 1
    assert count > 0, "No deformation branches enabled for smooth regularization."
    # Return the averaged loss across the branches used.
    return loss / (count if count > 0 else 1.0)

	import torch
	from utils.graphics_utils import apply_rotation, batch_quaternion_multiply

	def _compute_smooth_regularization_71(self,
	xyz,
	scales_emb,
	rotations_emb,
	opacity_emb,
	shs_emb,
	time,
	out_pts,
	out_scales,
	out_rot,
	out_opacity,
	out_shs,
	smooth_cfg,
	orig_cache=None):
	"""
	Compute smooth regularization loss on the deformation outputs.
	The idea is: for a small perturbation in the input (both spatial and temporal),
	the deformation network should produce similar outputs.

	This function re-computes the hidden representation and each branch's deformation
	for the original and perturbed inputs, then calculates an L2 loss for each branch,
	weighted by cfg, and returns the averaged loss.
	"""
	# Retrieve perturbation magnitudes from config (or use defaults)
	spatial_eps = smooth_cfg.get("spatial_eps", 5e-3)
	temporal_eps = smooth_cfg.get("temporal_eps", 5e-3)
	# temporal_eps = 0.0

	choice = None
	if smooth_cfg['downsample_ratio'] < 1.0:
	# == Choose Random Points ==
	device = xyz.device
	N_all = xyz.shape[0]
	subN = int(max(1, N_all * smooth_cfg['downsample_ratio']))
	choice = torch.randperm(N_all, device=device)[:subN]

	# (a) Subset + detach
	xyz = xyz[choice,:3].clone().detach()
	scales_emb = scales_emb[choice, :3].clone().detach()
	rotations_emb = rotations_emb[choice, :4].clone().detach()
	time = time[choice,:1].clone().detach() # [subN,1]

	# === Compute original deformation outputs ===
	if orig_cache is None:
	hidden_orig = self.query_time(xyz, scales_emb, rotations_emb, None, time)
	dx_orig = self.pos_deform(hidden_orig) if not self.args.no_dx else None
	ds_orig = self.scales_deform(hidden_orig) if not self.args.no_ds else None
	dr_orig = self.rotations_deform(hidden_orig) if not self.args.no_dr else None
	do_orig = self.opacity_deform(hidden_orig) if not self.args.no_do else None
	dshs_orig = (
	self.shs_deform(hidden_orig).reshape(xyz.shape[0], 16, 3)
	if not self.args.no_dshs else None
	)
	else:
	dx_orig = orig_cache.get("dx") if orig_cache else None
	ds_orig = orig_cache.get("ds") if orig_cache else None
	dr_orig = orig_cache.get("dr") if orig_cache else None
	do_orig = orig_cache.get("do") if orig_cache else None
	dshs_orig = orig_cache.get("dshs") if orig_cache else None
	if not self.args.no_dx and dx_orig is None:
	raise ValueError("orig_cache missing 'dx' for smooth regularization.")
	if not self.args.no_ds and ds_orig is None:
	raise ValueError("orig_cache missing 'ds' for smooth regularization.")
	if not self.args.no_dr and dr_orig is None:
	raise ValueError("orig_cache missing 'dr' for smooth regularization.")
	if not self.args.no_do and do_orig is None:
	raise ValueError("orig_cache missing 'do' for smooth regularization.")
	if not self.args.no_dshs and dshs_orig is None:
	raise ValueError("orig_cache missing 'dshs' for smooth regularization.")
	if choice is not None:
	if dx_orig is not None:
	dx_orig = dx_orig.index_select(0, choice)
	if ds_orig is not None:
	ds_orig = ds_orig.index_select(0, choice)
	if dr_orig is not None:
	dr_orig = dr_orig.index_select(0, choice)
	if do_orig is not None:
	do_orig = do_orig.index_select(0, choice)
	if dshs_orig is not None:
	dshs_orig = dshs_orig.index_select(0, choice)

	# === Generate perturbed inputs ===
	# Perturb the spatial (rays_pts_emb) and temporal (time_emb) inputs.
	xyz = self.normalize_aabb(xyz, self.grid.get_aabb)
	perturbed_rays = xyz + (torch.rand_like(xyz) * 2 - 1) * spatial_eps
	perturbed_time = time + (torch.rand_like(time) * 2 - 1) * temporal_eps

	# Compute the hidden representation for the perturbed inputs.
	hidden_pert = self.query_time(perturbed_rays, scales_emb, rotations_emb, None, perturbed_time, need_normalize=False)
	# hidden_pert = self.feature_out(hidden_pert)

	dx_pert = self.pos_deform(hidden_pert) if dx_orig is not None else None
	ds_pert = self.scales_deform(hidden_pert) if ds_orig is not None else None
	dr_pert = self.rotations_deform(hidden_pert) if dr_orig is not None else None
	do_pert = self.opacity_deform(hidden_pert) if do_orig is not None else None
	dshs_pert = (self.shs_deform(hidden_pert)
	.reshape(xyz.shape[0], 16, 3)) if dshs_orig is not None else None

	# === Compute L2 differences per branch ===
	loss = 0.0
	count = 0
	if dx_orig is not None:
	loss_dx = torch.mean((dx_orig - dx_pert)**2)
	loss += smooth_cfg.get("w_dx", 1.0) * loss_dx
	count += 1
	if ds_orig is not None:
	loss_ds = torch.mean((ds_orig - ds_pert)**2)
	loss += smooth_cfg.get("w_ds", 1.0) * loss_ds
	count += 1
	if dr_orig is not None:
	loss_dr = torch.mean((dr_orig - dr_pert)**2)
	loss += smooth_cfg.get("w_dr", 1.0) * loss_dr
	count += 1
	if do_orig is not None:
	loss_do = torch.mean((do_orig - do_pert)**2)
	loss += smooth_cfg.get("w_do", 1.0) * loss_do
	count += 1
	if dshs_orig is not None:
	loss_dshs = torch.mean((dshs_orig - dshs_pert)**2)
	loss += smooth_cfg.get("w_dshs", 1.0) * loss_dshs
	count += 1
	assert count > 0, "No deformation branches enabled for smooth regularization."
	# Return the averaged loss across the branches used.
	return loss / (count if count > 0 else 1.0)

	def _compute_smooth_regularization_72(self,
	xyz,
	scales_emb,
	rotations_emb,
	opacity_emb,
	shs_emb,
	time,
	out_pts,
	out_scales,
	out_rot,
	out_opacity,
	out_shs,
	smooth_cfg):
	"""
	Compute smooth regularization loss on the deformation outputs.
	The idea is: for a small perturbation in the input (both spatial and temporal),
	the deformation network should produce similar outputs.

	This function re-computes the hidden representation and each branch's deformation
	for the original and perturbed inputs, then calculates an L2 loss for each branch,
	weighted by cfg, and returns the averaged loss.
	"""
	# Retrieve perturbation magnitudes from config (or use defaults)
	spatial_eps = smooth_cfg.get("spatial_eps", 5e-3)
	temporal_eps = smooth_cfg.get("temporal_eps", 5e-3)
	# temporal_eps = 0.0

	if smooth_cfg['downsample_ratio'] < 1.0:
	# == Choose Random Points ==
	device = xyz.device
	N_all = xyz.shape[0]
	subN = int(max(1, N_all * smooth_cfg['downsample_ratio']))
	choice = torch.randperm(N_all, device=device)[:subN]

	# (a) Subset + detach
	xyz = xyz[choice,:3].clone().detach()
	scales_emb = scales_emb[choice, :3].clone().detach()
	rotations_emb = rotations_emb[choice, :4].clone().detach()
	time = time[choice,:1].clone().detach() # [subN,1]

	# === Compute original deformation outputs ===
	# For consistency, re-run query_time and feature_out on the original input.
	hidden_orig = self.query_time(xyz, scales_emb, rotations_emb, None, time)

	dx_orig = self.pos_deform(hidden_orig) if not self.args.no_dx else None
	ds_orig = self.scales_deform(hidden_orig) if not self.args.no_ds else None
	dr_orig = self.rotations_deform(hidden_orig) if not self.args.no_dr else None
	do_orig = self.opacity_deform(hidden_orig) if not self.args.no_do else None
	dshs_orig = (self.shs_deform(hidden_orig)
	.reshape(xyz.shape[0], 16, 3)) if not self.args.no_dshs else None

	# === Generate perturbed inputs ===
	# Perturb the spatial (rays_pts_emb) and temporal (time_emb) inputs.
	xyz = self.normalize_aabb(xyz, self.grid.get_aabb)
	perturbed_rays = xyz + (torch.rand_like(xyz) * 2 - 1) * spatial_eps
	perturbed_time = time + (torch.rand_like(time) * 2 - 1) * temporal_eps

	# Compute the hidden representation for the perturbed inputs.
	hidden_pert = self.query_time(perturbed_rays, scales_emb, rotations_emb, None, perturbed_time, need_normalize=False)
	# hidden_pert = self.feature_out(hidden_pert)

	dx_pert = self.pos_deform(hidden_pert) if dx_orig is not None else None
	ds_pert = self.scales_deform(hidden_pert) if ds_orig is not None else None
	dr_pert = self.rotations_deform(hidden_pert) if dr_orig is not None else None
	do_pert = self.opacity_deform(hidden_pert) if do_orig is not None else None
	dshs_pert = (self.shs_deform(hidden_pert)
	.reshape(xyz.shape[0], 16, 3)) if dshs_orig is not None else None

	# === Compute L2 differences per branch ===
	loss = 0.0
	count = 0
	if dx_orig is not None:
	loss_dx = torch.mean((dx_orig - dx_pert)**2)
	loss += smooth_cfg.get("w_dx", 1.0) * loss_dx
	count += 1
	if ds_orig is not None:
	loss_ds = torch.mean((ds_orig - ds_pert)**2)
	loss += smooth_cfg.get("w_ds", 0.3) * loss_ds
	count += 1
	if dr_orig is not None:
	loss_dr = torch.mean((dr_orig - dr_pert)**2)
	loss += smooth_cfg.get("w_dr", 0.3) * loss_dr
	count += 1
	if do_orig is not None:
	loss_do = torch.mean((do_orig - do_pert)**2)
	loss += smooth_cfg.get("w_do", 0.1) * loss_do
	count += 1
	if dshs_orig is not None:
	loss_dshs = torch.mean((dshs_orig - dshs_pert)**2)
	loss += smooth_cfg.get("w_dshs", 0.1) * loss_dshs
	count += 1
	assert count > 0, "No deformation branches enabled for smooth regularization."
	# Return the averaged loss across the branches used.
	return loss / (count if count > 0 else 1.0)