本文主要是介绍NeRF从入门到放弃5: Neurad代码实现细节,希望对大家解决编程问题提供一定的参考价值,需要的开发者们随着小编来一起学习吧!
Talk is cheap, show me the code。
CNN Decoder
如patch设置为32x32,patch_scale设置为3,则先在原图上采样96x96大小的像素块,然后每隔三个取一个像素,降采样成32x32的块。
用这32x32个像素render feature,再经过CNN反卷积预测出96x96的像素,与真值对比。
def _patches_from_centers(self,image: torch.Tensor,patch_center_indices: torch.Tensor,rgb_size: int,device: Union[torch.device, str] = "cpu",
):"""Convert patch center coordinates to the full set of ray indices and image patches."""offsets = torch.arange(-(rgb_size // 2), (rgb_size // 2) + rgb_size % 2, device=device)zeros = offsets.new_zeros((rgb_size, rgb_size))relative_indices = torch.stack((zeros, *torch.meshgrid(offsets, offsets, indexing="ij")), dim=-1)[None] # 1xKxKx3,原图采样大小rgb_indices = patch_center_indices[:, None, None] + relative_indices # NxKxKx3ray_indices = rgb_indices[:, self.patch_scale // 2 :: self.patch_scale, self.patch_scale // 2 :: self.patch_scale] # NxKfxKfx3,降采样ray_indices = ray_indices.reshape(-1, 3) # (N*Kf*Kf)x3img_patches = image[rgb_indices[..., 0], rgb_indices[..., 1], rgb_indices[..., 2]]return ray_indices, img_patches
相机位姿优化
参考nerfstudio/cameras/camera_optimizers.py
每迭代一次优化一次
- 初始化
self.pose_adjustment = torch.nn.Parameter(torch.zeros((num_cameras, 6), device=device)) # Nx6,前3维表示平移,后三维表示后3维表示切向量,再通过exp_map_SO3xR3,把6维变量映射为位姿和位移变量。相当于优化的是每个相机的标定参数
- 计算位姿偏移量
def forward(self,indices: Int[Tensor, "camera_indices"],) -> Float[Tensor, "camera_indices 3 4"]:correction_matrices = exp_map_SO3xR3(self._get_pose_adjustment()[indices, :])
- 应用到相机的原始位姿上
def apply_to_raybundle(self, raybundle: RayBundle) -> None:"""Apply the pose correction to the raybundle"""if self.config.mode != "off":correction_matrices = self(raybundle.camera_indices.squeeze()) # type: ignoreraybundle.origins = raybundle.origins + correction_matrices[:, :3, 3]raybundle.directions = (torch.bmm(correction_matrices[:, :3, :3], raybundle.directions[..., None]).squeeze().to(raybundle.origins))
- 可学习的6维向量如何转成旋转矩阵
# nerfstudio/cameras/lie_groups.py
# We make an exception on snake case conventions because SO3 != so3.
def exp_map_SO3xR3(tangent_vector: Float[Tensor, "b 6"]) -> Float[Tensor, "b 3 4"]:"""Compute the exponential map of the direct product group `SO(3) x R^3`.This can be used for learning pose deltas on SE(3), and is generally faster than `exp_map_SE3`.Args:tangent_vector: Tangent vector; length-3 translations, followed by an `so(3)` tangent vector.Returns:[R|t] transformation matrices."""# code for SO3 map grabbed from pytorch3d and stripped down to bare-boneslog_rot = tangent_vector[:, 3:]nrms = (log_rot * log_rot).sum(1)rot_angles = torch.clamp(nrms, 1e-4).sqrt()rot_angles_inv = 1.0 / rot_anglesfac1 = rot_angles_inv * rot_angles.sin()fac2 = rot_angles_inv * rot_angles_inv * (1.0 - rot_angles.cos())skews = torch.zeros((log_rot.shape[0], 3, 3), dtype=log_rot.dtype, device=log_rot.device)skews[:, 0, 1] = -log_rot[:, 2]skews[:, 0, 2] = log_rot[:, 1]skews[:, 1, 0] = log_rot[:, 2]skews[:, 1, 2] = -log_rot[:, 0]skews[:, 2, 0] = -log_rot[:, 1]skews[:, 2, 1] = log_rot[:, 0]skews_square = torch.bmm(skews, skews)ret = torch.zeros(tangent_vector.shape[0], 3, 4, dtype=tangent_vector.dtype, device=tangent_vector.device)ret[:, :3, :3] = (fac1[:, None, None] * skews+ fac2[:, None, None] * skews_square+ torch.eye(3, dtype=log_rot.dtype, device=log_rot.device)[None])# Compute the translationret[:, :3, 3] = tangent_vector[:, :3]return ret
Apperance embedding
就是简单的使用torch.nn.Embedding(num_embeds, self.config.appearance_dim)
# Appearance embedding settings
# num_sensor指的是相机个数,如果配置temporal,则每一帧都有单独的embedding
if self.config.use_temporal_appearance:self._num_embeds_per_sensor = math.ceil(self._duration * self.config.temporal_appearance_freq)num_embeds = num_sensors * self._num_embeds_per_sensor
else:num_embeds = num_sensors# num_embeds=6,self.config.appearance_dim=16,表示6个相机,每个相机有16维的Embedding特征
self.appearance_embedding = torch.nn.Embedding(num_embeds, self.config.appearance_dim)def _get_appearance_embedding(self, ray_bundle, features):sensor_idx = ray_bundle.metadata.get("sensor_idxs")if sensor_idx is None:assert not self.training, "Sensor sensor_idx must be present in metadata during training"sensor_idx = torch.full_like(features[..., :1], self.fallback_sensor_idx.value, dtype=torch.long)if self.config.use_temporal_appearance:time_idx = ray_bundle.times / self._duration * (embd_per_sensor := self._num_embeds_per_sensor)before_idx = time_idx.floor().clamp(0, embd_per_sensor - 1)after_idx = (before_idx + 1).clamp(0, embd_per_sensor - 1)ratio = time_idx - before_idx# unwrap to true embedding indices, which also account for the sensor index, not just the time indexbefore_idx, after_idx = (x + sensor_idx * embd_per_sensor for x in (before_idx, after_idx))before_embed = self.appearance_embedding(before_idx.squeeze(-1).long())after_embed = self.appearance_embedding(after_idx.squeeze(-1).long())embed = before_embed * (1 - ratio) + after_embed * ratioelse:embed = self.appearance_embedding(sensor_idx.squeeze(-1))return embed
lidar建模和采样
lidar发射射线和camer类似,只需要根据世界坐标系下lidar原点的坐标和点云的坐标,就能确定一条射线了,沿这条射线采样点,真值是这条射线上真正扫描到的点。
采样时,根据每次迭代设置的采样点数N如16384,平均到每帧的每个点上。
采样方式是把全部帧的点云concate起来,每个点有个全局的序号和帧的idx,假设总点数为100万,采样时在0-100万之间随机生成N个随机数。
def get_lidar_batch_and_ray_bundle(self):if not len(self.lidar_dataset.lidars):return None, Nonebatch = self.point_sampler.sample(self.cached_points)ray_indices = batch.pop("indices") # Nx2, 0: lidar index, 1: point index,共采样16384个点,每帧采样点数一样ray_bundle: RayBundle = self.lidar_ray_generator(ray_indices, points=batch["lidar"]) #把所有的点都concate起来了return batch, ray_bundle # batch存储lidar原始点,ray_bundle存储采样的方向,原点信息
另外,pixel_area的作用没太看懂,有点像是MipNerf里面的用锥形体界面去积分,而不是直接的射线?
dx = self.horizontal_beam_divergence[lidar_indices.squeeze(-1)] # ("num_rays":...,)dy = self.vertical_beam_divergence[lidar_indices.squeeze(-1)] # ("num_rays":...,)pixel_area = dx * dy # ("num_rays":..., 1)
sdf实现
如果使用sdf,直接根据下面公式预测出不透明度α;否则便是先预测出密度density,再根据density积分得到不透明度。
因此两种render weight的方式是不同的。
if self.config.use_sdf:signed_distance = geo_out # 直接把mlp的输出当作signed distanceoutputs[FieldHeadNames.SDF] = signed_distanceoutputs[FieldHeadNames.ALPHA] = self.sdf_to_density(signed_distance)
else:outputs[FieldHeadNames.DENSITY] = trunc_exp(geo_out) # 调用了torch.exp(), 为什么不能直接用geo_out作为density?有两个原因:1.因为density的物理意义是大于0的,geo_out不保证大于0 2. 网络输出的值可能非常小,使用epx放大,可以保持数值稳定性self.sdf_to_density = SigmoidDensity(self.config.sdf_beta, learnable_beta=self.config.learnable_beta)
这个名字应该叫SigmoidAlpha,最后输出的被当做α,不是density了
class SigmoidDensity(nn.Module):"""Learnable sigmoid density"""def __init__(self, init_val, beta_min=0.0001, learnable_beta=False):super().__init__()self.register_buffer("beta_min", torch.tensor(beta_min))self.register_parameter("beta", nn.Parameter(init_val * torch.ones(1), requires_grad=learnable_beta))def forward(self, sdf: Tensor, beta: Union[Tensor, None] = None) -> Tensor:"""convert sdf value to density value with beta, if beta is missing, then use learable beta"""if beta is None:beta = self.get_beta()# negtive sdf will have large densityreturn torch.sigmoid(-sdf * beta) #这里就是上面的公式,这里叫α,和density不是一个东西def get_beta(self):"""return current beta value"""beta = self.beta.abs() + self.beta_minreturn beta
render_weight_from_alpha()
直接处理不透明度,而[render_weight_from_density()
]则需要先从密度计算不透明度。
def _render_weights(self, outputs, ray_samples):value = outputs[FieldHeadNames.ALPHA if self.config.field.use_sdf else FieldHeadNames.DENSITY].squeeze(-1)if self.device.type in ("cpu", "mps"):# Note: for debugging on devices without cudaweights = torch.zeros_like(value) + 0.5elif self.config.field.use_sdf:weights, _ = nerfacc.render_weight_from_alpha(value)else:weights, _, _ = nerfacc.render_weight_from_density(t_ends=ray_samples.frustums.ends.squeeze(-1),t_starts=ray_samples.frustums.starts.squeeze(-1),sigmas=value,)return weights
这篇关于NeRF从入门到放弃5: Neurad代码实现细节的文章就介绍到这儿,希望我们推荐的文章对编程师们有所帮助!