3DGS Code Review

3D Gaussian Splatting for Real-Time Radiance Field Rendering (SIGGRAPH 2023)

code :
https://github.com/graphdeco-inria/gaussian-splatting
reference :
https://charlieppark.kr
NeRF and 3DGS Study
https://arxiv.org/abs/2401.03890

Code Flow Diagram

reference : https://charlieppark.kr

train.py

Fig 1. train.py Algorithm 
gaussians = GaussianModel(dataset.sh_degree)
scene = Scene(dataset, gaussians)

Gaussian Initialize

Fig 1.의 빨간 박스 : SfM pcd로부터 Gaussian param. 초기화

class Scene:
  def __init__(...):
    ...
    self.gaussians.create_from_pcd(scene_info.point_cloud, self.cameras_extent)
class Scene.__init__()
sceneLoadTypeCallbacks(dict) > readColmapSceneInfo
GaussianModel.create_from_pcd()

Densification

Fig 1.의 초록 박스 : densification (clone and split)

train.py
GaussianModel.add_densification_stats()
GaussianModel.densify_and_prune()
GaussianModel.densify_and_clone(), GaussianModel.densify_and_split()
GaussianModel.densification_postfix()

GS Rasterize

Fig 2. GS Rasterize Algorithm

Fig 2.의 노란 박스 : GS rasterization by CUDA

render()
diff_gaussian_rasterization.__init__.py.GaussianRasterizer(nn.Module)
diff_gaussian_rasterization.__init__.py._RasterizeGaussians(torch.autograd.Function).forward()
diff_gaussian_rasterization.__init__.py._RasterizeGaussians(torch.autograd.Function).backward()
C++/CUDA code를 컴파일하여 diff_gaussian_rasterization._C 패키지를 만들기 위한 setup.py
python 함수와 rasterize_points.cu 의 CUDA 함수를 PYBIND11_MODULE()로 연결 
// create a tile grid (a group of blocks)
dim3 tile_grid((width + BLOCK_X - 1) / BLOCK_X, (height + BLOCK_Y - 1) / BLOCK_Y, 1);
// create a block (a group of threads(pixels))
dim3 block(BLOCK_X, BLOCK_Y, 1);

FORWARD

/* preprocessCUDA()의 main code flow */

// global rank of current thread (Gaussian)
auto idx = cg::this_grid().thread_rank();

in_frustum(/*...*/) // near culling

// compute 3D covariance in world-coord.
computeCov3D(/*...*/)
// compute 2D covariance in image-coord.(z=1) by EWA Splatting
computeCov2D(/*...*/)

// compute inverse of 2D covariance
float3 conic = { cov.z * det_inv, -cov.y * det_inv, cov.x * det_inv }; 

// image from ndc to pixel-coord.
float2 point_image = { ndc2Pix(p_proj.x, W), ndc2Pix(p_proj.y, H) };

// overlap 기준 : 반지름 = 3 * sigma
float my_radius = ceil(3.f * sqrt(max(lambda1, lambda2)));

// idx-th Gaussian과 겹치는 가장 왼/오른쪽, 위/아래쪽 tile의 index를 계산
getRect(point_image, my_radius, rect_min, rect_max, grid);

// idx-th Gaussian의 depth, radius 저장
depths[idx] = p_view.z; radii[idx] = my_radius;
// pixel-coord. image 저장
points_xy_image[idx] = point_image; 
// idx-th Gaussian의 inverse covariance, opacity 저장
conic_opacity[idx] = { conic.x, conic.y, conic.z, opacities[idx] }; 
// idx-th Gaussian과 겹치는 tile 개수 저장
tiles_touched[idx] = (rect_max.y - rect_min.y) * (rect_max.x - rect_min.x);
__global__ void preprocessCUDA()
__device__ bool in_frustum()
__device__ void computeCov3D() 
// world-to-camera of 3D mean to obtain J
float3 t = transformPoint4x3(mean, viewmatrix);
__device__ float3 computeCov2D()
__device__ float ndc2Pix() and __device__ void getRect() 
/*
// 각 Gaussian이 touch한 tile 개수
geomState.tiles_touched : [2, 3, 0, 2, 1]의 주소

// inclusive sum 수행한 후의 output array
geomState.point_offsets : [2, 5, 5, 7, 8]의 주소

// Gaussian이 touch한 총 tile 개수 또는 duplicated Gaussian instance 수
num_rendered = *(geomState.point_offsets + P-1) : 8
*/
CHECK_CUDA(cub::DeviceScan::InclusiveSum(geomState.scanning_space, geomState.scan_size, geomState.tiles_touched, geomState.point_offsets, P), debug)

from collections import deque

# bit sequence 말고 자연수를 정렬한다고 할 때
# 1의 자릿수 기준으로 정렬한 뒤
# 10의 자릿수 기준으로 정렬한 뒤
# ...
def radixSort():
    nums = list(map(int, input().split(' ')))
    buckets = [deque() for _ in range(10)] # 각 자릿수(0~9)에 대응되는 10개의 empty deque()
    
    max_val = max(nums)
    queue = deque(nums) # 정렬할 숫자들
    digit = 1 # 정렬 기준이 되는 자릿수
    
    while (max_val >= digit): # 가장 큰 수의 자릿수일 때까지만 실행
        while queue:
            num = queue.popleft() # 정렬할 숫자
            buckets[(num // digit) % 10].append(num) # 각 자릿수(0~9)에 따라 buckets에 num을 넣는다.
        
        # 해당 정렬 기준 자릿수에서 buckets에 다 넣었으면, buckets에 담겨있는 순서대로 꺼내와서 queue에 넣음
        for bucket in buckets:
            while bucket:
                queue.append(bucket.popleft())

        digit *= 10 # 정렬 기준이 되는 자릿수 증가시키기
    
    print(list(queue))

/* renderCUDA()의 main code flow */

// current block
// block.group_index() : current block의 (x, y) index
auto block = cg::this_thread_block();

// current block's first pixel
uint2 pix_min = { block.group_index().x * BLOCK_X, block.group_index().y * BLOCK_Y };
// current thread
// pixf : current thread의 (x, y) index
float2 pixf = { (float)(pix_min.x + block.thread_index().x), (float)(pix_min.y + block.thread_index().y) };

// check if current thread corresponds to a valid pixel
bool inside = pix.x < W&& pix.y < H;

// current tile의 Gaussians 수 (range.y - range.x)를 처리하기 위한 batch 수
const int rounds = ((range.y - range.x + BLOCK_SIZE - 1) / BLOCK_SIZE);

// current block 내 모든 threads가 __syncthreads_count()라는 barrier에 도달해서
// done==True를 만족하는 threads 수 반환
int num_done = __syncthreads_count(done);

for (int i = 0; i < rounds; i++, toDo -= BLOCK_SIZE)
{
    /* fetch step */
    // current block 내 모든 threads가 각자 progress-th Gaussian을 병렬적으로 fetch하여
    // BLOCK_SIZE=256개(batch)만큼씩 fetch into shared memory
    // block.thread_rank() : current thread의 local rank
    int progress = i * BLOCK_SIZE + block.thread_rank();
    if (range.x + progress < range.y){
        int coll_id = point_list[range.x + progress];
        collected_id[block.thread_rank()] = coll_id;
        collected_xy[block.thread_rank()] = points_xy_image[coll_id];
        collected_conic_opacity[block.thread_rank()] = conic_opacity[coll_id];
    }
    // current block 내 모든 threads가 여기 barrier에 도달할 때까지 대기
    block.sync();

    /* rasterization step */
    // thread(pixel)마다 병렬적으로 BLOCK_SIZE=256개(batch)의 Gaussians를 앞에서부터 alpha-compositing
    // accumulated opacity T 가 너무 작아지면 해당 threads는 alpha-compositing 종료
    for (int ch = 0; ch < CHANNELS; ch++)
	    C[ch] += features[collected_id[j] * CHANNELS + ch] * alpha * T;
    for (int ch = 0; ch < CHANNELS; ch++)
		// 마지막에 bg color까지 alpha-compositing
        out_color[ch * H * W + pix_id] = C[ch] + T * bg_color[ch];
}
__global__ void __launch_bounds__(BLOCK_X * BLOCK_Y) renderCUDA()

BACKWARD