| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| import numpy as np |
| import sys |
| import os |
| from external.pointnet2.pointnet2_modules import PointnetSAModuleVotes, PointnetFPModule |
| from .utils import zero_module |
| from .Positional_Embedding import PositionalEmbedding |
|
|
| class Pointnet2Encoder(nn.Module): |
| def __init__(self,input_feature_dim=0,npoints=[2048,1024,512,256],radius=[0.2,0.4,0.6,1.2],nsample=[64,32,16,8]): |
| super().__init__() |
| self.sa1 = PointnetSAModuleVotes( |
| npoint=npoints[0], |
| radius=radius[0], |
| nsample=nsample[0], |
| mlp=[input_feature_dim, 64, 64, 128], |
| use_xyz=True, |
| normalize_xyz=True |
| ) |
|
|
| self.sa2 = PointnetSAModuleVotes( |
| npoint=npoints[1], |
| radius=radius[1], |
| nsample=nsample[1], |
| mlp=[128, 128, 128, 256], |
| use_xyz=True, |
| normalize_xyz=True |
| ) |
|
|
| self.sa3 = PointnetSAModuleVotes( |
| npoint=npoints[2], |
| radius=radius[2], |
| nsample=nsample[2], |
| mlp=[256, 256, 256, 512], |
| use_xyz=True, |
| normalize_xyz=True |
| ) |
|
|
| self.sa4 = PointnetSAModuleVotes( |
| npoint=npoints[3], |
| radius=radius[3], |
| nsample=nsample[3], |
| mlp=[512, 512, 512, 512], |
| use_xyz=True, |
| normalize_xyz=True |
| ) |
| def _break_up_pc(self, pc): |
| xyz = pc[..., 0:3].contiguous() |
| features = ( |
| pc[..., 3:].transpose(1, 2).contiguous() |
| if pc.size(-1) > 3 else None |
| ) |
|
|
| return xyz, features |
| def forward(self,pointcloud,end_points=None): |
| if not end_points: end_points = {} |
| batch_size = pointcloud.shape[0] |
|
|
| xyz, features = self._break_up_pc(pointcloud) |
|
|
| end_points['org_xyz'] = xyz |
| |
| xyz1, features1, _ = self.sa1(xyz, features) |
| end_points['sa1_xyz'] = xyz1 |
| end_points['sa1_features'] = features1 |
|
|
| xyz2, features2, _ = self.sa2(xyz1, features1) |
| end_points['sa2_xyz'] = xyz2 |
| end_points['sa2_features'] = features2 |
|
|
| xyz3, features3, _ = self.sa3(xyz2, features2) |
| end_points['sa3_xyz'] = xyz3 |
| end_points['sa3_features'] = features3 |
| |
| xyz4, features4, _ = self.sa4(xyz3, features3) |
| end_points['sa4_xyz'] = xyz4 |
| end_points['sa4_features'] = features4 |
| |
| return end_points |
|
|
|
|
|
|
| class PointUNet(nn.Module): |
| r""" |
| Backbone network for point cloud feature learning. |
| Based on Pointnet++ single-scale grouping network. |
| |
| Parameters |
| ---------- |
| input_feature_dim: int |
| Number of input channels in the feature descriptor for each point. |
| e.g. 3 for RGB. |
| """ |
|
|
| def __init__(self): |
| super().__init__() |
|
|
| self.noisy_encoder=Pointnet2Encoder() |
| self.cond_encoder=Pointnet2Encoder() |
| self.fp1_cross = PointnetFPModule(mlp=[512 + 512, 512, 512]) |
| self.fp1 = PointnetFPModule(mlp=[512 + 512, 512, 512]) |
| |
| self.fp2_cross = PointnetFPModule(mlp=[512 + 512, 512, 256]) |
| self.fp2 = PointnetFPModule(mlp=[256 + 256, 512, 256]) |
| |
| self.fp3_cross= PointnetFPModule(mlp=[256 + 256, 256, 128]) |
| self.fp3 = PointnetFPModule(mlp=[128 + 128, 256, 128]) |
| |
| self.fp4_cross=PointnetFPModule(mlp=[128+128, 128, 128]) |
| self.fp4 = PointnetFPModule(mlp=[128, 128, 128]) |
| |
|
|
| self.output_layer=nn.Sequential( |
| nn.LayerNorm(128), |
| zero_module(nn.Linear(in_features=128,out_features=3,bias=False)) |
| ) |
| self.t_emb_layer = PositionalEmbedding(256) |
| self.map_layer0 = nn.Linear(in_features=256, out_features=512) |
| self.map_layer1 = nn.Linear(in_features=512, out_features=512) |
|
|
| def forward(self, noise_points, t,cond_points): |
| r""" |
| Forward pass of the network |
| |
| Parameters |
| ---------- |
| pointcloud: Variable(torch.cuda.FloatTensor) |
| (B, N, 3 + input_feature_dim) tensor |
| Point cloud to run predicts on |
| Each point in the point-cloud MUST |
| be formated as (x, y, z, features...) |
| |
| Returns |
| ---------- |
| end_points: {XXX_xyz, XXX_features, XXX_inds} |
| XXX_xyz: float32 Tensor of shape (B,K,3) |
| XXX_features: float32 Tensor of shape (B,K,D) |
| XXX-inds: int64 Tensor of shape (B,K) values in [0,N-1] |
| """ |
| t_emb = self.t_emb_layer(t) |
| t_emb = F.silu(self.map_layer0(t_emb)) |
| t_emb = F.silu(self.map_layer1(t_emb)) |
| t_emb = t_emb[:, :, None] |
| noise_end_points=self.noisy_encoder(noise_points) |
| cond=self.cond_encoder(cond_points) |
| |
| features = self.fp1_cross(noise_end_points['sa4_xyz'],cond['sa4_xyz'],noise_end_points['sa4_features']+t_emb, |
| cond['sa4_features']) |
| features = self.fp1(noise_end_points['sa3_xyz'], noise_end_points['sa4_xyz'], noise_end_points['sa3_features'], |
| features) |
| features = self.fp2_cross(noise_end_points['sa3_xyz'],cond['sa3_xyz'],features, |
| cond["sa3_features"]) |
| features = self.fp2(noise_end_points['sa2_xyz'], noise_end_points['sa3_xyz'], noise_end_points['sa2_features'], |
| features) |
| features = self.fp3_cross(noise_end_points['sa2_xyz'],cond['sa2_xyz'],features, |
| cond['sa2_features']) |
| features = self.fp3(noise_end_points['sa1_xyz'],noise_end_points['sa2_xyz'],noise_end_points['sa1_features'],features) |
| features = self.fp4_cross(noise_end_points['sa1_xyz'],cond['sa1_xyz'],features, |
| cond['sa1_features']) |
| features = self.fp4(noise_end_points['org_xyz'], noise_end_points['sa1_xyz'], None, features) |
| features=features.transpose(1,2) |
|
|
| |
| |
| |
| |
| |
| |
| |
| |
| |
| output_points=self.output_layer(features) |
|
|
| return output_points |
|
|
|
|
| if __name__ == '__main__': |
| net=PointUNet().cuda().float() |
| net=net.eval() |
| noise_points=torch.randn(16,4096,3).cuda().float() |
| cond_points=torch.randn(16,4096,3).cuda().float() |
| t=torch.randn(16).cuda().float() |
| cond_encoder=Pointnet2Encoder().cuda().float() |
|
|
| out = net(noise_points,cond_points) |
| print(out.shape) |
