Spaces:
Sleeping
Sleeping
File size: 1,507 Bytes
78d2329 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 | from sklearn.neighbors import NearestNeighbors
import torch
from torch import Tensor
from optgs.scene_trainer.common.gaussian_adapter import RGB2SH
def knn(x: Tensor, K: int = 4) -> Tensor:
x_np = x.cpu().numpy()
model = NearestNeighbors(n_neighbors=K, metric="euclidean").fit(x_np)
distances, _ = model.kneighbors(x_np)
return torch.from_numpy(distances).to(x)
def points_to_gaussians(
points_dict: dict[str, Tensor],
sh_degree: int = 3,
device: torch.device = torch.device("cpu"),
) -> dict[str, Tensor]:
xyz = points_dict["xyz"].clone().to(device)
N = xyz.shape[0]
# color is SH coefficients
rgbs = points_dict["rgb"].clone().to(device) # [N, 3], in [0, 1]
# if sh_degree > 0:
shs = torch.zeros((N, (sh_degree + 1) ** 2, 3), device=device) # [N, K, 3]
shs[:, 0, :] = RGB2SH(rgbs)
sh0 = shs[:, :1, :] # [N, 1, 3]
if sh_degree > 0:
shN = shs[:, 1:, :] # [N, K-1, 3]
else:
shN = None
quats_unnorm = torch.rand((N, 4), device=device) # [N, 4]
scales = points_dict["scales"].clone().to(device) # [N, 3]
scales_raw = torch.log(scales)
opacities = points_dict["opacities"].clone().to(device) # [N,]
opacities_raw = torch.logit(opacities)
return {
"xyz": xyz,
"sh0": sh0, # [N, 1, 3]
"shN": shN, # [N, sh_d-1, 3] or None
"scales_raw": scales_raw,
"rotations_unnorm": quats_unnorm,
"opacities_raw": opacities_raw,
}
|