make_ply_sperate_frames.py

import os
import shutil

# import ipdb
import numpy as np
from plyfile import PlyData, PlyElement

from lib.config import cfg
from lib.datasets.dataset import Dataset

# from lib.models.gaussian_model import GaussianModel
from lib.models.scene import Scene
from lib.models.street_gaussian_model import StreetGaussianModel


def inverse_opacity(x):
    return np.log(x / (1 - x))


def inverse_scale(x):
    return np.log(x)


if __name__ == '__main__':
    # 初始化配置和数据集
    dataset = Dataset()
    gaussians = StreetGaussianModel(dataset.scene_info.metadata)
    scene = Scene(gaussians=gaussians, dataset=dataset)
    
    # 获取训练和测试相机，并按 ID 排序
    train_cameras = scene.getTrainCameras()
    test_cameras = scene.getTestCameras()
    cameras = train_cameras + test_cameras
    cameras = sorted(cameras, key=lambda x: x.id)
    
    # 获取所有唯一的 frame_ids
    frame_ids = sorted(set(camera.meta['frame_idx'] for camera in cameras))
    
    # 设置 Gaussians 的可见性
    keys_list = [key for key in gaussians.model_name_id.keys() if key != 'background']
    gaussians.set_visibility(keys_list)
    # 如果你希望仅设置为背景，可使用以下代码：
    # gaussians.set_visibility(["background"])
    
    # 遍历所有 frame_ids 及其对应的相机
    for frame_id in frame_ids:
        # 找到当前 frame_id 对应的所有相机（假设每个 frame_id 可能有多个相机）
        viewpoint_cameras = [camera for camera in cameras if camera.meta['frame_idx'] == frame_id]
        
        if not viewpoint_cameras:
            print(f'警告: 找不到 frame_idx 为 {frame_id} 的相机，跳过此帧。')
            continue  # 跳过没有对应相机的帧
        
        # 仅处理每个 frame_id 的第一个相机
        viewpoint_camera = viewpoint_cameras[0]
        print(f'正在处理 frame_id: {frame_id}, camera_id: {viewpoint_camera.id}')
        
        # 解析当前相机
        gaussians.parse_camera(camera=viewpoint_camera)
        
        # 获取 XYZ 数据
        xyz = gaussians.get_xyz.detach().cpu().numpy()
        normals = np.zeros_like(xyz)  # 如果有法线数据，可以替换此行
        
        # 获取特征数据
        f = gaussians.get_features.detach().transpose(1, 2).contiguous()  # [n, 3, sh_degree]
        f_dc = f[..., :1].flatten(start_dim=1).cpu().numpy()
        f_rest = f[..., 1:].flatten(start_dim=1).cpu().numpy()
        
        # 获取不透明度，并进行反变换
        opacities = gaussians.get_opacity.detach().cpu().numpy()
        opacities = np.clip(opacities, a_min=1e-6, a_max=1.-1e-6)
        opacities = inverse_opacity(opacities)
        
        # 获取缩放和旋转数据，并进行反变换
        scale = gaussians.get_scaling.detach().cpu().numpy()
        scale = inverse_scale(scale)
        
        rotation = gaussians.get_rotation.detach().cpu().numpy()
        
        # 定义 PLY 文件的字段
        field_names = ['x', 'y', 'z', 'nx', 'ny', 'nz']
        for i in range(f_dc.shape[1]):
            field_names.append(f'f_dc_{i}')
        for i in range(f_rest.shape[1]):
            field_names.append(f'f_rest_{i}')
        field_names.append('opacity')
        for i in range(scale.shape[1]):
            field_names.append(f'scale_{i}')
        for i in range(rotation.shape[1]):
            field_names.append(f'rot_{i}')
        
        dtype_full = [(field, 'f4') for field in field_names]
        
        # 创建结构化数组
        elements = np.empty(xyz.shape[0], dtype=dtype_full)
        attributes = np.concatenate((xyz, normals, f_dc, f_rest, opacities, scale, rotation), axis=1)
        elements[:] = list(map(tuple, attributes))
        
        # 定义保存路径
        save_dir = os.path.join(cfg.model_path, 'point_cloud_for_mesh')
        pointcloud_dir = os.path.join(save_dir, 'point_cloud')
        os.makedirs(pointcloud_dir, exist_ok=True)
        
        # 复制必要的文件到当前帧的保存目录
        shutil.copyfile(os.path.join(cfg.model_path, 'cameras.json'), os.path.join(save_dir, 'cameras.json'))
        shutil.copyfile(os.path.join(cfg.model_path, 'cfg_args'), os.path.join(save_dir, 'cfg_args'))
        shutil.copyfile(os.path.join(cfg.model_path, 'input.ply'), os.path.join(save_dir, 'input.ply'))
        
        # 保存每个 frame_id 相关的 PLY 点云文件
        ply_element = PlyElement.describe(elements, 'vertex')
        ply_data = PlyData([ply_element])
        ply_filename = os.path.join(pointcloud_dir, f'point_cloud_{frame_id:06d}.ply')  # 文件名与 frame_id 相关
        ply_data.write(ply_filename)
        
        print(f'已保存点云到: {ply_filename}')
    
    print('所有帧处理完成。')