Tri-MipRF: Tri-Mip Representation for Efficient Anti-Aliasing Neural Radiance Fields

环境配置

Python:3.9.0, Pytorch:1.13.1, Cuda:11.7, Tinycudann:1.7，按照这个版本安装，tinycudann不在报错。（好烦啊，主页上还说要安装cu11.6，结果根本不可以）
这里记录一个和NeRO安装一样的包nvdiffrast

cd nvdiffrast
pip install .

另一个包tiny-cuda-nn，需要自行跑到github页面安装给的步骤安装，可能还会遇到cmake版本过低的问题，下载cmake

tar -zxvf cmake-3.24.1.tar.gz
cd cmake-3.24.1
./bootstrap

# 中间如果报错
apt-get install libssl-dev

make
make install


# 配置cmake路径
export ....
source ~/.bashrc

# 查看版本 
 cmake --version

记录windows 系统出现的问题

记录nvdiffrast配置时出现的问题
顺序是：先装cuda，再装vsstudio，再新建虚拟环境，安装pytorch，在cmd对应目录下启动vscode，配置nvdiffrast，这一套下来应该没什么问题了，但注意如果你电脑上本来就有cuda的话，也都要重装，不然还是会报错。

subprocess.CalledProcessError: Command '['where', 'cl']' returned non-zero exit status 1.

参考这个链接能解决，如果在windows上你的nvdiffrast能正常使用的话，VSstudio应该没问题，所以就按照他给的目录找到大概相似的路径就可以~
配置好环境后，在cmd重启vscode(我才知道原来还能这样！每天学到新知识！)

code .

Update

AttributeError: 'NoneType' object has no attribute 'ContractionType'会报这个错误，太窒息了！！
这里是因为NerfAcc: No CUDA toolkit found. NerfAcc will be disabled.这个问题，nvcc没有找到，所以配置环境变量，方法如下：

vim ~/.bashrc 
export LD_LIBRARY_PATH=/usr/local/cuda/lib64
export PATH=$PATH:/usr/local/cuda/bin
source ~/.bashrc
nvcc -V

配置成功后再新开一个终端运行比较保险。

python main.py --ginc config_files/ms_blender/TriMipRF.gin 

代码细节

数据输入

ray_dataset.py 调用数据集，parsers文件夹下写不同数据集的读取方式。

Init dataloader ...
Init model ...
Init trainer ...

主函数trimipRF.py调用nerfacc包，先看TriMipRFModel的初始化：

# 创建occ grid，参数有region of interest, resolution
self.ray_sampler = nerfacc.OccupancyGrid(roi_aabb=self.aabb, resolution=occ_grid_resolution)  

接下来需要自己写一个函数，用来更新这个occupancy grid，这里对应的是train.py中，每次迭代前都调用一下 self.model.before_iter(step)

# 调用内容是更新occ grid函数
self.ray_sampler.every_n_step(
            step=step,
            occ_eval_fn=lambda x: self.field.query_density(
                x=self.contraction(x),
                level_vol=torch.empty_like(x[..., 0]).fill_(self.occ_level_vol),
            )['density']
            * self.render_step_size,
            occ_thre=5e-3,
        )

关键函数

Tri-MipRF-main/neural_field/field/trimipRF.py 定义了MLP和编码方式，通过以下两个函数调用MLP并获得density和color

self.field.query_density(x=positions, level_vol=level_vol, return_feat=True)
输入：
	position 采样点坐标(n,3)
	level_vol 不知道是什么，和mip相关的参数
	return_feat 是否需要返回特征（如果后续要输入给color mlp就需要）
输出：包含以下两个keys的字典
	density (n,1)
	feature (n,15)
作用：输入position获得density（这里不需要知道方向）

self.field.query_rgb(self, dir, embedding)
输入：
	dir 方向
	embedding 前面和density一起得到的特征
输出：
	{"rgb": rgb} 颜色
作用：有feature和dir得到颜色

Tri-MipRF-main/neural_field/model/trimipRF.py 输入rays_o rays_d，调用前述的函数以获得density和color

sigma_fn(t_starts, t_ends, ray_indices)
输入：如上
输出：self.field.query_density(positions, level_vol)['density']
作用：给出position，调用query_density获得density

rgb_sigma_fn(t_starts, t_ends, ray_indices)
输入：如上
输出：rgb, density
作用：给出position，调用query_density获得density和feature，接着调用query_rgb获得color

rendering函数调用了以上两个函数

self.rendering(t_starts,t_ends,ray_indices,rays,rgb_sigma_fn=rgb_sigma_fn,render_bkgd=background_color)
输入：如上
输出：color, alpha, depths, ...
作用：根据density求出weights，然后做体渲染求出color，为求loss做准备

train和eval对比

iter_train_loader = iter(self.train_loader)
iter_eval_loader = iter(self.eval_loader)
metrics = self.train_iter(step, data=next(iter_train_loader), ...)
metrics, final_rb, target = self.eval_img(
                    next(iter_eval_loader) if self.varied_eval_img else eval_0,
                    compute_metrics=True,
                )

rb = self.model(step, cam_rays, target.render_bkgd)
rb = self.model(
                flatten_rays[i : i + self.test_chunk_size],
                flatten_target[i : i + self.test_chunk_size].render_bkgd,
            )

添加输入x输出level的函数

# 每次调用sdf都需要知道 level_vol
def density(x, level_vol)
	self.log2_plane_size = math.log2(plane_size)	# plane_size自己定义的plane分辨率
	level = (level_vol if level_vol is None else level_vol + self.log2_plane_size)
	enc = self.encoding(x.view(-1, 3), level=level.view(-1, 1),)	# 这里的level怎么获得？

# density函数中的level_vol
sample_ball_radii = self.compute_ball_radii(distance, radiis, cos)
level_vol = torch.log2(sample_ball_radii / self.feature_vol_radii)  # real level should + log2(feature_resolution)

# 其中self.compute_ball_radii()是函数，而self.feature_vol_radii是变量
self.feature_vol_radii = self.aabb_size[0] / 2.0

self.aabb_size[0]=3 # 表示放入的正方体的场景长宽高

综上所示，由pts和dirs先求出distance, radiis, cos，接着求出level_vol，然后求出level

流程

对着流程图再从头到尾介绍一遍，敲重点！

1. 由公式(3)出发获得的圆锥内每个内切球S(x,r)半径r，可以看出半径r和x的位置相关，即：与t相关。接着将球投影到3平面上获得3个圆片Discs={ $D_{xy}$, $D_{xz}$, $D_{yz}$ }。
2. 以$D_{xy}$为例，由公式(5) 查询$D_{xy}$的半径对应Tri-Mip编码的哪一层。能查询出它位于$M_{xy}^{L_i}$和$M_{xy}^{L_{i+1}}$之间，然后将以$(x,y)$为中心的$D_{xy}$投影到两层M上，获得8个点，对这8个点进行三线性插值获得特征$f_{XY}$
3. 获得完所有特征后，进行concat然后输入给MLP
4. 需要强调的是，这里的Tri-Mip编码是进行随机初始化的，可学习的显式表达，一共有N层，每一层是上一层的下采样1/2。相当于第一层$M_{xy}^{L_0}$维度是$W\times H\times C$，第二层$M_{xy}^{L_1}$维度是$W/2\times H/2\times C$…以此类推。 整个过程就是，查询内切球对应的编码，然后concat作为输入。

公式3我推了一晚上也没推出来，实在看不出是哪里不对，有大佬讲一下吗？？

记录细节

self.compute_ball_radii()函数对应公式（3），虽然我也没看出来怎么对应上的，这里似乎假设f=1：

sample_ball_radii = self.compute_ball_radii(
                    distance, radiis, cos
                )

def compute_ball_radii(distance, radiis, cos):
    inverse_cos = 1.0 / cos
    tmp = (inverse_cos * inverse_cos - 1).sqrt() - radiis
    sample_ball_radii = distance * radiis * cos / (tmp * tmp + 1.0).sqrt()
    return sample_ball_radii

输入：t, r^., 方向向量与z轴负半轴夹角的余弦(z是距离)
输出：每个采样点的r
作用：计算每个球的半径

不重要的随记

radiis: 2D平面上，一个像素所占半径，当相机参数只有一种时，radiis只有一种
ray_cos: 光线与水平轴的夹角
ray_cos=torch.matmul(
                directions,
                torch.tensor([[0.0, 0.0, self.sign_z]], device=device).T,
            )	# self.sign_z=1/-1
sample_ball_radii: 每个内切球的半径，需要知道distance, radiis, ray_cos
sample_ball_radii = self.compute_ball_radii(distance, radiis, cos)
level/level_vol: 由内切球半径求得每个点所在的level

sdf(x, encoding, level) # 需要知道level