""" from EMER-NeRF """ from torch import Tensor import torch.nn as nn from typing import Optional, Tuple import torch import torch.nn.functional as F from .encodings import SinusoidalEncoder from typing import Dict, Optional, Union import math from scene.gaussian_model import GaussianModel from utils.general_utils import sample_on_aabb_surface from utils.sh_utils import SH2RGB import numpy as np from utils.general_utils import inverse_sigmoid, get_expon_lr_func, build_rotation 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 import matplotlib.pyplot as plt def min_max(x, name='x', print_info=True): print(f"{name} range : {x.min().item():.7f} ~ {x.max().item():.7f}") # 如果x 是一维 if x.dim() == 1: x = x.unsqueeze(1) out = torch.cat((x.min(dim=0).values, x.max(dim=0).values), dim=0) return out def get_rays( x: Tensor, y: Tensor, c2w: Tensor, intrinsic: Tensor ) -> Tuple[Tensor, Tensor, Tensor]: """ Args: x: the horizontal coordinates of the pixels, shape: (num_rays,) y: the vertical coordinates of the pixels, shape: (num_rays,) c2w: the camera-to-world matrices, shape: (num_cams, 4, 4) intrinsic: the camera intrinsic matrices, shape: (num_cams, 3, 3) Returns: origins: the ray origins, shape: (num_rays, 3) viewdirs: the ray directions, shape: (num_rays, 3) direction_norm: the norm of the ray directions, shape: (num_rays, 1) """ if len(intrinsic.shape) == 2: intrinsic = intrinsic[None, :, :] if len(c2w.shape) == 2: c2w = c2w[None, :, :] camera_dirs = torch.nn.functional.pad( torch.stack( [ (x - intrinsic[:, 0, 2] + 0.5) / intrinsic[:, 0, 0], (y - intrinsic[:, 1, 2] + 0.5) / intrinsic[:, 1, 1], ], dim=-1, ), (0, 1), value=1.0, ) # [num_rays, 3] # rotate the camera rays w.r.t. the camera pose directions = (camera_dirs[:, None, :] * c2w[:, :3, :3]).sum(dim=-1) origins = torch.broadcast_to(c2w[:, :3, -1], directions.shape) # TODO: not sure if we still need direction_norm direction_norm = torch.linalg.norm(directions, dim=-1, keepdims=True) # normalize the ray directions viewdirs = directions / (direction_norm + 1e-8) return origins, viewdirs, direction_norm def get_directions( x: Tensor, y: Tensor, c2w: Tensor, intrinsic: Tensor ) -> Tuple[Tensor, Tensor, Tensor]: """ Args: x: the horizontal coordinates of the pixels, shape: (num_rays,) y: the vertical coordinates of the pixels, shape: (num_rays,) c2w: the camera-to-world matrices, shape: (num_cams, 4, 4) intrinsic: the camera intrinsic matrices, shape: (num_cams, 3, 3) Returns: origins: the ray origins, shape: (num_rays, 3) viewdirs: the ray directions, shape: (num_rays, 3) direction_norm: the norm of the ray directions, shape: (num_rays, 1) """ if len(intrinsic.shape) == 2: intrinsic = intrinsic[None, :, :] if len(c2w.shape) == 2: c2w = c2w[None, :, :] camera_dirs = torch.nn.functional.pad( torch.stack( [ (x - intrinsic[:, 0, 2] + 0.5) / intrinsic[:, 0, 0], (y - intrinsic[:, 1, 2] + 0.5) / intrinsic[:, 1, 1], ], dim=-1, ), (0, 1), value=1.0, ) # [num_rays, 3] # rotate the camera rays w.r.t. the camera pose directions = (camera_dirs[:, None, :] * c2w[:, :3, :3]).sum(dim=-1) # TODO: not sure if we still need direction_norm direction_norm = torch.linalg.norm(directions, dim=-1, keepdims=True) # normalize the ray directions viewdirs = directions / (direction_norm + 1e-8) return viewdirs def ray_box_intersect(rays_o: torch.Tensor, rays_d: torch.Tensor, r: Union[float, torch.Tensor]) -> torch.Tensor: """ Calculate intersections of each rays and each boxes Args: rays_o (torch.Tensor): [B,3] rays_d (torch.Tensor): [B,3] r (Union[float, torch.Tensor]): [(B,)P] half side-lengths of boxes Returns: torch.Tensor: [B,P] depths of intersections along rays """ # Expand to (B, 1, 3) rays_o = rays_o.unsqueeze(1) rays_d = rays_d.unsqueeze(1) # Expand to ([B, ]P, 1) r = r[..., None] # t_min, t_max: (B, P, 3) t_min = (-r - rays_o) / rays_d t_max = (r - rays_o) / rays_d # t_near, t_far: (B, P) t_near = torch.minimum(t_min, t_max).max(dim=-1).values t_far = torch.maximum(t_min, t_max).min(dim=-1).values # Check if rays are inside boxes and boxes are in front of the ray origin mask_intersect = (t_far > t_near) & (t_far > 0) t_far[~mask_intersect] = math.nan return t_far class MLP(nn.Module): """A simple MLP with skip connections.""" def __init__( self, in_dims: int, out_dims: int, num_layers: int = 3, hidden_dims: Optional[int] = 256, skip_connections: Optional[Tuple[int]] = [0], ) -> None: super().__init__() self.in_dims = in_dims self.hidden_dims = hidden_dims self.n_output_dims = out_dims self.num_layers = num_layers self.skip_connections = skip_connections layers = [] if self.num_layers == 1: layers.append(nn.Linear(in_dims, out_dims)) else: for i in range(self.num_layers - 1): if i == 0: layers.append(nn.Linear(in_dims, hidden_dims)) elif i in skip_connections: layers.append(nn.Linear(in_dims + hidden_dims, hidden_dims)) else: layers.append(nn.Linear(hidden_dims, hidden_dims)) layers.append(nn.Linear(hidden_dims, out_dims)) self.layers = nn.ModuleList(layers) def forward(self, x: Tensor) -> Tensor: input = x for i, layer in enumerate(self.layers): if i in self.skip_connections: x = torch.cat([x, input], -1) x = layer(x) if i < len(self.layers) - 1: x = nn.functional.relu(x) return x class ShellModel(nn.Module): """ find_intersection + grid_interpolation --> sky_feat """ def __init__(self, model_cfg: dict): super().__init__() pass class SkyModel(nn.Module): """ dir_encoding + MLP --> sky_feat """ def __init__(self, model_cfg: dict = {"mlp_width": 256,}): super().__init__() self.model_cfg = model_cfg self.direction_encoding = SinusoidalEncoder( n_input_dims=3, min_deg=0, max_deg=4 ) self.sky_head = MLP( in_dims=self.direction_encoding.n_output_dims, out_dims=3, num_layers=3, hidden_dims=model_cfg['mlp_width'], skip_connections=[1], ) def forward(self, directions): dd = self.direction_encoding(directions) sky_feat = self.sky_head(dd) #rgb_sky = F.sigmoid(sky_feat) return sky_feat class BackgroundGaussianModel(GaussianModel): """ Gaussian model for background - 输入 aabb - 增加了 xyz 的映射 , 将坐标 固定在指定范围内: aabb 之外 - """ def print_gs_info(self, hist_bins='auto', hist_density=True, hist_img_path=None): pass print('------------------ GS info ------------------') print("Number of points : ", self.get_xyz.shape[0]) # -------------------------- get data ----------------------------------- # scale scale_before_activation = self._scaling scale = self.get_scaling # rotation: 四元数 用 norm 来归一化 rotation_before_activation = self._rotation rotation = self.get_rotation # opacity opacity_before_activation = self._opacity opacity = self.get_opacity # xyz xyz = self.get_xyz # -------------------------- static-info ----------------------------------- # ----------------- min-max range ----------------- #scale_before_activation_range = min_max(scale_before_activation, 'scale_before_activation') scale_x_before_activation_range = min_max(scale_before_activation[:, 0:1], 'scale_x_before_activation') scale_y_before_activation_range = min_max(scale_before_activation[:, 1:2], 'scale_y_before_activation') scale_z_before_activation_range = min_max(scale_before_activation[:, 2:3], 'scale_z_before_activation') #scale_range = min_max(scale, 'scale') scale_x_range = min_max(scale[:, 0:1], 'scale_x') scale_y_range = min_max(scale[:, 1:2], 'scale_y') scale_z_range = min_max(scale[:, 2:3], 'scale_z') # scale ratio # 长轴 与 短轴之比 min_axis = torch.min(scale, dim=1)[0] max_axis = torch.max(scale, dim=1)[0] scale_ratio = max_axis / min_axis scale_ratio_range = min_max(scale_ratio, 'scale_ratio') # rot 暂不考虑 #rotation_before_activation_range = min_max(rotation_before_activation, 'rotation_before_activation') #rotation_range = min_max(rotation, 'rotation') opacity_before_activation_range = min_max(opacity_before_activation, 'opacity_before_activation') opacity_range = min_max(opacity, 'opacity') #xyz_range = min_max(xyz, 'xyz') x_range = min_max(xyz[:, 0:1], 'x') y_range = min_max(xyz[:, 1:2], 'y') z_range = min_max(xyz[:, 2:3], 'z') # ----------------- hist ----------------- #scale_before_activation_hist = np.histogram(scale_before_activation.cpu().numpy(), bins=hist_bins, density=hist_density) scale_x_before_activation_hist = np.histogram(scale_before_activation[:, 0:1].cpu().numpy(), bins=hist_bins, density=hist_density) scale_y_before_activation_hist = np.histogram(scale_before_activation[:, 1:2].cpu().numpy(), bins=hist_bins, density=hist_density) scale_z_before_activation_hist = np.histogram(scale_before_activation[:, 2:3].cpu().numpy(), bins=hist_bins, density=hist_density) #scale_hist = np.histogram(scale.cpu().numpy(), bins=hist_bins, density=hist_density) scale_x_hist = np.histogram(scale[:, 0:1].cpu().numpy(), bins=hist_bins, density=hist_density) scale_y_hist = np.histogram(scale[:, 1:2].cpu().numpy(), bins=hist_bins, density=hist_density) scale_z_hist = np.histogram(scale[:, 2:3].cpu().numpy(), bins=hist_bins, density=hist_density) scale_ratio_hist = np.histogram(scale_ratio.cpu().numpy(), bins=hist_bins, density=hist_density) #rotation_before_activation_hist = np.histogram(rotation_before_activation.cpu().numpy(), bins=hist_bins, density=hist_density) #rotation_hist = np.histogram(rotation.cpu().numpy(), bins=hist_bins, density=hist_density) opacity_before_activation_hist = np.histogram(opacity_before_activation.cpu().numpy(), bins=hist_bins, density=hist_density) opacity_hist = np.histogram(opacity.cpu().numpy(), bins=hist_bins, density=hist_density) #xyz_hist = np.histogram(xyz.cpu().numpy(), bins=hist_bins, density=hist_density) x_hist = np.histogram(xyz[:, 0:1].cpu().numpy(), bins=hist_bins, density=hist_density) y_hist = np.histogram(xyz[:, 1:2].cpu().numpy(), bins=hist_bins, density=hist_density) z_hist = np.histogram(xyz[:, 2:3].cpu().numpy(), bins=hist_bins, density=hist_density) # collect_info info_dict = { 'scale_x_before_activation_range': scale_x_before_activation_range, 'scale_y_before_activation_range': scale_y_before_activation_range, 'scale_z_before_activation_range': scale_z_before_activation_range, 'scale_x_range': scale_x_range, 'scale_y_range': scale_y_range, 'scale_z_range': scale_z_range, 'scale_ratio_range': scale_ratio_range, 'opacity_before_activation_range': opacity_before_activation_range, 'opacity_range': opacity_range, 'x_range': x_range, 'y_range': y_range, 'z_range': z_range, 'scale_x_before_activation_hist': scale_x_before_activation_hist, 'scale_y_before_activation_hist': scale_y_before_activation_hist, 'scale_z_before_activation_hist': scale_z_before_activation_hist, 'scale_x_hist': scale_x_hist, 'scale_y_hist': scale_y_hist, 'scale_z_hist': scale_z_hist, 'scale_ratio_hist': scale_ratio_hist, 'opacity_before_activation_hist': opacity_before_activation_hist, 'opacity_hist': opacity_hist, 'x_hist': x_hist, 'y_hist': y_hist, 'z_hist': z_hist, } if hist_img_path is not None: # plot hist fig, axes = plt.subplots(4, 3, figsize=(15, 15)) axes[0, 0].hist(scale_x_before_activation_hist[1][:-1], scale_x_before_activation_hist[1], weights=scale_x_before_activation_hist[0]) axes[0, 0].set_title('scale_x_before_activation_hist') axes[0, 1].hist(scale_y_before_activation_hist[1][:-1], scale_y_before_activation_hist[1], weights=scale_y_before_activation_hist[0]) axes[0, 1].set_title('scale_y_before_activation_hist') axes[0, 2].hist(scale_z_before_activation_hist[1][:-1], scale_z_before_activation_hist[1], weights=scale_z_before_activation_hist[0]) axes[0, 2].set_title('scale_z_before_activation_hist') axes[1, 0].hist(scale_x_hist[1][:-1], scale_x_hist[1], weights=scale_x_hist[0]) axes[1, 0].set_title('scale_x_hist') axes[1, 1].hist(scale_y_hist[1][:-1], scale_y_hist[1], weights=scale_y_hist[0]) axes[1, 1].set_title('scale_y_hist') axes[1, 2].hist(scale_z_hist[1][:-1], scale_z_hist[1], weights=scale_z_hist[0]) axes[1, 2].set_title('scale_z_hist') axes[2, 0].hist(scale_ratio_hist[1][:-1], scale_ratio_hist[1], weights=scale_ratio_hist[0]) axes[2, 0].set_title('scale_ratio_hist') axes[2, 1].hist(opacity_before_activation_hist[1][:-1], opacity_before_activation_hist[1], weights=opacity_before_activation_hist[0]) axes[2, 1].set_title('opacity_before_activation_hist') axes[2, 2].hist(opacity_hist[1][:-1], opacity_hist[1], weights=opacity_hist[0]) axes[2, 2].set_title('opacity_hist') axes[3, 0].hist(x_hist[1][:-1], x_hist[1], weights=x_hist[0]) axes[3, 0].set_title('x_hist') axes[3, 1].hist(y_hist[1][:-1], y_hist[1], weights=y_hist[0]) axes[3, 1].set_title('y_hist') axes[3, 2].hist(z_hist[1][:-1], z_hist[1], weights=z_hist[0]) axes[3, 2].set_title('z_hist') plt.savefig(hist_img_path) plt.close() print(f'hist_img_path: {hist_img_path}') return info_dict def training_setup_no_feat(self, training_args): self.percent_dense = training_args.bg_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") # 相当于是 基于相机运动轨迹 or 基于场景大小 的一定比例 来确定 xyz 的学习率 # 但如果 我对 xyz 进行了归一化, 那么就应该避免 场景尺度的学习 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 training_setup(self, training_args): self.percent_dense = training_args.bg_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") if self._language_feature is None or self._language_feature.shape[0] != self._xyz.shape[0]: # 开始feature训练的时候，往模型中加入language feature参数 language_feature = torch.zeros((self._xyz.shape[0], training_args.feat_dim), device="cuda") self._language_feature = nn.Parameter(language_feature.requires_grad_(True)) # 初始化 language feature self._language_feature.data.uniform_(-1.0 / self.num_panoptic_objects, 1.0 / self.num_panoptic_objects) # 与 codebook-embedding 的 dim-norm 的情况 保持一致 self.with_dim_norm = training_args.gs_dim_norm self.with_render_dim_norm = training_args.render_gs_dim_norm 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"}, {'params': [self._language_feature], 'lr': training_args.language_feature_lr, "name": "language_feature"}, # TODO: training_args.language_feature_lr # bg-gs 不需要 codebook #{'params': self.feat_conv.parameters(), 'lr': training_args.feat_conv_lr, "name": "feat_conv"}, ] # bg-gs 不需要 背景模型 self.bg_model = self.bg_gs = None self.optimizer = torch.optim.Adam(l, lr=0.0, eps=1e-15) # bg-gs 不需要 multi-scheduler 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 create_from_pcd(self, pcd : BasicPointCloud, spatial_lr_scale : float, init_opacity = 0.1): 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 BackGround 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(init_opacity * 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 min_max(x, name='x', print_info=True): print(f"{name} range : {x.min().item():.7f} ~ {x.max().item():.7f}") # 如果x 是一维 if x.dim() == 1: x = x.unsqueeze(1) out = torch.cat((x.min(dim=0).values, x.max(dim=0).values), dim=0) return out class Background_Model(nn.Module): """ NeRF-like bg-model """ def __init__(self, img_width = 960, img_height = 640, aabb = None, cam_norm = None, model_type = "mlp", # mlp, rec_shell, sphere_shell model_cfg: dict = {"head_mlp_layer_width": 256,}, model_args: dict = {}, ) -> None: super().__init__() self.aabb = aabb self.img_width, self.img_height = img_width, img_height self.model_type = model_type self.model_cfg = model_cfg if model_type == 'mlp': self.sky_model = SkyModel(model_cfg) elif model_type == 'rec_shell' or model_type == 'sphere_shell': # TODO: has not been implemented self.sky_model = ShellModel(model_cfg) elif model_type == 'gs': self.sky_model = BackgroundGaussianModel(**model_cfg) else: raise NotImplementedError("model_type {} not implemented".format(model_type)) if model_type == 'gs': # 先构造高斯点 fg_aabb_center, fg_aabb_size = (aabb[0] + aabb[1]) / 2, aabb[1] - aabb[0] # cam-frustum aabb # use bg_scale to scale the aabb bg_gs_aabb = np.stack([fg_aabb_center - fg_aabb_size * model_args.bg_aabb_scale / 2, fg_aabb_center + fg_aabb_size * model_args.bg_aabb_scale / 2], axis=0) bg_aabb_center, bg_aabb_size = (bg_gs_aabb[0] + bg_gs_aabb[1]) / 2, bg_gs_aabb[1] - bg_gs_aabb[0] # add bg_gs_aabb SURFACE points bg_points = sample_on_aabb_surface(bg_aabb_center, bg_aabb_size, model_args.bg_gs_num, above_half=True) print("bg_gs_points min:",np.min(bg_points,axis=0)) print("bg_gs_points max:",np.max(bg_points,axis=0)) # DO NOT add bg_gs_points to points #points = np.concatenate([points, bg_points], axis=0) #shs = np.concatenate([shs, np.random.random((len(bg_points), 3)) / 255.0], axis=0) bg_shs = np.random.random((len(bg_points), 3)) / 255.0 # visualize #from utils.general_utils import visualize_points #visualize_points(points, fg_aabb_center, fg_aabb_size) #bg_ply_path = os.path.join(data_root, "ds-bg-points3d.ply") #storePly(bg_ply_path, bg_points, SH2RGB(bg_shs) * 255) bg_pcd = BasicPointCloud(points=bg_points, colors=SH2RGB(bg_shs), normals=np.zeros((len(bg_points), 3))) self.sky_model.create_from_pcd(bg_pcd, spatial_lr_scale = cam_norm['radius'] * model_args.bg_aabb_scale, #spatial_lr_scale = 1.0 # TODO 设置 xyz 在 0-1之内变化 ， 后续映射到 inf 空间里 ) # cam 距离 轨迹中心的 最大距离 # prepare bg-hyper-params: 暂且与 fg-gs 保持一致 #bg_model_args = model_args[''] # gs-training-setup self.sky_model.training_setup(model_args) else: # pixel coords x = torch.arange(0, img_width, dtype=torch.float32, device="cuda") y = torch.arange(0, img_height, dtype=torch.float32, device="cuda") x, y = x.long(), y.long() self.pixel_coord = torch.stack(torch.meshgrid(x, y, indexing='xy')).float().cuda() def get_ray_directions(self, c2w, intrinsic): x, y= self.pixel_coord[0].flatten(), self.pixel_coord[1].flatten() viewdirs = get_directions(x, y, c2w, intrinsic) viewdirs = viewdirs.reshape(self.img_height, self.img_width, 3) return viewdirs def forward(self, c2w, intrinsic, img_idx=None) -> Tensor: directions = self.get_ray_directions(c2w, intrinsic) sky_feat = self.sky_model(directions) rgb_sky = torch.sigmoid(sky_feat) return rgb_sky.permute(2, 0, 1)