CMU 3DGS Tutorial의 Tutorial

Robot navigation에서 Gaussian Splatting의 활용 관련 공부를 요새 많이 하고 있는데, 지도해주시는 선배님이 이 레포를 따라해보면 GS 학습을 간략하게 맛볼 수 있다고 추천해주셨다. 논문에 나오는 핵심적인 수식들을 직접 코드로 구현해 볼 수 있게 빈칸을 뚫어 놓았기 떄문에 연습하기 좋다. 

https://github.com/learning3d/assignment4

GitHub - learning3d/assignment4: 3D Gaussian Splatting and Diffusion Guided Optimization

 

환경 세팅

 

pip install --no-build-isolation git+https://github.com/NVlabs/nvdiffrast.git
# in assignment4
pip install --no-deps $(grep -v nvdiffrast requirements.txt | tr '\n' ' ')

 

파일 구조

assignment4/
├─ README.md
├─ requirements.txt
├─ Q1/
│  ├─ data/                # 예제 데이터(PLY/포인트클라우드 등)
│  ├─ data_utils.py
│  ├─ data_utils_harder_scene.py
│  ├─ model.py
│  ├─ render.py
│  ├─ train.py
│  ├─ train_harder_scene.py
│  └─ unit_test_gaussians.py
├─ Q2/
│  ├─ Q21_image_optimization.py
│  ├─ Q22_mesh_optimization.py
│  ├─ Q23_nerf_optimization.py
│  ├─ SDS.py
│  ├─ activation.py
│  ├─ implicit.py
│  ├─ meshutils.py
│  ├─ optimizer.py
│  ├─ data/                # 메쉬(cow) 리소스
│  ├─ freqencoder/         # 주파수 인코더 확장 모듈
│  ├─ gridencoder/         # 그리드 인코더 확장 모듈
│  ├─ nerf/                # NeRF 관련 모듈(모델/렌더러/유틸)
│  └─ raymarching/         # 레이마칭 확장 모듈
├─ ref_output/             # 참고 출력 예시
└─ .gitignore

 

• Q1
  • data_utils.py : 데이터 전처리, 카메라 파라미터 처리 등
  • data_utils_harder_scene.py : 난이도 높은 씬 용 데이터 유틸
  • model.py : 3D 가우시안 표현, 투영/가중치 계산, Splatting renderer 등 핵심 로직
  • render.py : 사전 학습된 가우시안으로 렌더링 실행 스크립트
  • train.py : 기본 씬 학습 파이프라인, Optimizer / loss 정의
  • train_harder_scene.py : 더 어려운 씬 학습 (랜덤 초기화, 추가 실험 옵션)
  • unit_test_gaussians.py : 가우시안 수학 함수 단위 테스트
• Q2
  • SDS.py : Score Distillation Sampling 손실 구현
  • Q21_image_optimization.py : 텍스트 프롬프트 기반 2D 이미지 최적화 실험
  • Q22_mesh_optimization.py : 고정 기하 텍스처 최적화 파이프라인
  • freqencoder/, gridencoder/ : 커스텀 인코더 빌드/래퍼 모듈

 

자료형 (model.py)

 

class Gaussians

Gaussians 클래스는 가우시안 파라미터 (학습 변수) 관리 및 수학 계산을 지원한다. N개의 3D Gaussian을 파라미터 텐서 묶음으로 관리하고, 학습 시에는 이 안의 텐서들이 Optimizer의 파라미터가 된다.

• means : (N,3)
  • 각 Gaussian의 3D 중심
• pre_act_scales : isotrophic이면 (N, 1), anisotrophic이면 (N,3)
  • 로그 스케일 형태로 저장됨
• pre_act_quats : (N, 4)
  • 회전 쿼터니언 
  • isotrohpic 모드에서는 회전이 의미가 없어서 최적화 대상이 아님 
• colours : (N,3)
  • Gaussian의 RGB 기여도
• pre_act_opacities : (N,)
  • 시그모이드 전 값 -> 시그모이드 후에는 (0,1) opacity 가 됨

초기화 방식

• init_types = "gaussians"
  • .ply에서 pretrained gaussian을 로드 (_load_gaussians)
• init_types = "random"
  • 평균/스케일/색 등을 랜덤 초기화
• init_types = "points"
  • .npy 포인트 클라우드 로드
  • Means를 그 점들로 초기화

렌더링 수학 메소드

• apply_activations(pre_act_quats, pre_act_scales, pre_act_opacities)
  • scales  = exp(pre_act_scales)
  • quats = normalize(pre_act_quats)
  • opacities = sigmoid (pre_act_opacities)
• compute_cov_3D
  • 3D Gaussian covariance를 계산 (Paper Eq.6)
• compute_cov_2D
  • 3D covariance를 카메라로 투영해 2D covariance 계산 (Paper Eq. 5)
• compute_means_2D
  • 3D mean을 카메라로 투영해 픽셀 좌표계의 2D mean으로 변환

class Scene

이 Gaussian을 어떤 카메라에서 렌더링할 것인가를 담당하는 렌더링 파이프라인 컨트롤러. Scen.render()

• render(camera, per_splat=-1, img_size=(W,H), bg_colour=(...))
  • 깊이 계산 / 정렬 준비
    • z < 0 제거  + z 기준 오름차순 정렬 인덱스 생성
  • 정렬된 순서로 Gaussian 파라미터 재배열
  • activation 적용
  • 배경 합성
• splat
  • 주어진 N개 Gaussian을 이미지 평면에 뿌려서 최종 결과를 만듦
    • 2D 파라미터 계산
    • Alpha 맵 계산
    • transmittance 계산
    • 색/깊이/마스크 합성
• compute_alphas
  • 내부에서 모든 픽셀 좌표를 생성 
• compute_depth_values

 

과제에서 구현해야 하는 것

1. Class Gaussian의 멤버 함수

compute_cov_3D

• 역할
  • 3D Gaussian의 covariance를 만드는 부분 (3x3)
• 왜 필요한지?
  • 3D에서의 가우시안 모양이 2D로 투영될 때 알파/가중치에 직접 영향

        # HINT: Are quats ever used or optimized for isotropic gaussians? What will their value be?
        # Based on your answers, can you write a more efficient code for the isotropic case?
        if self.is_isotropic:

            ### YOUR CODE HERE ###
            s2 = (scales.squeeze(-1) ** 2).to(self.device)
            cov_3D = torch.zeros((len(s2), 3, 3), device=self.device, dtype=scales.dtype)
            cov_3D[:, 0, 0] = s2
            cov_3D[:, 1, 1] = s2
            cov_3D[:, 2, 2] = s2

        # HINT: You can use a function from pytorch3d to convert quaternions to rotation matrices.
        else:

            ### YOUR CODE HERE ###
            q = quats
            q = torch.nn.functional.normalize(q, dim=-1)
            qw, qx, qy, qz = q[:, 0], q[:, 1], q[:, 2], q[:, 3]

            two_s = 2.0
            R = torch.empty((q.shape[0], 3, 3), device=q.device, dtype=q.dtype)
            R[:, 0, 0] = 1 - two_s * (qy * qy + qz * qz)
            R[:, 0, 1] = two_s * (qx * qy - qz * qw)
            R[:, 0, 2] = two_s * (qx * qz + qy * qw)
            R[:, 1, 0] = two_s * (qx * qy + qz * qw)
            R[:, 1, 1] = 1 - two_s * (qx * qx + qz * qz)
            R[:, 1, 2] = two_s * (qy * qz - qx * qw)
            R[:, 2, 0] = two_s * (qx * qz - qy * qw)
            R[:, 2, 1] = two_s * (qy * qz + qx * qw)
            R[:, 2, 2] = 1 - two_s * (qx * qx + qy * qy)

            s2 = scales ** 2
            S = torch.zeros((len(scales), 3, 3), device=scales.device, dtype=scales.dtype)
            S[:, 0, 0] = s2[:, 0]
            S[:, 1, 1] = s2[:, 1]
            S[:, 2, 2] = s2[:, 2]

            cov_3D = torch.bmm(R, torch.bmm(S, R.transpose(1, 2)))

compute_cov_2D

• 역할
  • 3D covariance를 카메라에 투영한 2D covariance 계산
• 왜 필요한지?
  • 2D 가우시안의 퍼짐(타원 모양)이 정해져야 픽셀별 gaussian 값을 계산 가능
• 내부적으로 채워야 하는 것
  • Jacobian J (N,2,3) : 투영의 미분
  • World -> Camera 회전 W (N,3,3)
  • cov_3D (N,3,3)

compute_means_2D

• 역할
  • 3D mean을 픽셀 좌표계 2D mean으로 투영
• 왜 필요한지?
  • 2D Gaussian의 중심이 그려질 위치를 결정

evaluate_gaussian_2D

• 역할
  • 각 픽셀 위치에서 각 gaussianDml Exponent(power) 계산
• 왜 필요한지? 
  • alpha = opacity * exp(power)

2. Class Scene의 멤버 함수

compute_depth_values

• 각 3D Gaussian이 카메라에서 얼마나 앞/뒤에 있는지 Depth 계산

get_idxs_to_filter_and_sort

• 카메라 뒤쪽 가우시안 제거

compute_transmittance

• 앞 Gaussian에 의해 이미 가려진 정도 구함

splat

• 렌더링의 최종 합성 단계
  • 2D 파라미터 계산 (평균/공분산)
  • alpha/transmittance 계산
  • 색/깊이/실루엣을 가중합으로 생성

이렇게 쭉 하고  python render.py --device cuda 을 실행하면 약 3분 동안 7GB 정도의 메모리를 쓰면서

이런 모델을 만든다.