BUG to be fixed

[liruilong940607]

Rendering is bugged with the following code (I'm looking into it rn):

import math

import imageio
import numpy as np
import torch
from gsplat import ProjectGaussians, RasterizeGaussians

torch.manual_seed(42)

device = torch.device("cuda:0")
num_points = 3

BLOCK_X, BLOCK_Y = 16, 16
fov_x = math.pi / 2.0
H, W = 256, 256
focal = 0.5 * float(W) / math.tan(0.5 * fov_x)
tile_bounds = (
    (W + BLOCK_X - 1) // BLOCK_X,
    (H + BLOCK_Y - 1) // BLOCK_Y,
    1,
)
img_size = torch.tensor([W, H, 1], device=device)
block = torch.tensor([BLOCK_X, BLOCK_Y, 1], device=device)

bd = 2
means = torch.rand(num_points, 3, device=device)
scales = torch.rand(num_points, 3, device=device)
rgbs = torch.rand(num_points, 3, device=device)

u = torch.rand(num_points, 1, device=device)
v = torch.rand(num_points, 1, device=device)
w = torch.rand(num_points, 1, device=device)

quats = torch.cat(
    [
        torch.sqrt(1.0 - u) * torch.sin(2.0 * math.pi * v),
        torch.sqrt(1.0 - u) * torch.cos(2.0 * math.pi * v),
        torch.sqrt(u) * torch.sin(2.0 * math.pi * w),
        torch.sqrt(u) * torch.cos(2.0 * math.pi * w),
    ],
    -1,
)
opacities = torch.ones((num_points, 1), device=device)

viewmat = torch.tensor(
    [
        [1.0, 0.0, 0.0, 0.0],
        [0.0, 1.0, 0.0, 0.0],
        [0.0, 0.0, 1.0, 8.0],
        [0.0, 0.0, 0.0, 1.0],
    ],
    device=device,
)
background = torch.zeros(3, device=device)


xys, depths, radii, conics, num_tiles_hit, cov3d = ProjectGaussians.apply(
    means,
    scales,
    1,
    quats,
    viewmat,
    viewmat,
    focal,
    focal,
    W / 2,
    H / 2,
    H,
    W,
    tile_bounds,
)
render = RasterizeGaussians.apply(
    xys,
    depths,
    radii,
    conics,
    num_tiles_hit,
    torch.sigmoid(rgbs),
    torch.sigmoid(opacities),
    H,
    W,
    background,
)

canvas = (render * 255.0).detach().cpu().numpy()
imageio.imwrite("ref.png", canvas.astype(np.uint8))