125 lines
3.4 KiB
Python
125 lines
3.4 KiB
Python
from dataclasses import dataclass
|
|
from typing import Optional
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
from einops import rearrange
|
|
|
|
from ..utils import BaseModule
|
|
|
|
|
|
class TriplaneUpsampleNetwork(BaseModule):
|
|
@dataclass
|
|
class Config(BaseModule.Config):
|
|
in_channels: int
|
|
out_channels: int
|
|
|
|
cfg: Config
|
|
|
|
def configure(self) -> None:
|
|
self.upsample = nn.ConvTranspose2d(
|
|
self.cfg.in_channels, self.cfg.out_channels, kernel_size=2, stride=2
|
|
)
|
|
|
|
def forward(self, triplanes: torch.Tensor) -> torch.Tensor:
|
|
triplanes_up = rearrange(
|
|
self.upsample(
|
|
rearrange(triplanes, "B Np Ci Hp Wp -> (B Np) Ci Hp Wp", Np=3)
|
|
),
|
|
"(B Np) Co Hp Wp -> B Np Co Hp Wp",
|
|
Np=3,
|
|
)
|
|
return triplanes_up
|
|
|
|
|
|
class NeRFMLP(BaseModule):
|
|
@dataclass
|
|
class Config(BaseModule.Config):
|
|
in_channels: int
|
|
n_neurons: int
|
|
n_hidden_layers: int
|
|
activation: str = "relu"
|
|
bias: bool = True
|
|
weight_init: Optional[str] = "kaiming_uniform"
|
|
bias_init: Optional[str] = None
|
|
|
|
cfg: Config
|
|
|
|
def configure(self) -> None:
|
|
layers = [
|
|
self.make_linear(
|
|
self.cfg.in_channels,
|
|
self.cfg.n_neurons,
|
|
bias=self.cfg.bias,
|
|
weight_init=self.cfg.weight_init,
|
|
bias_init=self.cfg.bias_init,
|
|
),
|
|
self.make_activation(self.cfg.activation),
|
|
]
|
|
for i in range(self.cfg.n_hidden_layers - 1):
|
|
layers += [
|
|
self.make_linear(
|
|
self.cfg.n_neurons,
|
|
self.cfg.n_neurons,
|
|
bias=self.cfg.bias,
|
|
weight_init=self.cfg.weight_init,
|
|
bias_init=self.cfg.bias_init,
|
|
),
|
|
self.make_activation(self.cfg.activation),
|
|
]
|
|
layers += [
|
|
self.make_linear(
|
|
self.cfg.n_neurons,
|
|
4, # density 1 + features 3
|
|
bias=self.cfg.bias,
|
|
weight_init=self.cfg.weight_init,
|
|
bias_init=self.cfg.bias_init,
|
|
)
|
|
]
|
|
self.layers = nn.Sequential(*layers)
|
|
|
|
def make_linear(
|
|
self,
|
|
dim_in,
|
|
dim_out,
|
|
bias=True,
|
|
weight_init=None,
|
|
bias_init=None,
|
|
):
|
|
layer = nn.Linear(dim_in, dim_out, bias=bias)
|
|
|
|
if weight_init is None:
|
|
pass
|
|
elif weight_init == "kaiming_uniform":
|
|
torch.nn.init.kaiming_uniform_(layer.weight, nonlinearity="relu")
|
|
else:
|
|
raise NotImplementedError
|
|
|
|
if bias:
|
|
if bias_init is None:
|
|
pass
|
|
elif bias_init == "zero":
|
|
torch.nn.init.zeros_(layer.bias)
|
|
else:
|
|
raise NotImplementedError
|
|
|
|
return layer
|
|
|
|
def make_activation(self, activation):
|
|
if activation == "relu":
|
|
return nn.ReLU(inplace=True)
|
|
elif activation == "silu":
|
|
return nn.SiLU(inplace=True)
|
|
else:
|
|
raise NotImplementedError
|
|
|
|
def forward(self, x):
|
|
inp_shape = x.shape[:-1]
|
|
x = x.reshape(-1, x.shape[-1])
|
|
|
|
features = self.layers(x)
|
|
features = features.reshape(*inp_shape, -1)
|
|
out = {"density": features[..., 0:1], "features": features[..., 1:4]}
|
|
|
|
return out
|