Custom Neural Rendering Viewer 개발기

들어가며

최근 2D Gaussian Splatting 에 대한 remote viewer 를 side project 로 작업하여 github 에 배포하였다.

처음 작업해본 side project 라 걱정이 많았지만, 나름 star 도 많이 수집하고, official repo 에 언급되는 등 소소하게 만족스러운 경험이었다.

잊기 전에 Neural Rendering 관련 viewer 개발하면서 얻게 된 교훈들과 느낀 점을 공유해보려 한다.

1. Viewer 설계하기

1.1. Original 2D GS projects 의 기능 활용하기

Rendering 함수나 Gaussian model 에 대한 정의를 viewer 에 적합하도록, 혹은 구현하려는 transform 이나 삭제, mesh export 기능 등에 호환되도록 바꿔야 할 필요가 있었다.

하지만 original project 자체가 아직 live project 이며, 미리 정의된 class 를 상속 받거나 overloading 해서 개발하는 것이 편하기 때문에 original project 를 submodule 로써 추가하여 viewer 개발과 독립적으로 사용할 수 있도록 하였다.

이를 위해

git submodule add <original-2dgs-repo-url> 2d_gaussian_splatting

로 projects 에 submodule 을 추가하였으며, 다음 구조로 개발을 진행하였다.

my_viewer/
├── 2d_gaussian_splatting/
│   ├── utils/
│   │   └── __init__.py
│   │   └── some_utils.py
│   ├── scene/
│   │   └── __init__.py
│   │   └── GaussianModel.py
│   ├── ...
├── internal/
│   └── viewer/
│       ├── client.py
│       └── model.py 
│       └── ui
│           ├── edit.py
│           ├── transform.py
├── viewer.py

1.2. Logic 설계하기

Viewer 구현에는

1. local viewer (3D GS's SIBR viewer 류)
2. Remote Viewer (NeRFStudio viewer 류)

의 선택지가 있었는데, 보통 display 장치가 없는 GPU server 에 docker / conda 등으로 작업할 일이 많았기 때문에 자연스럽게 Remote Viewer 를 개발하고자 했다.

직접적으로 참고한 project 는 "NeRFStudio" 와 "3D Gaussian Splatting PyTorch Lightning" 인데, 둘 모두

• Web Server: remote GPU server
• Web Client: local web browser

의 구성으로 되어 있다.

즉 rendering (여기선 rasterization) 의 작업은 server 에서 진행하되, client (browser) 내의 mouse / keyboard interaction 등을 통해 camera parameter 나 rendering option 을 load 하는 구성이다.

Viser Project 에서 client camera 의 wxyz quaternion 을 제공하기 때문에 camera param 을 계산하기 위해선 다음의 pseudo code 정도만 구현하면 되며,

def make_camera(self, wxyz, position):
	R = quaternion_to_SO3(wxyz=wxyz)
	c2w = torch.eye(4)
	c2w[:3, :3] = R
	c2w[:3, 3] = position
    return c2w

(make camera function 은 실제로는 단순히 cam2world projection 뿐만 아니라 near, far plane 을 고려한 ndc parameterization 이 구현되어야 한다)

camera refresh 가 일어날 때마다 이 camera 계산을 update 하여 render 함수를 호출하도록 하는 기본적인 토대를 설계했다.

client.camera.on_update
def _(cam: viser.CameraHandle) -> None:
	with self.client.atomic():
		self.render_trigger.set()

즉,

1. Client (browser) 에서 camera refresh (mouse / keyboard / rendering option change) 등이 일어난다.
2. Server 에 client 에서 얻은 정보 (camera param / rendering option 등) 을 넘겨서 rendering 을 진행한다
3. rendering 한 image 를 client (browser) 에 보여준다

의 logic 으로 요약할 수 있겠다.

2. 구현 이모저모

2.1. Server 와 Client

server 와 client thread 가 독립적으로 실행되기 때문에, thread 간 통신할 때 sync 를 맞추는 것 또한 중요했다.

따라서 적절히 atomic 옵션을 통해서 contiguous 하고 synchronized 된 data 를 유지할 필요가 있었는데, 가령 예를 들어 client camera 의 camera param 을 renderer 로 넘겨주는 다음과 같은 부분이다.

with self.client.atomic():
	camera = self.make_camera(self.client.camera)
    image = self.render_image(camera)

이 외에도 통신 과정에서의 비동기 문제가 발생하는 경우가 있었기 때문에 최대한 client 에서의 option 편집이 실제 rendering 과정과 충돌을 일으키지 않게 하려 노력했다.

2.2. Edit & Transform 을 위한 Model 정의 변경

직접적으로 참고한 3D GS PL 에서 gaussian ply 에 대한 편집 기능을 제공하긴 하지만, 이 기능의 작동 방식이 일반적인 직관과 약간 거리가 있었다.

그 이유는 바로 transform 에 대한 기능이 다음과 같이 구현되어 있었기 때문인데,

def transform_with_vectors(model, rescale_factor, rotation, position):
        xyz = model.xyz
        scaling = model.scaling
        rotation = model.rotation
        features = model.features 

        # transform 
        xyz, scaling, rotation, features = TransformGaussian(xyz, 
        													scaling, 
                                                            rotation, 
                                                            features, 
                                                            rescale_factor, 
                                                            rotation, position)
                                                            
        model.xyz = xyz
        model.scaling = scaling
        model.rotation = rotation
        model.features = new_features

얼핏 봐선 문제가 없어보이지만, 이 구현은 직관과 어긋나는 큰 문제점이 있었다.

예를 들어 모델에 대해 다음 두 단계에 걸쳐서 변환을 진행하면,

1. viewer 상에서 euler x-axis 로 30도 rotation 입력
2. 30도 입력했던 숫자를 50으로 바꿈

직관적으로는 원래 model 에서 50도 회전한 결과가 나타나야하지만, 위 구현을 따르면 80도 변환을 하게 된다.

따라서 GaussianModel 에 original attributes 들을 추가적으로 저장하여, transform 이 일어날 때는 현재 model 의 xyz, scale 값 등이 아니라 original 값에 일어나도록 구현하였다.

def transform_with_vectors(model, rescale_factor, rotation, position):
        xyz = model.org_xyz
        scaling = model.org_scaling
        rotation = model.org_rotation
        features = model.org_features 

        # transform 
        xyz, scaling, rotation, features = TransformGaussian(xyz, 
        													scaling, 
                                                            rotation, 
                                                            features, 
                                                            rescale_factor, 
                                                            rotation, position)
                                                            
        model.xyz = xyz
        model.scaling = scaling
        model.rotation = rotation
        model.features = new_features

변환된 상태를 default 상태로 model 을 저장하고 싶을 때에는 'Set to Default' 기능을 만들어 set to default 가 호출되었을 때만 original feature 값들을 업데이트 하도록 하였다.

Blog 를 쓰면서 생각하게된 사실인데, 위 방식은 original attributes 들을 추가적으로 저장해야하기 때문에 memory consumption 이 늘어난다.
대신에 last transform 정보만 저장하고 current transform 을 last transform 과 비교해서 변화한 정도만 실제 transform 이 일어나도록 코드를 수정하는 것이 좋을 것 같다.

2.3. Multi-thread 활용하기

Client, server 가 독립적인 thread 에서 실행되는 것 이외에도 특정 기능은 thread 를 하나 더 열어서 작업하도록 구현할 필요가 있었다.

예를 들어 train 상황을 볼 수 있는 train_w_viewer.py 등을 구현할 때는 viewer thread 를 독립적으로 관리해야 했으며, 특정 iteration 마다 viewer 의 gaussian model 을 training current state 를 반영하도록 구현하였다.

(in train_w_viewer.py)

class Trainer:
    ...
    
    def set_viewer(self, viewer_attributes):
        self.viewer = Viewer(viewer_attributes)
        
    def viewer_thread(self):
        self.viser_viewer.start()
        
    def training(~):
    	... 
        self.set_viewer(args)
        self.viewer._get_training_gaussians(gaussians)
        
        for iteration in range(first_iter, opt.iterations + 1):   
        	...
            if iteration % 100 == 0:
                with self.lock:
                    self.viser_viewer._get_training_gaussians(gaussians)
    	

즉 trainer setting 중에 viewer thread 를 별도로 호출하여 실행하고, viewer thread 에 gaussian model 정보를 update 할 때마다 동기화를 위하여 self.lock 을 이용해 asynchronize problem 을 방지하였다.

이 외에도 2D GS 에서 제공하는 TSDF mesh recon 을 viewer 내에서 실행하고, 결과를 visualize 하는 기능들도 별도의 thread 를 열어서 작업하고 보여줄 수 있도록 작업하였다.

마치며

AI Researcher 로서 논문을 위한 구현 정도가 아니면 한 가지 프로젝트를 진득히 코딩하는 경우가 드물었는데, side project 구현하면서 오랜만에 '코딩'이란 행위에 즐거움을 느꼈다.

또한 학부 수업 이외엔 실제로 체감할 일이 없었던 비동기 문제라든지, multi thread 라든지 등을 간단한 수준에서라도 프로젝트에 활용해보고 문제점을 느껴보는 것도 좋은 경험이었다.

Radiance Fields / Neural Rendering 분야에서 연구를 지속하고 싶기 때문에, custom 으로 viewer 를 만들어본 이번 경험이 꽤나 값진 경험으로 오래 남을 것 같다.