从零开始的Taichi Gaussian Splatting复现

从零开始的Taichi Gaussian Splatting

最近在准备保研夏令营,所以我想从零开始复现一遍Gaussian Splatting练练手。由于主要是工程的复现,这里我就默认这篇blog的阅读者已经读过原论文3D Gaussian Splatting for Real-Time Radiance Field Rendering (inria.fr)且对这项技术有较深入的理解了。

在这份实现里,为了保证项目的易更改,同时做到“开箱即用”,我主要通过PyTorch来进行训练,Taichi来进行渲染。

0. 准备

在构建框架之前,我们需要先从一组图片中提取空间信息,这包括估计相机的位置和姿态,以及相机的内部参数。原论文使用 Colmap 来完成这项任务,我们也将采用同样的方法处理图片。

实际上,我曾尝试过使用 Open3D 结合 OpenCV 或 PyColmap 的方法来实现,但是发现这两种方式都存在一些问题:在 Windows 平台下,Open3D 不支持 CUDA 实现;通过pip安装的PyColmap也只能在 CPU 上运行,而要部署支运行在 GPU 上的 PyColmap,又需要先部署 Colmap 本体。在这种情况下,与其继续折腾 PyColmap,倒不如直接使用 Colmap + os 更为方便。

为了创建一个简单的图片数据集处理脚本,我们需要实现图像的读取和简单校正,并最终输出稀疏重建的三维点云及相机参数。首先,我们要获取一些必需的参数:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
import argparse
import os

parser = argparse.ArgumentParser(
description="Create dataset from pictures for sparse reconstruction. Pre-processing arguments must be provided.")

parser.add_argument('--name',
type=str,
required=True,
help='Specify the name of the dataset.')
parser.add_argument('--path',
type=str,
required=True,
help='Specify the path where the images are located.')
parser.add_argument('--resize',
type=float,
default=1.0,
help='Specify the resizing scale for input images.')

args = parser.parse_args()

colmap = 'colmap.exe'

script_dir = os.path.dirname(os.path.abspath(__file__))
output_path = os.path.join(script_dir, f'data/{args.name}/')
input_path = args.path
database_path = output_path + r'\database.db'

if not os.path.exists(output_path):
os.makedirs(output_path)

这段代码中,name 参数指定数据集名称,path 参数指定图片所在路径,resize 参数则用于设置图片的缩放倍率。与原始处理方法相比,我们采用的方法是直接读取图像集和数据集名称,然后在项目的 “data” 文件夹下生成数据集。

随后,我们需要对指定的图像进行多目重建,即通过多张照片,估计相机拍摄时的位置与姿态信息,并生成一个场景的大致点云。Colmap提供了易用的接口来进行基本的多目重建,包括特征点提取、特征匹配以及相机姿态优化。我们只需要安装Colmap,并在终端中运行这些指令:

1
2
3
4
5
colmap feature_extractor --database_path path/to/database.db --image_path path/to/input

colmap exhaustive_matcher --database_path path/to/database.db

colmap mapper --database_path path/to/database.db --image_path path/to/input --output_path path/to/output

其中,feature_extractor*指令用于在一组图片中放置特征点,并生成一个特征点数据库database.db;exhaustive_mathcher用于读取特征点数据库的内容并进行特征点之间的匹配,mapper则根据特征点匹配结果进行三维重建。为了提升数据集的质量,我们有必要在这些指令中增加额外的参数:

  • feature_extractor
    • ImageReader.single_camera:这个参数可以让Colmap在构造特征点时,指定场景中所有图片均由一台相机拍摄。如果不进行这项设置,由于Colmap的识别存在误差,重建的数据中可能认为场景中存在很多台相机,这会为后续的渲染带来很大的困难。
    • camera_model:设置拍摄图像的相机的类型。Colmap在构造特征点时需要指定一种相机类别,常用的有PIN_HOLE, RADIAL_FISHEYE, OPENCV等。默认情况下一般使用PIN_HOLE模型,即小孔成像模型,适用于没有镜头畸变的图像。由于我们并不提前知道拍摄的照片是否具有镜头畸变,因此我们使用OPENCV模型,在这个模型下Colmap会自行检测输入图片是否具有镜头畸变,并自动生成相机内参。
  • mapper
    • Mapper.ba_global_function_tolerance:三维重建的全局宽容度。值越小,宽容度越低,重建精度越高,重建的收敛也越慢。值过低的时候重建可能无法收敛,导致重建失败。

在Python中,我们可以使用os.system()来运行终端命令:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
# os.system(command) : 在终端运行command命令

# feature extraction
os.system(f'{colmap} feature_extractor '
f'--database_path {database_path} '
f'--image_path {input_path} '
f'--ImageReader.single_camera 1 '
f'camera_model OPENCV')
os.system(f'{colmap} exhaustive_matcher '
f'--database_path {database_path}')
os.system(f'{colmap} mapper '
f'--database_path {database_path} '
f'--image_path {input_path} '
f'--output_path {output_path} + /sparse '
f'--Mapper.ba_global_function_tolerance=0.000001')

由于我们进行的是三维场景的合成,因此我们还需要去除相机的畸变,以还原场景的原始三维信息:

1
2
3
4
5
6
os.system(f'{colmap} image_undistorter '
f'--image_path {input_path} '
f'--input_path {output_path}/0 '
f'--output_path {output_path}/undistorter '
f'--output_type COLMAP')
# 去镜头畸变的图片会自动被放置在输出路径的/images路径下

此时,运行脚本,就可以成功生成数据集了。*/data/数据集名称/undistorter/*路径下保存了我们需要的所有信息。在下一章,我们会将生成的数据全部读入Taichi中,并进行渲染。

为了排版好看,我在这里并未实现调整图像大小的功能。我将这部分内容移至了6.优化与调整章节中。

1. 数据录入

数据集加载

首先,在拥有一个数据集并准备开始训练时,我们需要将数据集中的初始信息读入到程序中。考虑到我们使用Taichi进行渲染,因此我们直接将数据读入到Taichi.field中。

由于Colmap生成的三维重建数据是二进制文件,因此我们需要构建一个简单的二进制文件读取函数。由于读取二进制文件并不是本文的重点,因此我在这里不过多赘述,而是直接提供一个可用的函数。

1
2
3
def read_next_bytes(fid, num_bytes, format_char_sequence, endian_character="<"):
data = fid.read(num_bytes)
return struct.unpack(endian_character + format_char_sequence, data)

输入参数的具体含义如下:

  • fid:表示具体的二进制文件
  • format_char_sequence:表示读入二进制数据的类型。
    • c: 一个字节的字符串
    • e: 一个 C 语言风格的浮点数(32位)
    • f: 单精度浮点数(32位)
    • d: 双精度浮点数(64位)
    • h: 短整型(2字节)
    • H: 无符号短整型(2字节)
    • i: 整型(4字节)
    • I: 无符号整型(4字节)
    • l: 长整型(4字节)
    • L: 无符号长整型(4字节)
    • q: 长长整型(8字节)
    • Q: 无符号长长整型(8字节)
  • endian_character:字节的表示顺序

随后,我们根据Colmap项目文档中描述的数据构造方式,依次读取各二进制文件中的数据即可。我们首先构造一个CameraData类,从相机参数的读取开始。相机参数由图像宽高和相机内参组成。我们使用Taichi的@ti.data_oriented来修饰这个类。这个修饰符表示将类中的数据存储在 GPU 中。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
import taichi as ti
from taichi.math import vec3, vec4

# Taichi提供了两种类修饰符,ti.dataclass和ti.data_oriented。其中,前者被称为数据类,适用于同时需要在Taichi作用域和Python作用域内使用的类;后者是面向数据类,用于需要高频更新和访问、且仅在Taichi作用域内使用的类。
@ti.data_oriented
class CameraData:
def __init__(self):
self.camera_id = 0
self.width = 0
self.height = 0
self.params = vec4(0, 0, 0, 0)

def read(self, camera_path: str):
with open(camera_path, "rb") as fid:
num_cameras = read_next_bytes(fid, 8, "Q")[0]

camera_properties = read_next_bytes(
fid,
num_bytes=24,
format_char_sequence="iiQQ"
)
self.camera_id = camera_properties[0]
self.width = camera_properties[2]
self.height = camera_properties[3]

# 相机内参fx,fy,cx,cy
num_params = 4
params = read_next_bytes(
fid,
num_bytes=8 * num_params,
format_char_sequence="d" * num_params
)
self.params = vec4(params[0], params[1], params[2], params[3])

相机内参的具体解释详解本节的数据验证与简单可视化部分

接下来导入点云信息。点云的二进制信息包含点的位置、颜色、编号、对应的相机编号、容差,以及在图像中相匹配的二维点编号。我们只需要记录点的位置与颜色信息即可。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
@ti.data_oriented
class PointData:
def __init__(self):
self.id = None
self.position = None
self.color = None
self.num = 0

def read(self, position_path):
with open(position_path, 'rb') as fid:
num_points = read_next_bytes(fid, 8, 'Q')[0]

self.num = num_points
self.id = ti.field(dtype=ti.i32, shape=num_points)
self.position = ti.Vector.field(3, dtype=ti.f64, shape=num_points)
self.color = ti.Vector.field(3, dtype=ti.f64, shape=num_points)
for _ in range(num_points):
point_properties = read_next_bytes(
fid,
num_bytes=43,
format_char_sequence='QdddBBBd'
)
i = point_properties[0]
xyz = point_properties[1:4]
rgb = point_properties[4:7]
self.position[_] = vec3(xyz[0], xyz[1], xyz[2])
self.color[_] = vec3(rgb[0] / 255.0, rgb[1] / 255.0, rgb[2] / 255.0)
self.id[_] = i # DEBUG

track_length = read_next_bytes(
fid,
num_bytes=8,
format_char_sequence='Q'
)
track_elems = read_next_bytes(
fid,
num_bytes=8 * track_length[0],
format_char_sequence='ii' * track_length[0]
)

现在我们可以正确读入点云数据了。然而现在这份代码的运行效率极底,因为Taichi会将数据全部存在显存上,而用 CPU 对显存进行频繁访问显然会很大程度地影响效率。为了解决这一问题,我们引入几个临时的numpy数组作为数据交换的中间件,先将数据读入numpy数组,再将numpy数组导入Taichi域中。numpy在CPU与GPU的交互上速度远超Taichi,而且Taichi的field可以很轻松地与numpy数组进行类型转换:

1
2
3
4
5
6
7
x = ti.field(float, shape=(3, 3))
a = np.arange(9).reshape(3, 3).astype(np.int32)
x.from_numpy(a)
print(x)
#[[0 1 2]
# [3 4 5]
# [6 7 8]]

在这里,我们定义三个临时的numpy数组,然后将数据存入numpy中,就可以快速地把数据置入taichi.field中了。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
def read(self, position_path):
with open(position_path, 'rb') as fid:
num_points = read_next_bytes(fid, 8, 'Q')[0]

self.num = num_points
self.id = ti.field(dtype=ti.i32, shape=(num_points, 1))
self.position = ti.Vector.field(3, dtype=ti.f64, shape=(num_points, 1))
self.color = ti.Vector.field(3, dtype=ti.f64, shape=(num_points, 1))

id_np = np.empty((num_points, 1))
pos_np = np.empty((num_points, 1, 3))
color_np = np.empty((num_points, 1, 3))

for _ in range(num_points):
point_properties = read_next_bytes(
fid,
num_bytes=43,
format_char_sequence='QdddBBBd'
)
id_np[_] = point_properties[0]
pos_np[_] = point_properties[1:4]
color_np[_] = point_properties[4:7]

track_length = read_next_bytes(
fid,
num_bytes=8,
format_char_sequence='Q'
)
track_elems = read_next_bytes(
fid,
num_bytes=8 * track_length[0],
format_char_sequence='ii' * track_length[0]
)
self.id.from_numpy(id_np)
self.position.from_numpy(pos_np)
self.color.from_numpy(color_np)

思考题:为什么从numpy导入taichi要比直接导入taichi速度更快?

最后导入图像数据。图像数据包含图像的拍摄位姿、拍摄相机编号、图像名称和关键点信息组成。我们只需要提取相机位姿和图像编号即可。注意,这里导入的拍摄位姿不是拍摄照片时相机在世界坐标下的位置与方向,而是世界坐标系到相机坐标系的投影。换言之,在后续计算相机位置时,我们需要还原相机的投影矩阵,并据此来计算相机的世界坐标。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
@ti.data_oriented
class ImageData:
def __init__(self):
self.num = 0
self.id = None
self.translate = None
self.rotation = None

def read(self, image_path):
with open(image_path, 'rb') as fid:
num_images = read_next_bytes(fid, 8, 'Q')[0]

self.num = num_images
self.id = ti.field(dtype=ti.i32, shape=num_images)

self.translate = ti.Vector.field(3, dtype=ti.f64, shape=num_images)
self.rotation = ti.Vector.field(4, dtype=ti.f64, shape=num_images)

id_np = np.empty((num_images,))
pos_np = np.empty((num_images, 3))
rot_np = np.empty((num_images, 4))

for _ in range(num_images):
image_properties = read_next_bytes(
fid,
num_bytes=64,
format_char_sequence='idddddddi'
)
id_np[_] = image_properties[0]
rot_np[image_properties[0] - 1] = image_properties[2:6]
rot_np[image_properties[0] - 1, 3] = image_properties[1]
pos_np[image_properties[0] - 1] = image_properties[5:8]

current_char = read_next_bytes(fid, 1, "c")[0]
while current_char != b"\x00":
current_char = read_next_bytes(fid, 1, "c")[0]
num_points = read_next_bytes(fid, num_bytes=8,
format_char_sequence="Q")[0]
x_y_id_s = read_next_bytes(fid, num_bytes=24 * num_points,
format_char_sequence="ddq" * num_points)

self.id.from_numpy(id_np)
self.translate.from_numpy(pos_np)
self.rotation.from_numpy(rot_np)

需要注意的是,Colmap数据中四元数的存储格式为(w,x,y,z),为了方便后续计算和我个人习惯,这里我将其以(x,y,z,w)的形式读入。

这样我们就完成了数据集的导入操作。最后,我们封装一个数据录入类,用于管理所有数据:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
@ti.dataclass
class DataReader:
camera: CameraData
points: PointData
images: ImageData

def __init__(self, camera_path, points_path, images_path):
self.camera = CameraData()
self.points = PointData()
self.images = ImageData()

self.camera_path = camera_path
self.points_path = points_path
self.images_path = images_path

def read(self):
self.camera.read(self.camera_path)
self.points.read(self.points_path)
self.images.read(self.images_path)

接下来我们要使用导入的数据,对原始点云进行渲染。

数据验证与简单可视化

我们首先使用Taichi提供的场景渲染框架对读入的数据集信息进行渲染,以此检查是否正确读入了数据。我们首先创建一个新的renderer.py,并在其中定义一个基本的渲染器:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
import argparse

import taichi as ti
from taichi.math import vec3
from data import DataReader

import config

parser = argparse.ArgumentParser(
description='Load a dataset.')
parser.add_argument('--path',
type=str,
default=r'.\data\test',
help='Specify the dataset to load.')
args = parser.parse_args()


@ti.data_oriented
class Renderer:
def __init__(self):
self.data = DataReader(args.path)
self.data.read() # 读入数据

self.window = ti.ui.Window("Test Renderer", (self.data.camera.width, self.data.camera.height),
vsync=True) # 将窗口尺寸设置的和图像大小一样
self.scene = self.window.get_scene() # 获取窗口附带的场景
self.canvas = self.window.get_canvas() # 渲染窗口画板
self.canvas.set_background_color((0.1, 0.1, 0.1)) # 渲染窗口画板背景颜色

self.camera = ti.ui.Camera()
self.camera_pos = vec3(0, 0, 0)
self.camera_lookat = vec3(0, 0, 0)

def rendering(self):
while self.window.running:
self.camera.position(self.camera_pos.x, self.camera_pos.y, self.camera_pos.z) # 调整相机位置
self.camera.lookat(self.camera_lookat.x, self.camera_lookat.y, self.camera_lookat.z) # 调整相机方向
self.scene.set_camera(self.camera) # 将相机参数绑定至场景
self.canvas.scene(self.scene) # 将场景信息绑定至渲染窗口画板
self.window.show()

此时,只需要实例化渲染器类,并调用rendering方法,就可以看见渲染窗口。由于还没有导入任何数据,所以渲染窗口中空无一物。

渲染窗口

接下来,我们需要将读入的数据输入到渲染器的场景和相机中。

我们先处理点云数据。Taichi为在场景中添加粒子提供了很方便的接口,我们只需要在rendering方法中增加一句:

1
self.scene.particles(self.data.points.position, 0.05, (1.0, 0.2, 0.2))

粒子数据就已经进入到场景中了。注意,此时运行代码,你仍然有可能只能看见一个空场景,因为我们还没有设置相机参数,此时相机有一定概率完全看不见任何粒子。

随后,我们需要将用于表示相机旋转的四元数data.images.rotation、表示相机位置的向量data.images.position,以及相机的内参data.camera.params输入到渲染器的相机中。我们希望相机的内参固定不变,而位姿可以随意变化。因此我们在初始化时直接固定相机的内参,并新建一个方法用于将图片的位姿输入到相机中。

我们首先处理相机的位姿。我们记在拍摄每张图像时,相机在世界坐标系下的坐标是$W_{camera}$,在自身坐标系下的坐标是$C_{camera}=(0,0,0)$。此时,相机距世界坐标系原点的位置偏移量为三维向量$T(x,y,z)$,角度偏移量为四元数$Q(w,x,y,z)$。由于我们读入的是世界坐标系到相机坐标系的投影,则显然有:
$$
C_{camera}=QW_{camera}+T
$$
由于相机坐标为$(0,0,0)$,我们对原式变形,可得:
$$
Q^{-1}(-T)=W_{camera}
$$
显然,我们只需要对偏移量取负值,并左乘上四元数的逆即可。由于四元数与三维向量运算不能直接进行数学乘法,我们需要先将四元数转换为旋转矩阵。Colmap所导出的四元数遵循Hamilton约定,对于四元数$q$可以化为如下旋转矩阵:
$$
\begin{bmatrix}
1 - 2q_y^2 - 2q_z^2 & 2q_xq_y - 2q_wq_z & 2q_xq_z + 2q_wq_y \
2q_xq_y + 2q_wq_z & 1 - 2q_x^2 - 2q_z^2 & 2q_yq_z - 2q_wq_x \
2q_xq_z - 2q_wq_y & 2q_yq_z + 2q_wq_x & 1 - 2q_x^2 - 2q_y^2
\end{bmatrix}
$$
这样我们就完成了相机坐标的确定。我们可以ImageData类中新增一个方法,用于预处理求出每张照片对应的相机坐标。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
@ti.kernel
def get_position(self):
for i in range(self.num):
q = self.rotation[i]
t = self.translate[i]

w_p = -t
r = mat3(
1 - 2 * q.y * q.y - 2 * q.z * q.z, 2 * q.x * q.y - 2 * q.w * q.z, 2 * q.x * q.z + 2 * q.w * q.y,
2 * q.x * q.y + 2 * q.w * q.z, 1 - 2 * q.x * q.x - 2 * q.z * q.z, 2 * q.y * q.z - 2 * q.w * q.x,
2 * q.x * q.z - 2 * q.w * q.y, 2 * q.y * q.z + 2 * q.w * q.x, 1 - 2 * q.x * q.x - 2 * q.y * q.y
)

w_p = r.inverse() @ w_p
self.position[i] = w_p

Taichi的@ti.kernel修饰符用于指定某方法在 GPU 上进行运算。需要注意的是,被@ti.kernel修饰的方法只能调用被@ti.func修饰的方法,表示这个方法也在 GPU 上运行;同时,@ti.func方法也只能被带@ti.kernel修饰的方法调用。

此时,我们在ImageReader中调用该方法,并在场景中添加对相机位置的渲染,就可以看见,相机的位置信息成功转换到世界坐标系了。

1
self.scene.particles(self.data.images.position, 0.01, (0.2, 1, 0.2))

同时,为了方便观察,我们可以在场景渲染中添加如下代码,以启用Taichi自带的相机移动:

1
self.camera.track_user_inputs(self.window, movement_speed=0.03, hold_key=ti.ui.LMB)

即可很容易地在场景中检查相机位置的运算是否正确了。

绿色的点代表相机位置,红色的点代表物体位置

接下来需要设置相机的旋转信息。由于Taichi并未提供直接对摄像机进行旋转的接口,我们需要使用通过设置相机的lookatup来调整相机的方向。这也十分容易实现——我们直接对(0,0,1)和(0,1,0)应用相机的投影矩阵即可。

再次修改ImageReader类,我们提前计算好相机的朝向与上方向,并在渲染时将相机的lookatup修改为对应数据值,点击运行,即可看见相机到达了正确的位置。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
@ti.kernel
def get_position(self):
for i in range(self.num):
q = self.rotation[i]
t = self.translate[i]

w_p = -t

s = mat3(
1, 0, 0,
0, -1, 0,
0, 0, 1
)

r = mat3(
1 - 2 * q.y * q.y - 2 * q.z * q.z, 2 * q.x * q.y - 2 * q.w * q.z, 2 * q.x * q.z + 2 * q.w * q.y,
2 * q.x * q.y + 2 * q.w * q.z, 1 - 2 * q.x * q.x - 2 * q.z * q.z, 2 * q.y * q.z - 2 * q.w * q.x,
2 * q.x * q.z - 2 * q.w * q.y, 2 * q.y * q.z + 2 * q.w * q.x, 1 - 2 * q.x * q.x - 2 * q.y * q.y
)

f = r.inverse()

w_l = f @ (vec3(0, 0, 1) - t)
w_u = f @ (vec3(0, -1, 0) - t)
w_p = f @ w_p
self.position[i] = w_p
self.lookat[i] = w_l
self.up[i] = w_u

左图为点云渲染结果,右图为实际图片

不幸的是,此时我们会发现,渲染的点云似乎上下颠倒了。我们并不会在这里解决这个问题,因为这一节主要是为了验证读入数据的正确性,图像的上下颠倒在这里无关紧要。

思考题:为什么这里图像会上下颠倒?

在这里,我们并不对相机的内参进行设置。在 Taichi 的 ti.ui.camera.Camera 类中,没有直接的方法来设置内参参数。这是因为 Taichi 的摄像机模型以更高级的概念(如位置、视点、上方向等)来配置,而非直接设置内参矩阵的具体值。到这一步的完整渲染器代码如下:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
@ti.data_oriented
class Renderer:
def __init__(self):
self.data = DataReader(args.path)
self.data.read() # 读入数据

self.window = ti.ui.Window("Test Renderer", (self.data.camera.width, self.data.camera.height),
vsync=True) # 将窗口尺寸设置的和图像大小一样
self.scene = self.window.get_scene() # 获取窗口附带的场景
self.canvas = self.window.get_canvas() # 渲染窗口画板
self.canvas.set_background_color((0.1, 0.1, 0.1)) # 渲染窗口画板背景颜色

self.camera = ti.ui.Camera()
self.camera.up(0, -1, 0)
self.camera_pos = vec3(0, 0, 0)
self.camera_lookat = vec3(0, 0, 0.5)
self.coordinate = 0

self.refresh = True

print(f"params:{self.data.camera.params}")

t_v2 = vec2(self.data.camera.width, self.data.camera.height)
d_fov = 80 * atan2(sqrt((t_v2.x * t_v2.x) + (t_v2.y * t_v2.y)), 2 * self.data.camera.params[0])
self.camera.fov(d_fov)

def input_image_posture(self, i):
print(f"current coordinate:{i}")
p = self.data.images.position[i]
l = self.data.images.lookat[i]
u = self.data.images.up[i]
self.camera.position(p.x, p.y, p.z)
self.camera.lookat(l.x, l.y, l.z)
self.camera.up(u.x, u.y, u.z)

def rendering(self):
while self.window.running:
self.camera.track_user_inputs(self.window, movement_speed=0.03, hold_key=ti.ui.LMB)

self.scene.set_camera(self.camera) # 将相机参数绑定至场景

self.scene.ambient_light((0.5, 0.5, 0.5))

self.scene.particles(self.data.points.position, 0.005, (1, 0.2, 0.2))
self.scene.particles(self.data.images.position, 0.01, (0.2, 1, 0.2))

if self.window.is_pressed(ti.ui.LEFT) and self.refresh:
self.refresh = False
self.coordinate -= 1
self.coordinate = max(self.coordinate, 0)
self.input_image_posture(self.coordinate)
elif self.window.is_pressed(ti.ui.RIGHT) and self.refresh:
self.refresh = False
self.coordinate += 1
self.coordinate = min(self.coordinate, self.data.images.num - 1)
self.input_image_posture(self.coordinate)
elif not self.window.is_pressed(ti.ui.LEFT) and not self.window.is_pressed(ti.ui.RIGHT):
self.refresh = True

self.canvas.scene(self.scene) # 将场景信息绑定至渲染窗口画板
self.window.show()

现在我们完成了对录入数据的基本配置和调整。接下来我们抛弃Taichi的渲染器,自己实现一个Gaussian Splatting渲染器。

2. 渲染

在开始实现自己的渲染器之前,我们先移除上一章Renderer类中的Scene、Camera以及相关参数,以方便我们自行实现渲染器。同时,我们新增一个image_pixel的field存储我们的渲染画面,并将渲染画面传递给canvas,以显示在屏幕上。:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
@ti.data_oriented
class Renderer:
def __init__(self):
self.data = DataReader(args.path)
self.data.read() # 读入数据

self.window = ti.ui.Window(
"Test Renderer",
(self.data.camera.width, self.data.camera.height)
)
self.canvas = self.window.get_canvas()
self.image_pixel = ti.Vector.field(
3,
dtype=ti.f32,
shape=(self.data.camera.width, self.data.camera.height)
)

def input_image_posture(self, i):
p = self.data.images.position[i]
l = self.data.images.lookat[i]
u = self.data.images.up[i]

def draw(self):
while self.window.running:
self.rendering()
self.canvas.set_image(self.image_pixel) # 将场景信息绑定至渲染窗口画板
self.window.show()

@ti.kernel
def rendering(self):
# 渲染一个简单的画面,以检查代码的正确性
w = self.data.camera.width
h = self.data.camera.height
for i, j in self.image_pixel:
self.image_pixel[i, j] = vec3(i / w, j / h, 1)

为了让整体结构更清晰,我将把画面绘制到窗口这一步的方法名改为了draw,将把场景渲染到画面这一步的方法名改为了rendering

基本渲染器

Splatting(抛雪球)的渲染方法,就是将各个点云映射到屏幕上,然后将其按一定的规则扩散。具体地,Gaussian Splatting就是把点渲染到屏幕上,然后用高斯函数来计算其扩散。

我们先实现一个可以以Gaussian的形式渲染数据的简单渲染器。首先定义一个Gaussian类,在Gaussian Splatting中,Gaussina的定义如下:
$$
G(x) = exp(-\frac{1}{2}(x)^T\Sigma^{-1})
$$
其中,$\Sigma$是由旋转矩阵和缩放矩阵组成的协方差矩阵:
$$
\Sigma = RSS^TR^T
$$
同时,论文中提到,要将三维的协方差矩阵映射到二维平面,我们可以直接去掉矩阵的z轴,以此将矩阵降为到二维。

基于Gaussian的定义,我们需要在我们的高斯类中定义包含带透明通道的rgb颜色、位置信息、缩放信息以及旋转信息。为了减小开销,我们可以用三维向量来表示位置信息和缩放信息,同时用四元数来表示旋转信息。在必要时,我们再将其化为矩阵形式即可。对于协方差矩阵的计算,我们在这里同时定义两个版本,一个基于 GPU 运算,一个基于 CPU 运算,以便于调试。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
@ti.dataclass
class Gaussian:
color: vec4
position: vec3
scale: vec3
rotation: vec4
covariance: mat3

def __init__(self, color, position, scale, rotation):
self.color = color
self.position = position
self.scale = scale
self.rotation = rotation
self.covariance = mat3(
1, 0, 0,
0, 1, 0,
0, 0, 1
)
self.cov_inv = mat3(
1, 0, 0,
0, 1, 0,
0, 0, 1
)

@ti.func
def cal_covariance(self):
q = self.rotation

r = mat3(
1 - 2 * q.y * q.y - 2 * q.z * q.z, 2 * q.x * q.y - 2 * q.w * q.z, 2 * q.x * q.z + 2 * q.w * q.y,
2 * q.x * q.y + 2 * q.w * q.z, 1 - 2 * q.x * q.x - 2 * q.z * q.z, 2 * q.y * q.z - 2 * q.w * q.x,
2 * q.x * q.z - 2 * q.w * q.y, 2 * q.y * q.z + 2 * q.w * q.x, 1 - 2 * q.x * q.x - 2 * q.y * q.y
)
s = mat3(
self.scale.x, 0, 0,
0, self.scale.y, 0,
0, 0, self.scale.z
)

c = r @ s @ s.transpose() @ r.transpose()
self.cov = c
self.cov_inv = c.inverse()
self.cov2 = mat2(
c[0, 0], c[0, 1],
c[1, 0], c[1, 1]
)
self.cov2_inv = self.cov2.inverse()

def cal_covariance_cpu(self):
q = self.rotation

r = mat3(
1 - 2 * q.y * q.y - 2 * q.z * q.z, 2 * q.x * q.y - 2 * q.w * q.z, 2 * q.x * q.z + 2 * q.w * q.y,
2 * q.x * q.y + 2 * q.w * q.z, 1 - 2 * q.x * q.x - 2 * q.z * q.z, 2 * q.y * q.z - 2 * q.w * q.x,
2 * q.x * q.z - 2 * q.w * q.y, 2 * q.y * q.z + 2 * q.w * q.x, 1 - 2 * q.x * q.x - 2 * q.y * q.y
)
s = mat3(
self.scale.x, 0, 0,
0, self.scale.y, 0,
0, 0, self.scale.z
)

c = r @ s @ s.transpose() @ r.transpose()
self.cov = c
self.cov_inv = c.inverse()
self.cov2 = mat2(
c[0, 0], c[0, 1],
c[1, 0], c[1, 1]
)
self.cov2_inv = self.cov2.inverse()

接下来,我们定义一个全局Gaussian用于测试,对单个Gaussian进行渲染。

1
2
3
4
5
6
7
8
# 一个全局变量Gaussian,用于测试
test_gaussian = Gaussian(
position=vec3(350, 350, 0),
color=vec4(0.5, 1, 0.5, 1),
rotation=vec4(0, 0, 0, 0),
scale=vec3(1, 1, 1)
)
test_gaussian.cal_covariance_cpu()
1
2
3
4
5
6
7
8
9
10
11
@ti.kernel
def rendering(self):
# 渲染一个简单的画面,以检查代码的正确性
w = self.data.camera.width
h = self.data.camera.height
for i, j in self.image_pixel:
self.image_pixel[i, j] =
test_gaussian.color.rgb * test_gaussian.color.a *
exp(
-0.5 * length((vec2(i, j) - test_gaussian.position.xy) @ test_gaussian.cov2_inv)
)

此时运行代码,就能看见我们定义的Gaussian被渲染到了屏幕上。我们也能通过修改test_gaussian的数值,来控制Gaussian点的渲染情况。

不同参数下单个Gaussian的渲染效果

接下来我们将Gaussian置于世界坐标原点,并试着将其映射至屏幕空间,实现point-based渲染器。首先我们要将坐标表示从三维向量改成四维向量,应用齐次坐标系。这只需要修改ImageReaderPointReader中的position field和对应的numpy数组即可。随后,我们删除ImageReader中对相机属性的预计算,因为我们不再应用Taichi的相机,而是要转而实现自己的Camera类。

在我们定义的相机类中,为了后续做相机控制的方便性,我们不使用transform+rotation的形式来构造相机视角,而是用lookatup的坐标轴形式定义相机。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
@ti.data_oriented
class Camera:
position: vec3
lookat: vec3
up: vec3

width: int
height: int

mat_vp: mat4 # 从世界空间到屏幕空间的变换矩阵

fov = float
z_near = float
z_far = float

def __init__(self, width, height, params, position, lookat, up, fov=30, z_near=0.05, z_far=1000):
self.width = width
self.height = height
self.position = position
self.lookat = lookat
self.up = up
self.right = vec3(0, 0, 0)
self.mat_vp = mat4(
1, 0, 0, 0,
0, 1, 0, 0,
0, 0, 1, 0,
0, 0, 0, 1
)
self.fov = fov
self.z_near = z_near
self.z_far = z_far

随后,我们需要实现一个方法,建立场景中Gaussian到屏幕的映射。在常规光栅化方法中,我们需要进行mvp变换+适口变换;而在我们的场景中,由于Gaussian直接由世界空间下的坐标定义,我们只需要做vp变换+适口变换即可。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
@ti.kernel
def cal_mat_vp_kernel(self, l: vec3, u: vec3, p: vec3, fov: float, z_near: float, z_far: float) -> mat4:
g = normalize(l)
t = normalize(u)
gxt = cross(l, u)

# view matrix
r_view = mat4(
gxt.x, gxt.y, gxt.z, 0,
t.x, t.y, t.z, 0,
g.x, g.y, g.z, 0,
0, 0, 0, 1
)
t_view = mat4(
1, 0, 0, -p.x,
0, 1, 0, -p.y,
0, 0, 1, -p.z,
0, 0, 0, 1
)
view = r_view @ t_view

# projection matrix
hvalue = fov * 0.00872664 # 角度 * 0.00872664
asp = self.width / self.height
top = tan(hvalue) * z_near
bottom = -top
right = top * asp
left = -right

p_m = mat4(
z_near, 0, 0, 0,
0, z_near, 0, 0,
0, 0, z_near + z_far, -z_near * z_far,
0, 0, 1, 0
)
p_n = mat4(
2 / (right - left), 0, 0, 0,
0, 2 / (top - bottom), 0, 0,
0, 0, -2 / (z_near - z_far), 0,
0, 0, 0, 1
)
p_p = mat4(
1, 0, 0, 0,
0, 1, 0, 0,
0, 0, 1, -(z_near + z_far) / 2,
0, 0, 0, 1
)
projection = p_n @ p_p @ p_m

# screen matrix
screen = mat4(
.5 * self.width, 0, 0, .5 * self.width,
0, .5 * self.height, 0, .5 * self.height,
0, 0, 1, 0,
0, 0, 0, 1
)

return screen @ projection @ view

随后,我们可以结合之前的内容,导入数据集,适当移动摄像机的位置并进行渲染。由于我们目前的投影变换并未考虑到Colmap所给出的相机参数,也并未进行屏幕空间的裁剪,因此在渲染时可能会存在较大程度的变形。我们会在随后的章节解决这个问题。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
@ti.kernel
def rendering(self, mat_vp: mat4):
for i, j in self.image_pixel:
self.image_pixel[i, j] = self.cal_color(vec2(i, j), mat_vp, self.data.points.num, self.gaussian)

@ti.func
def cal_color(self, pixel_pos: vec2, mat_vp: mat4, n: int, gaussian) -> vec3:
res = vec3(0.0)
for i in range(n):
mt4 = mat_vp @ gaussian[i].position
mt = mt4.xyz / mt4.w
if mt.x > self.camera.width * 2 or mt.y > self.camera.height * 2:
continue
res += (
gaussian[i].color.rgb * gaussian[i].color.a *
exp(
-0.5 * length((pixel_pos - mt.xy) @ gaussian[i].cov2_inv)
)
)
return res

对Hotdog数据集的渲染。可以注意到,我们当前的渲染效率极低,且使得物体产生了莫名其妙的发光效果。接下来我们将先优化渲染器的效率,并引入相机参数,最后再实现Gaussian Splatting渲染。

接下来,我们要根据data.camera.params,即相机内参,设置渲染器相机的各项参数。

Gaussian Splatting

3. 训练

4. 优化策略

5. 结果评估

6. 优化与调整

7. 附录


从零开始的Taichi Gaussian Splatting复现
https://tridiamond.tech/post/从零开始的Taichi-Gaussian-Splatting.html
作者
想变成一只白泽
发布于
2024年3月23日
许可协议