122 lines
4.5 KiB
Python
122 lines
4.5 KiB
Python
# modified from https://github.com/autonomousvision/convolutional_occupancy_networks/blob/master/src/encoder/pointnet.py
|
|
from dataclasses import dataclass
|
|
import torch
|
|
import torch.nn as nn
|
|
from torch_scatter import scatter_mean, scatter_max
|
|
|
|
from TriplaneGaussian.utils.base import BaseModule
|
|
from TriplaneGaussian.models.networks import ResnetBlockFC
|
|
from TriplaneGaussian.utils.ops import scale_tensor
|
|
|
|
class LocalPoolPointnet(BaseModule):
|
|
''' PointNet-based encoder network with ResNet blocks for each point.
|
|
Number of input points are fixed.
|
|
|
|
Args:
|
|
c_dim (int): dimension of latent code c
|
|
dim (int): input points dimension
|
|
hidden_dim (int): hidden dimension of the network
|
|
scatter_type (str): feature aggregation when doing local pooling
|
|
plane_resolution (int): defined resolution for plane feature
|
|
padding (float): conventional padding paramter of ONet for unit cube, so [-0.5, 0.5] -> [-0.55, 0.55]
|
|
n_blocks (int): number of blocks ResNetBlockFC layers
|
|
'''
|
|
|
|
@dataclass
|
|
class Config(BaseModule.Config):
|
|
input_channels: int = 3
|
|
c_dim: int = 128
|
|
hidden_dim: int = 128
|
|
scatter_type: str = "max"
|
|
plane_size: int = 32
|
|
n_blocks: int = 5
|
|
radius: float = 1.
|
|
|
|
cfg: Config
|
|
|
|
def configure(self) -> None:
|
|
super().configure()
|
|
self.fc_pos = nn.Linear(self.cfg.input_channels, 2 * self.cfg.hidden_dim)
|
|
self.blocks = nn.ModuleList([
|
|
ResnetBlockFC(2 * self.cfg.hidden_dim, self.cfg.hidden_dim) for i in range(self.cfg.n_blocks)
|
|
])
|
|
self.fc_c = nn.Linear(self.cfg.hidden_dim, self.cfg.c_dim)
|
|
|
|
self.actvn = nn.ReLU()
|
|
|
|
if self.cfg.scatter_type == 'max':
|
|
self.scatter = scatter_max
|
|
elif self.cfg.scatter_type == 'mean':
|
|
self.scatter = scatter_mean
|
|
else:
|
|
raise ValueError('incorrect scatter type')
|
|
|
|
|
|
def generate_plane_features(self, index, c):
|
|
# acquire indices of features in plane
|
|
# xy = normalize_coordinate(p.clone(), plane=plane, padding=self.padding) # normalize to the range of (0, 1)
|
|
# index = self.coordinate2index(x, self.cfg.plane_size)
|
|
|
|
# scatter plane features from points
|
|
fea_plane = c.new_zeros(index.shape[0], self.cfg.c_dim, self.cfg.plane_size ** 2)
|
|
c = c.permute(0, 2, 1) # B x 512 x T
|
|
fea_plane = scatter_mean(c, index, out=fea_plane) # B x 512 x reso^2
|
|
fea_plane = fea_plane.reshape(index.shape[0], self.cfg.c_dim, self.cfg.plane_size, self.cfg.plane_size) # sparce matrix (B x 512 x reso x reso)
|
|
|
|
return fea_plane
|
|
|
|
def pool_local(self, xy, index, c):
|
|
bs, fea_dim = c.shape[0], c.shape[2]
|
|
keys = xy.keys()
|
|
|
|
c_out = 0
|
|
for key in keys:
|
|
# scatter plane features from points
|
|
fea = self.scatter(c.permute(0, 2, 1), index[key], dim_size=self.cfg.plane_size ** 2)
|
|
if self.scatter == scatter_max:
|
|
fea = fea[0]
|
|
# gather feature back to points
|
|
fea = fea.gather(dim=2, index=index[key].expand(-1, fea_dim, -1))
|
|
c_out += fea
|
|
return c_out.permute(0, 2, 1)
|
|
|
|
def coordinate2index(self, x):
|
|
x = (x * self.cfg.plane_size).long()
|
|
index = x[..., 0] + self.cfg.plane_size * x[..., 1]
|
|
assert index.max() < self.cfg.plane_size ** 2
|
|
return index[:, None, :]
|
|
|
|
def forward(self, p):
|
|
batch_size, T, D = p.shape
|
|
|
|
# acquire the index for each point
|
|
coord = {}
|
|
index = {}
|
|
|
|
position = torch.clamp(p[..., :3], -self.cfg.radius + 1e-6, self.cfg.radius - 1e-6)
|
|
position_norm = scale_tensor(position, (-self.cfg.radius, self.cfg.radius), (0, 1))
|
|
coord["xy"] = position_norm[..., [0, 1]]
|
|
coord["xz"] = position_norm[..., [0, 2]]
|
|
coord["yz"] = position_norm[..., [1, 2]]
|
|
index["xy"] = self.coordinate2index(coord["xy"])
|
|
index["xz"] = self.coordinate2index(coord["xz"])
|
|
index["yz"] = self.coordinate2index(coord["yz"])
|
|
|
|
net = self.fc_pos(p)
|
|
|
|
net = self.blocks[0](net)
|
|
for block in self.blocks[1:]:
|
|
pooled = self.pool_local(coord, index, net)
|
|
net = torch.cat([net, pooled], dim=2)
|
|
net = block(net)
|
|
|
|
c = self.fc_c(net)
|
|
|
|
features = torch.stack([
|
|
self.generate_plane_features(index["xy"], c),
|
|
self.generate_plane_features(index["xz"], c),
|
|
self.generate_plane_features(index["yz"], c)
|
|
], dim=1)
|
|
|
|
return features
|