[3DGS 구현] 6. 이미지 전처리 및 GPU로딩

 

이전 글

[3DGS 구현] 5. gaussian 초기화 및 SIBR뷰어를 통한 시각화

[3DGS 구현] 5. gaussian 초기화 및 SIBR뷰어를 통한 시각화

---

 

학습을 진행할 때 원본과 얼마나 다른지를 계산해서 손실을 구해야한다.

이때 필요한 것은 원본이미지와, 이 원본 이미지를 찍은 카메라의 정보들, 현재 3D모델을 이 카메라 화면에 투영하는 행렬이다.

 

먼저 몇가지 연산에 도움을 주는 함수를 정의하자.

작업폴더 안에 utils.py라는 이름으로 새로 파일을 만들었고, 이 안에는 여러 계산식들을 넣어놓을 예정이다.

 

<utils.py에 추가>

qvec2rotmat : 4차원 쿼터니안을 3 x 3 회전행렬로 변환. (관련 글: 쿼터니안과 회전행렬)

get_world_view_matrix : 월드 좌표계를 카메라를 기준으로한 좌표계로 변환하는 행렬을 계산.

get_projection_matrix : 카메라 좌표계에서 해당 카메라의 센서 평면(2D)로 투영시키는 행렬을 계산.

# utils.py
import torch
import numpy as np

def qvec2rotmat(qvec):
    return np.array([
        [1 - 2 * qvec[2]**2 - 2 * qvec[3]**2,
         2 * qvec[1] * qvec[2] - 2 * qvec[0] * qvec[3],
         2 * qvec[1] * qvec[3] + 2 * qvec[0] * qvec[2]],
        [2 * qvec[1] * qvec[2] + 2 * qvec[0] * qvec[3],
         1 - 2 * qvec[1]**2 - 2 * qvec[3]**2,
         2 * qvec[2] * qvec[3] - 2 * qvec[0] * qvec[1]],
        [2 * qvec[1] * qvec[3] - 2 * qvec[0] * qvec[2],
         2 * qvec[2] * qvec[3] + 2 * qvec[0] * qvec[1],
         1 - 2 * qvec[1]**2 - 2 * qvec[2]**2]
    ])

def get_world_view_matrix(R, T):
    rt = np.zeros((4, 4))
    rt[:3, :3] = R
    rt[:3, 3] = T
    rt[3, 3] = 1.0
    return rt

def get_projection_matrix(znear, zfar, fovX, fovY):
    tanHalfFovY = np.tan(fovY / 2)
    tanHalfFovX = np.tan(fovX / 2)

    top = tanHalfFovY * znear
    bottom = -top
    right = tanHalfFovX * znear
    left = -right

    P = torch.zeros(4, 4)
    z_sign = 1.0 # 3DGS는 +Z 방향을 바라봄

    P[0, 0] = 2.0 * znear / (right - left)
    P[1, 1] = 2.0 * znear / (top - bottom)
    P[0, 2] = (right + left) / (right - left)
    P[1, 2] = (top + bottom) / (top - bottom)
    
    P[3, 2] = z_sign
    P[2, 2] = z_sign * zfar / (zfar - znear)
    P[2, 3] = -(zfar * znear) / (zfar - znear)
    
    return P

 

이제 Camera클래스라는 것을 만들어보자.

각 정답이미지 당 하나씩의 인스턴스를 가질 수 있게 만들 것이고, 위에서 얘기한 Loss를 계산하는데 필요한 정보들이 모두 들어있게끔 설계했다.

이미지를 GPU로 가져오고, 월드 → 카메라 좌표계로 변환하는데 필요한 행렬을 저장한다.

 

import os
import torch
import numpy as np
from PIL import Image
from tqdm import tqdm
import struct


import config
from utils import qvec2rotmat, get_projection_matrix, get_world_view_matrix

class Camera(torch.nn.Module):
    def __init__(self, R, T, foV_x, foV_y, image, image_name):
        super(Camera, self).__init__()
        self.R = R
        self.T = T
        self.foV_x = foV_x
        self.foV_y = foV_y
        self.image_name = image_name
        
        # 정답 이미지를 GPU로
        self.original_image = image.clamp(0.0, 1.0).cuda()
        self.image_width = self.original_image.shape[2]
        self.image_height = self.original_image.shape[1]

        # 월드에서 카메라 행렬
        self.world_view_transform = torch.tensor(get_world_view_matrix(R, T)).float().transpose(0, 1).cuda()
        self.projection_matrix = get_projection_matrix(znear=0.01, zfar=100.0, fovX=self.foV_x, fovY=self.foV_y).float().transpose(0, 1).cuda()
        self.full_proj_transform = (self.world_view_transform @ self.projection_matrix)
        
        # 월드 기준 카메라 중심
        self.camera_center = self.world_view_transform.inverse()[3, :3]

 

 

이어서 같은 파일에 load_cameras라는 함수도 만들자.

위에서 만든 camera클래스의 객체를 직접 만드는 함수이다.

이때 리사이즈까지 고려해 클래스의 모든 변수 값을 채워서 인스턴스를 생성한다.

 

이전 글에서 보았듯이 images.bin파일에 '월드→카메라'좌표계로 변환할 때 필요한 회전정보(쿼터니언)와 위치정보가 들어있기 때문에 이를 파싱해와야 한다.

또 cameras.bin에는 카메라의 속성 값들이 있다. 역시 파싱해와야한다.

def load_cameras( resolution_scale):
    cameras = []
    
    # cameras.bin
    with open(config.CAMERAS_BIN_PATH, "rb") as f:
        f.read(8) # num_cameras 스킵
        # cam_id, model_id, width, height 스킵
        f.read(24) 
        # f만 가져옴
        focal_length = struct.unpack("<d", f.read(8))[0] 

    # images.bin
    with open(config.IMAGES_BIN_PATH, "rb") as f:
        num_reg_images = struct.unpack("<Q", f.read(8))[0]
        
        for _ in tqdm(range(num_reg_images)):
            #헤더에서 정보뽑기
            _, qw, qx, qy, qz, tx, ty, tz, _ = struct.unpack("<I dddd ddd I", f.read(64))
            
            # 파일명
            image_name = ""
            while True:
                char = f.read(1).decode("utf-8")
                if char == "\0": break
                image_name += char
            
            # 포인트 데이터 스킵
            num_points = struct.unpack("<Q", f.read(8))[0]
            f.seek(num_points * 24, os.SEEK_CUR) 

            # 리사이즈 관련
            pil_image = Image.open(config.IMAGE_PATH+'/images/'+image_name)
            w, h = pil_image.size
            if resolution_scale != 1:
                w, h = int(w / resolution_scale), int(h / resolution_scale)
                pil_image = pil_image.resize((w, h), Image.LANCZOS)
            
            image_tensor = torch.from_numpy(np.array(pil_image)).permute(2, 0, 1).float() / 255.0
            
            # 리사이즈 고려 fov계산
            curr_focal = focal_length / resolution_scale
            fov_x = 2 * np.arctan(w / (2 * curr_focal))
            fov_y = 2 * np.arctan(h / (2 * curr_focal))
            
            
            cameras.append(Camera(
                R=qvec2rotmat(np.array([qw, qx, qy, qz])), 
                T=np.array([tx, ty, tz]), 
                foV_x=fov_x, foV_y=fov_y, 
                image=image_tensor, image_name=image_name
            ))

    return cameras

 

최종적으로 반환되는 cameras는 모든 정답 이미지에 대한 정보를 담고 있다.

예를 들어 bicycle이미지셋의 경우 총 194장의 정답이미지가 있는데, 

cameras[0]부터 cameras[193]까지 저장이 된다.

각 항목 cameras[i]는 클래스의 객체이기 때문에,

이 안의 내용들은 점(.)을 통해 접근한다.

 

original_image : (3,H,W) 크기의 텐서

full_proj_transform : (4,4) 크기의 텐서

camera_center : 3차원

등의 정보가 들어있다.

 

현재까지의 작업폴더 구조는 아래와 같다.

 

다음글에서는 실제로 렌더링을 했을 때 나올 각 픽셀의 값을 구하는 과정을 알아보자.

이 과정은 논문에서 얘기한 Rasterization 기법이 포함되는 중요한 부분이다.