utils/gs.py

import random
import torch
import numpy as np
from .scene import GaussianModel
from .scene.dataset_readers import SceneInfo, getNerfppNorm
from .scene.cameras import Camera
from .ops import focal2fov, fov2focal
from .scene.gaussian_model import BasicPointCloud
from easydict import EasyDict as edict
from PIL import Image

from tqdm.auto import tqdm

def get_blank_gs_bundle(h, w):
    return {
        "camera_angle_x": focal2fov(torch.tensor([w], dtype=torch.float32), w),
        "W": w,
        "H": h,
        "pcd_points": None,
        "pcd_colors": None,
        'frames': [],
    }

def read_cameras_from_optimization_bundle(optimization_bundle, white_background: bool = False):
    cameras = []

    fovx = optimization_bundle["camera_angle_x"]
    frames = optimization_bundle["frames"]

    # we flip the x and y axis to move from PyTorch3D's coordinate system to COLMAP's
    coordinate_system_transform = np.array([-1, -1, 1])

    for idx, frame in enumerate(frames):
        c2w = np.array(frame["transform_matrix"])
        c2w[:3, :3] = c2w[:3, :3] * coordinate_system_transform

        # get the world-to-camera transform and set R, T
        w2c = np.linalg.inv(c2w)
        R = np.transpose(w2c[:3, :3])  # R is stored transposed due to 'glm' in CUDA code
        T = c2w[-1, :3] * coordinate_system_transform

        image = frame["image"]

        im_data = np.array(image.convert("RGBA"))

        bg = np.array([1,1,1]) if white_background else np.array([0, 0, 0])

        norm_data = im_data / 255.0
        arr = norm_data[:,:,:3] * norm_data[:, :, 3:4] + bg * (1 - norm_data[:, :, 3:4])
        image = Image.fromarray(np.array(arr*255.0, dtype=np.byte), "RGB")

        fovy = focal2fov(fov2focal(fovx, image.size[0]), image.size[1])
        FovY = fovy 
        FovX = fovx

        image = torch.Tensor(arr).permute(2,0,1)

        cameras.append(Camera(colmap_id=idx, R=R, T=T, FoVx=FovX, FoVy=FovY, image=image, mask=frame.get("mask", None),
                                gt_alpha_mask=None, image_name='', uid=idx, data_device='cuda'))

    return cameras

class Scene:
    gaussians: GaussianModel

    def __init__(self, traindata, gaussians: GaussianModel, gs_options, shuffle: bool = True):
        self.traindata = traindata
        self.gaussians = gaussians
        
        train_cameras = read_cameras_from_optimization_bundle(traindata, gs_options.white_background)
        
        nerf_normalization = getNerfppNorm(train_cameras)

        pcd = BasicPointCloud(points=traindata['pcd_points'], colors=traindata['pcd_colors'], normals=None)
        
        scene_info = SceneInfo(point_cloud=pcd,
                               train_cameras=train_cameras,
                               test_cameras=[],
                               nerf_normalization=nerf_normalization,
                               ply_path='')

        if shuffle:
            random.shuffle(scene_info.train_cameras)  # Multi-res consistent random shuffling

        self.cameras_extent = scene_info.nerf_normalization["radius"]

        self.train_cameras = scene_info.train_cameras

        bg_color = np.array([1,1,1]) if gs_options.white_background else np.array([0, 0, 0])
        self.background = torch.tensor(bg_color, dtype=torch.float32, device='cuda')

        self.gaussians.create_from_pcd(scene_info.point_cloud, self.cameras_extent)
        self.gaussians.training_setup(gs_options)

    def getTrainCameras(self):
        return self.train_cameras
    
    def getPresetCameras(self, preset):
        assert preset in self.preset_cameras
        return self.preset_cameras[preset]

def run_gaussian_splatting(scene, gs_optimization_bundle):
    torch.cuda.empty_cache()

    return scene

    from random import randint
    from .gaussian_renderer import render as gs_render
    from .scene.utils.loss_utils import l1_loss, ssim

    pbar = tqdm(range(1, gs_options.iterations + 1))
    for iteration in pbar:
        scene.gaussians.update_learning_rate(iteration)

        # Every 1000 its we increase the levels of SH up to a maximum degree
        if iteration % 1000 == 0:
            scene.gaussians.oneupSHdegree()

        # Pick a random Camera
        random_idx = randint(0, len(gs_optimization_bundle["frames"])-1)
        viewpoint_cam = scene.getTrainCameras()[random_idx]

        # Render
        render_pkg = gs_render(viewpoint_cam, scene.gaussians, gs_options, scene.background)
        image, viewspace_point_tensor, visibility_filter, radii = (
            render_pkg['render'], render_pkg['viewspace_points'], render_pkg['visibility_filter'], render_pkg['radii'])

        # Loss
        gt_image = viewpoint_cam.original_image.cuda()
        Ll1 = l1_loss(image, gt_image, reduce=False)
        loss = (1.0 - gs_options.lambda_dssim) * Ll1

        if viewpoint_cam.mask is not None:
            mask = torch.from_numpy(viewpoint_cam.mask).to(loss.device)
        else:
            mask = 1

        loss = (loss * mask).mean()
        loss = loss + gs_options.lambda_dssim * (1.0 - ssim(image, gt_image))
        loss.backward()

        pbar.set_description(f"Loss: {loss.item():.4f}")

        with torch.no_grad():
            # Densification
            if iteration < gs_options.densify_until_iter:
                # Keep track of max radii in image-space for pruning
                scene.gaussians.max_radii2D[visibility_filter] = torch.max(
                    scene.gaussians.max_radii2D[visibility_filter], radii[visibility_filter])
                scene.gaussians.add_densification_stats(viewspace_point_tensor, visibility_filter)

                if iteration > gs_options.densify_from_iter and iteration % gs_options.densification_interval == 0:
                    size_threshold = 20 if iteration > gs_options.opacity_reset_interval else None
                    scene.gaussians.densify_and_prune(
                        gs_options.densify_grad_threshold, 0.005, scene.cameras_extent, size_threshold)
                
                if (iteration % gs_options.opacity_reset_interval == 0 
                    or (gs_options.white_background and iteration == gs_options.densify_from_iter)
                ):
                    scene.gaussians.reset_opacity()

            # Optimizer step
            if iteration < gs_options.iterations:
                scene.gaussians.optimizer.step()
                scene.gaussians.optimizer.zero_grad(set_to_none = True)

    return scene

gs_options = edict({
    "sh_degree": 3,
    "images": "images",
    "resolution": -1,
    "white_background": False,
    "data_device": "cuda",
    "eval": False,
    "use_depth": False,
    "iterations": 0,#250,
    "position_lr_init": 0.00016,
    "position_lr_final": 0.0000016,
    "position_lr_delay_mult": 0.01,
    "position_lr_max_steps": 2990,
    "feature_lr": 0.0,#0.0025,
    "opacity_lr": 0.0,#0.05,
    "scaling_lr": 0.0,#0.005,
    "rotation_lr": 0.0,#0.001,
    "percent_dense": 0.01,
    "lambda_dssim": 0.2,
    "densification_interval": 100,
    "opacity_reset_interval": 3000,
    "densify_from_iter": 10_000,
    "densify_until_iter": 15_000,
    "densify_grad_threshold": 0.0002,
    "convert_SHs_python": False,
    "compute_cov3D_python": False,
    "debug": False,
})