3D Gaussian splat PLY 파일 합치기/압축하기 (How to merge/compress 3DGS PLYs)

3D gaussian splatting으로 만든 결과물, ply 파일 여러개를 갖고 동시에 편집하는 일이 종종 있다. 예를 들어, scene ply와 object ply를 합쳐서 공간 안에 물체를 넣은 듯한 결과물을 만들고 싶을 때다. 

 

GUI 로는 아주 쉽게 할 수 있다. https://superspl.at/editor/ 여기에 모든 파일을 drag and drop으로 던진 뒤, 원하는 편집을 마친 뒤 내보내기를 하면 하나로 합쳐져서 저장된다.

 

문제는 이걸 코드로 작성해서 자동으로 하고 싶을 때였다. 이게 javascript 코드다 보니까 코드를 봐선 뭔지 모르겠고 그래서 직접 python으로 만들어봤는데 사소하지만 꽤 귀찮은 작업이었기에 메모해둔다. 

 

import torch
import numpy as np
from plyfile import PlyData, PlyElement
from gsplat import export_splats

def load_ply(ply_path):
    ply = PlyData.read(ply_path)
    vertex = ply['vertex'].data
    names = vertex.dtype.names

    if 'x' in names and 'f_rest_44' in names:
        print("[INFO] Detected float format")

        def extract_stack(fields):
            return np.stack([vertex[f] for f in fields], axis=1)

        vertex_data = {
            'position': extract_stack(['x', 'y', 'z']),
            'f_dc': extract_stack(['f_dc_0', 'f_dc_1', 'f_dc_2']),
            'f_rest': extract_stack([f'f_rest_{i}' for i in range(45)]),
            'opacity': vertex['opacity'],
            'scale': extract_stack(['scale_0', 'scale_1', 'scale_2']),
            'rotation': extract_stack(['rot_0', 'rot_1', 'rot_2', 'rot_3']),
        }
        return vertex_data

    else:
        raise ValueError("Unsupported .ply format. Use uncompressed ply file.")

def merge_ply(s1, s2):
    merged = {}
    for key in s1.keys():
        merged[key] = np.concatenate([s1[key], s2[key]], axis=0)
    return merged

def save_ply(vertex_data, save_path):
    N = vertex_data['position'].shape[0]

    # Define dtype in the exact order as in the header
    dtype = [
        ('x', 'f4'), ('y', 'f4'), ('z', 'f4'),
        ('f_dc_0', 'f4'),
    ] + [(f'f_rest_{i}', 'f4') for i in range(45)] + [
        ('f_dc_1', 'f4'), ('f_dc_2', 'f4'),
        ('opacity', 'f4'),
        ('scale_0', 'f4'), ('scale_1', 'f4'), ('scale_2', 'f4'),
        ('rot_0', 'f4'), ('rot_1', 'f4'), ('rot_2', 'f4'), ('rot_3', 'f4'),
    ]

    vertex_array = np.empty(N, dtype=dtype)

    # Fill the fields
    vertex_array['x'] = vertex_data['position'][:, 0]
    vertex_array['y'] = vertex_data['position'][:, 1]
    vertex_array['z'] = vertex_data['position'][:, 2]

    vertex_array['f_dc_0'] = vertex_data['f_dc'][:, 0]
    for i in range(45):
        vertex_array[f'f_rest_{i}'] = vertex_data['f_rest'][:, i]
    vertex_array['f_dc_1'] = vertex_data['f_dc'][:, 1]
    vertex_array['f_dc_2'] = vertex_data['f_dc'][:, 2]

    vertex_array['opacity'] = vertex_data['opacity']

    vertex_array['scale_0'] = vertex_data['scale'][:, 0]
    vertex_array['scale_1'] = vertex_data['scale'][:, 1]
    vertex_array['scale_2'] = vertex_data['scale'][:, 2]

    vertex_array['rot_0'] = vertex_data['rotation'][:, 0]
    vertex_array['rot_1'] = vertex_data['rotation'][:, 1]
    vertex_array['rot_2'] = vertex_data['rotation'][:, 2]
    vertex_array['rot_3'] = vertex_data['rotation'][:, 3]

    # Save as binary little endian .ply
    ply_element = PlyElement.describe(vertex_array, 'vertex')
    PlyData([ply_element], text=False).write(save_path)

    print(f"[INFO] Saved to {save_path} with {N} splats")


def process(subject_path, background_path, save_path):
    subject = load_ply(subject_path)
    background = load_ply(background_path)
    merged = merge_ply(subject, background)
    key_map = {
        "position": "means",
        "opacity": "opacities",
        "rotation": "quats",
        "f_dc": "sh0",
        "f_rest": "shN",
        "scale": "scales",
    }
    splats = {}
    num_splats = np.shape(merged["position"])[0]
    for key in list(merged.keys()):
        value = merged[key]
        new_key = key_map[key]
        if "sh" not in new_key:
            splats[new_key] = torch.from_numpy(value).cuda()
        else: # spherical harmonic
            splats[new_key] = torch.from_numpy(value).reshape(num_splats, -1, 3).cuda()

    export_splats(**splats, format="ply_compressed", save_to=save_path)

# 예시 경로

file1 = "/home/jseob/Downloads/blender_results/135_heavy.ply"
file2 = "/home/jseob/Downloads/blender_results/korea.ply"
file3 = "/home/jseob/Downloads/blender_results/merged_light.ply"

process(file1, file2, file3)
print("")

 

그냥 코드를 통째로 적어둔다. 중요한 건 그냥 key만 제대로 맞춰서 불러온 다음 concat해주고 다시 저장하면 된다. key가 뭔지 일일이 확인하는게 귀찮아서 문제다.

 

오히려 저장하는게 더 문제다. ply 파일은 byte로 저장하는 포맷이다보니 헤더를 정확히 명시해줘야 하고 띄어쓰기 오류나 byte 단위의 밀려쓰기가 있으면 파일이 통째로 다 깨진다. 근데 3DGS는 opacity, rotation, scale 같은 일반적으로 ply에서 쓰는 정보 외에도 형식을 맞춰야 하기 때문에 이걸 직접하는건 굉장히 번거롭다. 따라서 gsplat 에서 미리 짜둔 코드를 사용하는 걸 추천. 

pip install gsplat

 

def process(subject_path, background_path, save_path):
    subject = load_ply(subject_path)
    background = load_ply(background_path)
    merged = merge_ply(subject, background)
    key_map = {
        "position": "means",
        "opacity": "opacities",
        "rotation": "quats",
        "f_dc": "sh0",
        "f_rest": "shN",
        "scale": "scales",
    }
    splats = {}
    num_splats = np.shape(merged["position"])[0]
    for key in list(merged.keys()):
        value = merged[key]
        new_key = key_map[key]
        if "sh" not in new_key:
            splats[new_key] = torch.from_numpy(value).cuda()
        else: # spherical harmonic
            splats[new_key] = torch.from_numpy(value).reshape(num_splats, -1, 3).cuda()

    export_splats(**splats, format="ply_compressed", save_to=save_path)

 

이 대목에서 볼 수 있듯이 key를 다시 gsplat에 맞도록 remap 한 번 해주고 torch tensor화 + reshape만 섞어서 gsplat.export_splats를 통과시키면 된다. 

 

이 때 ply_compressed로 입력할 시 byte compression을 적용해서 supersplat 전용 형식으로 압축된다. (supersplat viewer 외에서는 안 열림)

 

압축이 필요하지 않을 경우 그냥 "ply"로 입력하면 된다.