| | 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) |
