wehub-resource-sync
120 lines
4.2 KiB
Python
120 lines
4.2 KiB
Python
import numpy as np
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
from pointnet2 import PointNet2FP, SAModule, SAMSGModule
|
|
from torch.autograd import Variable
|
|
|
|
|
|
class PointNet2SSGPartSeg(nn.Module):
|
|
def __init__(self, output_classes, batch_size, input_dims=6):
|
|
super(PointNet2SSGPartSeg, self).__init__()
|
|
# if normal_channel == true, input_dims = 6+3
|
|
self.input_dims = input_dims
|
|
|
|
self.sa_module1 = SAModule(
|
|
512, batch_size, 0.2, [input_dims, 64, 64, 128], n_neighbor=32
|
|
)
|
|
self.sa_module2 = SAModule(
|
|
128, batch_size, 0.4, [128 + 3, 128, 128, 256]
|
|
)
|
|
self.sa_module3 = SAModule(
|
|
None, batch_size, None, [256 + 3, 256, 512, 1024], group_all=True
|
|
)
|
|
|
|
self.fp3 = PointNet2FP(1280, [256, 256])
|
|
self.fp2 = PointNet2FP(384, [256, 128])
|
|
# if normal_channel == true, 128+16+6+3
|
|
self.fp1 = PointNet2FP(128 + 16 + 6, [128, 128, 128])
|
|
|
|
self.conv1 = nn.Conv1d(128, 128, 1)
|
|
self.bn1 = nn.BatchNorm1d(128)
|
|
self.drop1 = nn.Dropout(0.5)
|
|
self.conv2 = nn.Conv1d(128, output_classes, 1)
|
|
|
|
def forward(self, x, cat_vec=None):
|
|
if x.shape[-1] > 3:
|
|
l0_pos = x[:, :, :3]
|
|
l0_feat = x
|
|
else:
|
|
l0_pos = x
|
|
l0_feat = x
|
|
# Set Abstraction layers
|
|
l1_pos, l1_feat = self.sa_module1(l0_pos, l0_feat) # l1_feat: [B, N, D]
|
|
l2_pos, l2_feat = self.sa_module2(l1_pos, l1_feat)
|
|
l3_pos, l3_feat = self.sa_module3(l2_pos, l2_feat) # [B, N, C], [B, D]
|
|
# Feature Propagation layers
|
|
l2_feat = self.fp3(
|
|
l2_pos, l3_pos, l2_feat, l3_feat.unsqueeze(1)
|
|
) # l2_feat: [B, D, N]
|
|
l1_feat = self.fp2(l1_pos, l2_pos, l1_feat, l2_feat.permute(0, 2, 1))
|
|
l0_feat = torch.cat([cat_vec.permute(0, 2, 1), l0_pos, l0_feat], 2)
|
|
l0_feat = self.fp1(l0_pos, l1_pos, l0_feat, l1_feat.permute(0, 2, 1))
|
|
# FC layers
|
|
feat = F.relu(self.bn1(self.conv1(l0_feat)))
|
|
out = self.drop1(feat)
|
|
out = self.conv2(out) # [B, output_classes, N]
|
|
return out
|
|
|
|
|
|
class PointNet2MSGPartSeg(nn.Module):
|
|
def __init__(self, output_classes, batch_size, input_dims=6):
|
|
super(PointNet2MSGPartSeg, self).__init__()
|
|
|
|
self.sa_msg_module1 = SAMSGModule(
|
|
512,
|
|
batch_size,
|
|
[0.1, 0.2, 0.4],
|
|
[32, 64, 128],
|
|
[
|
|
[input_dims, 32, 32, 64],
|
|
[input_dims, 64, 64, 128],
|
|
[input_dims, 64, 96, 128],
|
|
],
|
|
)
|
|
self.sa_msg_module2 = SAMSGModule(
|
|
128,
|
|
batch_size,
|
|
[0.4, 0.8],
|
|
[64, 128],
|
|
[
|
|
[128 + 128 + 64 + 3, 128, 128, 256],
|
|
[128 + 128 + 64 + 3, 128, 196, 256],
|
|
],
|
|
)
|
|
self.sa_module3 = SAModule(
|
|
None, batch_size, None, [512 + 3, 256, 512, 1024], group_all=True
|
|
)
|
|
|
|
self.fp3 = PointNet2FP(1536, [256, 256])
|
|
self.fp2 = PointNet2FP(576, [256, 128])
|
|
# if normal_channel == true, 150 + 3
|
|
self.fp1 = PointNet2FP(150, [128, 128])
|
|
|
|
self.conv1 = nn.Conv1d(128, 128, 1)
|
|
self.bn1 = nn.BatchNorm1d(128)
|
|
self.drop1 = nn.Dropout(0.5)
|
|
self.conv2 = nn.Conv1d(128, output_classes, 1)
|
|
|
|
def forward(self, x, cat_vec=None):
|
|
if x.shape[-1] > 3:
|
|
l0_pos = x[:, :, :3]
|
|
l0_feat = x
|
|
else:
|
|
l0_pos = x
|
|
l0_feat = x
|
|
# Set Abstraction layers
|
|
l1_pos, l1_feat = self.sa_msg_module1(l0_pos, l0_feat)
|
|
l2_pos, l2_feat = self.sa_msg_module2(l1_pos, l1_feat)
|
|
l3_pos, l3_feat = self.sa_module3(l2_pos, l2_feat)
|
|
# Feature Propagation layers
|
|
l2_feat = self.fp3(l2_pos, l3_pos, l2_feat, l3_feat.unsqueeze(1))
|
|
l1_feat = self.fp2(l1_pos, l2_pos, l1_feat, l2_feat.permute(0, 2, 1))
|
|
l0_feat = torch.cat([cat_vec.permute(0, 2, 1), l0_pos, l0_feat], 2)
|
|
l0_feat = self.fp1(l0_pos, l1_pos, l0_feat, l1_feat.permute(0, 2, 1))
|
|
# FC layers
|
|
feat = F.relu(self.bn1(self.conv1(l0_feat)))
|
|
out = self.drop1(feat)
|
|
out = self.conv2(out)
|
|
return out
|