| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| import numpy as np |
| from ..octree import DfsOctree as Octree |
|
|
|
|
| class Strivec(Octree): |
| def __init__( |
| self, |
| resolution: int, |
| aabb: list, |
| sh_degree: int = 0, |
| rank: int = 8, |
| dim: int = 8, |
| device: str = "cuda", |
| ): |
| assert np.log2(resolution) % 1 == 0, "Resolution must be a power of 2" |
| self.resolution = resolution |
| depth = int(np.round(np.log2(resolution))) |
| super().__init__( |
| depth=depth, |
| aabb=aabb, |
| sh_degree=sh_degree, |
| primitive="trivec", |
| primitive_config={"rank": rank, "dim": dim}, |
| device=device, |
| ) |
|
|
