Spaces:
Running
on
Zero
Running
on
Zero
File size: 1,297 Bytes
917a889 |
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 |
from typing import *
import numpy as np
import torch
from ..modules import sparse as sp
from ..representations import Voxel
from .render_utils import render_video
def pca_color(feats: torch.Tensor, channels: Tuple[int, int, int] = (0, 1, 2)) -> torch.Tensor:
"""
Apply PCA to the features and return the first three principal components.
"""
feats = feats.detach()
u, s, v = torch.svd(feats)
color = u[:, channels]
color = (color - color.min(dim=0, keepdim=True)[0]) / (color.max(dim=0, keepdim=True)[0] - color.min(dim=0, keepdim=True)[0])
return color
def vis_sparse_tensor(
x: sp.SparseTensor,
num_frames: int = 300,
):
assert x.shape[0] == 1, "Only support batch size 1"
assert x.coords.shape[1] == 4, "Only support 3D coordinates"
coords = x.coords.cuda().detach()[:, 1:]
feats = x.feats.cuda().detach()
color = pca_color(feats)
resolution = max(list(x.spatial_shape))
resolution = int(2**np.ceil(np.log2(resolution)))
rep = Voxel(
origin=[-0.5, -0.5, -0.5],
voxel_size=1/resolution,
coords=coords,
attrs=color,
layout={
'color': slice(0, 3),
}
)
return render_video(rep, colors_overwrite=color, num_frames=num_frames)['color']
|