1. PointNet - Deep Learning on Point Sets for 3D Classification and Segmentation
Point cloud is an important type of geometric data structure. Due to its irregular format, most researchers transform such data to regular 3D voxel grids or collections of images. This, however, renders data unnecessarily voluminous and causes issues. In this paper, we design a novel type of neural network that directly consumes point clouds, which well respects the permutation invariance of points in the input. Our network, named PointNet, provides a unified architecture for applications ranging from object classification, part segmentation, to scene semantic parsing. Though simple, PointNet is highly efficient and effective. Empirically, it shows strong performance on par or even better than state of the art. Theoretically, we provide analysis towards understanding of what the network has learnt and why the network is robust with respect to input perturbation and corruption.
1.1. Problem Statement
We design a deep learning framework that directly consumes unordered point sets as inputs. A point cloud is represented as a set of 3D points ${P_i | i = 1, …, n}$, where each point $P_i$ is a vector of its $(x, y, z)$ coordinate plus extra feature channels such as color, normal etc. For simplicity and clarity, unless otherwise noted, we only use the $(x, y, z)$ coordinate as our point’s channels.
For the object classification task, the input point cloud is either directly sampled from a shape or pre-segmented from a scene point cloud. Our proposed deep network outputs $k$ scores for all the $k$ candidate classes. For semantic segmentation, the input can be a single object for part region segmentation, or a sub-volume from a 3D scene for object region segmentation. Our model will output $n$ × $m$ scores for each of the $n$ points and each of the $m$ semantic subcategories.
1.2. Deep Learning on Point Sets
1.2.1. Properties of Point Sets in $R^n$
Our input is a subset of points from an Euclidean space. It has three main properties:
• Unordered: Unlike pixel arrays in images or voxel arrays in volumetric grids, point cloud is a set of points without specific order. In other words, a network that consumes N 3D point sets needs to be invariant to N! permutations of the input set in data feeding order.
• Interaction among points: The points are from a space with a distance metric. It means that points are not isolated, and neighboring points form a meaningful subset. Therefore, the model needs to be able to capture local structures from nearby points, and the combinatorial interactions among local structures.
• Invariance under transformations: As a geometric object, the learned representation of the point set should be invariant to certain transformations. For example, rotating and translating points all together should not modify the global point cloud category nor the segmentation of the points.
1.2.2. PointNet Architecture
PointNet has three key modules: the max pooling layer as a symmetric function to aggregate information from all the points, a local and global information combination structure, and two joint alignment networks that align both input points and point features.
1. Symmetry Function for Unordered Input
In order to make a model invariant to input permutation, three strategies exist:
Sort input into a canonical order
Treat the input as a sequence to train an RNN, but augment the training data by all kinds of permutations
Use a simple symmetric function to aggregate the information from each point. Here, a symmetric function takes n vectors as input and outputs a new vector that is invariant to the input order. For example, + and ∗ operators are symmetric binary functions.
While sorting sounds like a simple solution, in high dimensional space there in fact does not exist an ordering that is stable w.r.t. point perturbations in the general sense. This can be easily shown by contradiction. If such an ordering strategy exists, it defines a bijection map between a high-dimensional space and a 1d real line. It is not hard to see, to require an ordering to be stable w.r.t point perturbations is equivalent to requiring that this map preserves spatial proximity as the dimension reduces, a task that cannot be achieved in the general case. Therefore, sorting does not fully resolve the ordering issue, and it’s hard for a network to learn a consistent mapping from input to output as the ordering issue persists. As shown in experiments, we find that applying a MLP directly on the sorted point set performs poorly, though slightly better than directly processing an unsorted input.
RNN not good.
=> approximate a general function defined on a point set by applying a symmetric function on transformed elements in the set.
Empirically, our basic module is very simple: we approximate $h$ by a multi-layer perceptron network and $g$ by a composition of a single variable function and a max pooling function. This is found to work well by experiments. Through a collection of $h$, we can learn a number of $f$’s to capture different properties of the set.
2. Local and Global Information Aggregation
The output from the above section forms a vector $[f_1, . . . , f_K]$, which is a global signature of the input set. We can easily train a SVM or multi-layer perceptron classifier on the shape global features for classification. However, point segmentation requires a combination of local and global knowledge. We can achieve this by a simple yet highly effective manner.
Our solution can be seen in Segmentation Network part. After computing the global point cloud feature vector, we feed it back to per point features by concatenating the global feature with each of the point features. Then we extract new per point features based on the combined point features - this time the per point feature is aware of both the local and global information.
=> With this modification our network is able to predict per point quantities that rely on both local geometry and global semantics .
For example we can accurately predict per-point normals (fig in supplementary), validating that the network is able to summarize information from the point’s local neighborhood.
3. Joint Alignment Network
The semantic labeling of a point cloud has to be invariant if the point cloud undergoes certain geometric transformations, such as rigid transformation. We therefore expect that the learnt representation by our point set is invariant to these transformations.
Our input form of point clouds allows us to achieve this goal in a much simpler way. We do not need to invent any new layers and no alias is introduced as in the image case. We predict an affine transformation matrix by a mini-network (T-net) and directly apply this transformation to the coordinates of input points. The mininetwork itself resembles the big network and is composed by basic modules of point independent feature extraction, max pooling and fully connected layers. More details about the T-net are in the supplementary. This idea can be further extended to the alignment of feature space, as well. We can insert another alignment network on point features and predict a feature transformation matrix to align features from different input point clouds. However, transformation matrix in the feature space has much higher dimension than the spatial transform matrix, which greatly increases the difficulty of optimization. We therefore add a regularization term to our softmax training loss. We constrain the feature transformation matrix to be close to orthogonal matrix:
where $A$ is the feature alignment matrix predicted by a mini-network. An orthogonal transformation will not lose information in the input, thus is desired. We find that by adding the regularization term, the optimization becomes more stable and our model achieves better performance.
Intuitively, PointNet learns to summarize a shape by a sparse set of key points
1.3. Experiments
Code Explanation
In this part, I use the code in Pytorch implementation to explain PointNet architecture based on this figure:
Input Transform Network
The first transformation network is a mini-PointNet that takes raw point cloud as input and regresses to a 3 × 3 matrix. It’s composed of a shared MLP(64, 128, 1024) network (with layer output sizes 64, 128, 1024) on each point, a max pooling across points and two fully connected layers with output sizes 512, 256. The output matrix is initialized as an identity matrix. All layers, except the last one, include ReLU and batch normalization.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
class STN3d(nn.Module):
def __init__(self):
super(STN3d, self).__init__()
self.conv1 = torch.nn.Conv1d(3, 64, 1)
self.conv2 = torch.nn.Conv1d(64, 128, 1)
self.conv3 = torch.nn.Conv1d(128, 1024, 1)
self.fc1 = nn.Linear(1024, 512)
self.fc2 = nn.Linear(512, 256)
self.fc3 = nn.Linear(256, 9)
self.relu = nn.ReLU()
self.bn1 = nn.BatchNorm1d(64)
self.bn2 = nn.BatchNorm1d(128)
self.bn3 = nn.BatchNorm1d(1024)
self.bn4 = nn.BatchNorm1d(512)
self.bn5 = nn.BatchNorm1d(256)
def forward(self, x):
batchsize = x.size()[0]
x = F.relu(self.bn1(self.conv1(x)))
x = F.relu(self.bn2(self.conv2(x)))
x = F.relu(self.bn3(self.conv3(x)))
x = torch.max(x, 2, keepdim=True)[0]
# Ref: https://pytorch.org/docs/master/tensors.html?highlight=view#torch.Tensor.view
# the size -1 is inferred from other dimensions
x = x.view(-1, 1024)
x = F.relu(self.bn4(self.fc1(x)))
x = F.relu(self.bn5(self.fc2(x)))
x = self.fc3(x)
iden = Variable(torch.from_numpy(np.array([1,0,0,0,1,0,0,0,1]).astype(np.float32))).view(1,9).repeat(batchsize,1)
if x.is_cuda:
iden = iden.cuda()
x = x + iden
x = x.view(-1, 3, 3)
return x
Feature Transform Network
The second transformation network has the same architecture as the first one except that the output is a 64 × 64 matrix. The matrix is also initialized as an identity. A regularization loss (with weight 0.001) is added to the softmax classification loss to make the matrix close to orthogonal.
We use dropout with keep ratio 0.7 on the last fully connected layer, whose output dimension 256, before class score prediction. The decay rate for batch normalization starts with 0.5 and is gradually increased to 0.99. We use adam optimizer with initial learning rate 0.001, momentum 0.9 and batch size 32.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
class STNkd(nn.Module):
def __init__(self, k=64):
super(STNkd, self).__init__()
self.conv1 = torch.nn.Conv1d(k, 64, 1)
self.conv2 = torch.nn.Conv1d(64, 128, 1)
self.conv3 = torch.nn.Conv1d(128, 1024, 1)
self.fc1 = nn.Linear(1024, 512)
self.fc2 = nn.Linear(512, 256)
self.fc3 = nn.Linear(256, k*k)
self.relu = nn.ReLU()
self.bn1 = nn.BatchNorm1d(64)
self.bn2 = nn.BatchNorm1d(128)
self.bn3 = nn.BatchNorm1d(1024)
self.bn4 = nn.BatchNorm1d(512)
self.bn5 = nn.BatchNorm1d(256)
self.k = k
def forward(self, x):
batchsize = x.size()[0]
x = F.relu(self.bn1(self.conv1(x)))
x = F.relu(self.bn2(self.conv2(x)))
x = F.relu(self.bn3(self.conv3(x)))
x = torch.max(x, 2, keepdim=True)[0]
x = x.view(-1, 1024)
x = F.relu(self.bn4(self.fc1(x)))
x = F.relu(self.bn5(self.fc2(x)))
x = self.fc3(x)
iden = Variable(torch.from_numpy(np.eye(self.k).flatten().astype(np.float32))).view(1,self.k*self.k).repeat(batchsize,1)
if x.is_cuda:
iden = iden.cuda()
x = x + iden
x = x.view(-1, self.k, self.k)
return x
PointNet Feature Extractor
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
class PointNetfeat(nn.Module):
def __init__(self, global_feat = True, feature_transform = False):
super(PointNetfeat, self).__init__()
self.stn = STN3d()
self.conv1 = torch.nn.Conv1d(3, 64, 1)
self.conv2 = torch.nn.Conv1d(64, 128, 1)
self.conv3 = torch.nn.Conv1d(128, 1024, 1)
self.bn1 = nn.BatchNorm1d(64)
self.bn2 = nn.BatchNorm1d(128)
self.bn3 = nn.BatchNorm1d(1024)
self.global_feat = global_feat
self.feature_transform = feature_transform
if self.feature_transform:
self.fstn = STNkd(k=64)
def forward(self, x):
n_pts = x.size()[2]
trans = self.stn(x)
x = x.transpose(2, 1)
x = torch.bmm(x, trans)
x = x.transpose(2, 1)
x = F.relu(self.bn1(self.conv1(x)))
if self.feature_transform:
trans_feat = self.fstn(x)
x = x.transpose(2,1)
x = torch.bmm(x, trans_feat)
x = x.transpose(2,1)
else:
trans_feat = None
pointfeat = x
x = F.relu(self.bn2(self.conv2(x)))
x = self.bn3(self.conv3(x))
x = torch.max(x, 2, keepdim=True)[0]
x = x.view(-1, 1024)
if self.global_feat:
return x, trans, trans_feat
else:
x = x.view(-1, 1024, 1).repeat(1, 1, n_pts)
return torch.cat([x, pointfeat], 1), trans, trans_feat
PointNet Classifier
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
class PointNetCls(nn.Module):
def __init__(self, k=2, feature_transform=False):
super(PointNetCls, self).__init__()
self.feature_transform = feature_transform
self.feat = PointNetfeat(global_feat=True, feature_transform=feature_transform)
self.fc1 = nn.Linear(1024, 512)
self.fc2 = nn.Linear(512, 256)
self.fc3 = nn.Linear(256, k)
self.dropout = nn.Dropout(p=0.3)
self.bn1 = nn.BatchNorm1d(512)
self.bn2 = nn.BatchNorm1d(256)
self.relu = nn.ReLU()
def forward(self, x):
x, trans, trans_feat = self.feat(x)
x = F.relu(self.bn1(self.fc1(x)))
x = F.relu(self.bn2(self.dropout(self.fc2(x))))
x = self.fc3(x)
return F.log_softmax(x, dim=1), trans, trans_feat
PointNet Segmentation
The segmentation network is an extension to the classification PointNet. Local point features (the output after the second transformation network) and global feature (output of the max pooling) are concatenated for each point. No dropout is used for segmentation network. Training parameters are the same as the classification network.
As to the task of shape part segmentation, we made a few modifications to the basic segmentation network architecture (Fig 2 in main paper) in order to achieve best performancepa. We add a one-hot vector indicating the class of the input and concatenate it with the max pooling layer’s output. We also increase neurons in some layers and add skip links to collect local point features in different layers and concatenate them to form point feature input to the segmentation network.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
class PointNetDenseCls(nn.Module):
def __init__(self, k = 2, feature_transform=False):
super(PointNetDenseCls, self).__init__()
self.k = k
self.feature_transform=feature_transform
self.feat = PointNetfeat(global_feat=False, feature_transform=feature_transform)
self.conv1 = torch.nn.Conv1d(1088, 512, 1)
self.conv2 = torch.nn.Conv1d(512, 256, 1)
self.conv3 = torch.nn.Conv1d(256, 128, 1)
self.conv4 = torch.nn.Conv1d(128, self.k, 1)
self.bn1 = nn.BatchNorm1d(512)
self.bn2 = nn.BatchNorm1d(256)
self.bn3 = nn.BatchNorm1d(128)
def forward(self, x):
batchsize = x.size()[0]
n_pts = x.size()[2]
x, trans, trans_feat = self.feat(x)
x = F.relu(self.bn1(self.conv1(x)))
x = F.relu(self.bn2(self.conv2(x)))
x = F.relu(self.bn3(self.conv3(x)))
x = self.conv4(x)
x = x.transpose(2,1).contiguous()
x = F.log_softmax(x.view(-1,self.k), dim=-1)
x = x.view(batchsize, n_pts, self.k)
return x, trans, trans_feat
Sample segmentation result: 
Training
2. PointNet++: Deep Hierarchical Feature Learning on Point Sets in a Metric Space
2.1. PointNet nad PointNet++
The basic idea of PointNet is to learn a spatial encoding of each point and then aggregate all individual point features to a global point cloud signature. By its design, PointNet does not capture local structure induced by the metric. However, exploiting local structure has proven to be important for the success of convolutional architectures. A CNN takes data defined on regular grids as the input and is able to progressively capture features at increasingly larger scales along a multi-resolution hierarchy. At lower levels neurons have smaller receptive fields whereas at higher levels they have larger receptive fields. The ability to abstract local patterns along the hierarchy allows better generalizability to unseen cases.
We introduce a hierarchical neural network, named as PointNet++, to process a set of points sampled in a metric space in a hierarchical fashion. The general idea of PointNet++ is simple. We first partition the set of points into overlapping local regions by the distance metric of the underlying space. Similar to CNNs, we extract local features capturing fine geometric structures from small neighborhoods; such local features are further grouped into larger units and processed to produce higher level features. This process is repeated until we obtain the features of the whole point set.
The design of PointNet++ has to address two issues: how to generate the partitioning of the point set, and how to abstract sets of points or local features through a local feature learner. The two issues are correlated because the partitioning of the point set has to produce common structures across partitions, so that weights of local feature learners can be shared, as in the convolutional setting. We choose our local feature learner to be PointNet. As demonstrated in that work, PointNet is an effective architecture to process an unordered set of points for semantic feature extraction. In addition, this architecture is robust to input data corruption. As a basic building block, PointNet abstracts sets of local points or features into higher level representations. In this view, PointNet++ applies PointNet recursively on a nested partitioning of the input set.
One issue that still remains is how to generate overlapping partitioning of a point set. Each partition is defined as a neighborhood ball in the underlying Euclidean space, whose parameters include centroid location and scale. To evenly cover the whole set, the centroids are selected among input point set by a farthest point sampling (FPS) algorithm. Compared with volumetric CNNs that scan the space with fixed strides, our local receptive fields are dependent on both the input data and the metric, and thus more efficient and effective.
A significant contribution of this paper is that PointNet++ leverages neighborhoods at multiple scales to achieve both robustness and detail capture. Assisted with random input dropout during training, the network learns to adaptively weight patterns detected at different scales and combine multi-scale features according to the input data.
2.2. Problem Statement
Suppose that $X = (M, d)$ is a discrete metric space whose metric is inherited from a Euclidean space $R^n$, where $M ⊆ R^n$ is the set of points and d is the distance metric. In addition, the density of $M$ in the ambient Euclidean space may not be uniform everywhere. We are interested in learning set functions $f$ that take such $X$ as the input (along with additional features for each point) and produce information of semantic interest regrading $X$ . In practice, such $f$ can be classification function that assigns a label to $X$ or a segmentation function that assigns a per point label to each member of $M$.
2.3. Method
2.3.1. Hierarchical Point Set Feature Learning
While PointNet uses a single max pooling operation to aggregate the whole point set, our new architecture builds a hierarchical grouping of points and progressively abstract larger and larger local regions along the hierarchy. Our hierarchical structure is composed by a number of set abstraction levels. At each level, a set of points is processed and abstracted to produce a new set with fewer elements. The set abstraction level is made of three key layers: Sampling layer, Grouping layer and PointNet layer. The Sampling layer selects a set of points from input points, which defines the centroids of local regions. Grouping layer then constructs local region sets by finding “neighboring” points around the centroids. PointNet layer uses a mini-PointNet to encode local region patterns into feature vectors.
Sampling layer: Given input points ${x_1, x_2, …, x_n}$, we use iterative farthest point sampling (FPS) to choose a subset of points ${x_{i_1}, x_{i_2}, …, x_{i_m}}$, such that $x_{i_j}$ is the most distant point (in metric distance) from the set ${x_{i_1}, x_{i_2}, …, x_{i_{j−1}}}$ with regard to the rest points. Compared with random sampling, it has better coverage of the entire point set given the same number of centroids. In contrast to CNNs that scan the vector space agnostic of data distribution, our sampling strategy generates receptive fields in a data dependent manner.
Grouping layer: The input to this layer is a point set of size $N × (d + C)$ and the coordinates of a set of centroids of size $N’ × d$. The output are groups of point sets of size $N’ × K × (d + C)$, where each group corresponds to a local region and $K$ is the number of points in the neighborhood of centroid points. Note that $K$ varies across groups but the succeeding PointNet layer is able to convert flexible number of points into a fixed length local region feature vector.
PointNet layer: In this layer, the input are $N’$ local regions of points with data size $N’×K×(d+C)$. Each local region in the output is abstracted by its centroid and local feature that encodes the centroid’s neighborhood. Output data size is $N’ × (d + C’)$.
The coordinates of points in a local region are firstly translated into a local frame relative to the centroid point: $x^{(j)}_i = x^{(j)}_i − \hat x^((j))$ for $i = 1, 2, …, K$ and $j = 1, 2, …, d$ where $\hat x$ is the coordinate of the centroid. We use PointNet as the basic building block for local pattern learning. By using relative coordinates together with point features we can capture point-to-point relations in the local region.
2.3.2. Robust Feature Learning under Non-Uniform Sampling Density
As discussed earlier, it is common that a point set comes with nonuniform density in different areas . Such non-uniformity introduces a significant challenge for point set feature learning. Features learned in dense data may not generalize to sparsely sampled regions. Consequently, models trained for sparse point cloud may not recognize fine-grained local structures .
Ideally, we want to inspect as closely as possible into a point set to capture finest details in densely sampled regions . However, such close inspect is prohibited at low density areas because local patterns may be corrupted by the sampling deficiency . In this case, we should look for larger scale patterns in greater vicinity . To achieve this goal we propose density adaptive PointNet layers that learn to combine features from regions of different scales when the input sampling density changes . We call our hierarchical network with density adaptive PointNet layers as PointNet++ .
Previously section, each abstraction level contains grouping and feature extraction of a single scale. In PointNet++, each abstraction level extracts multiple scales of local patterns and combine them intelligently according to local point densities. In terms of grouping local regions and combining features from different scales, we propose two types of density adaptive layers as listed below.
- Multi-scale grouping (MSG): A simple but effective way to capture multiscale patterns is to apply grouping layers with different scales followed by according PointNets to extract features of each scale. Features at different scales are concatenated to form a multi-scale feature.
We train the network to learn an optimized strategy to combine the multi-scale features. This is done by randomly dropping out input points with a randomized probability for each instance, which we call random input dropout. Specifically, for each training point set, we choose a dropout ratio $θ$ uniformly sampled from $[0, p]$ where $p ≤ 1$. For each point, we randomly drop a point with probability $θ$. In practice we set $p = 0.95$ to avoid generating empty point sets. In doing so we present the network with training sets of various sparsity (induced by $θ$) and varying uniformity (induced by randomness in dropout). During test, we keep all available points.
- Multi-resolution grouping (MRG): The MSG approach above is computationally expensive since it runs local PointNet at large scale neighborhoods for every centroid point. In particular, since the number of centroid points is usually quite large at the lowest level, the time cost is significant.
2.3.3. Point Feature Propagation for Set Segmentation
In set abstraction layer, the original point set is subsampled. However in set segmentation task such as semantic point labeling, we want to obtain point features for all the original points. One solution is to always sample all points as centroids in all set abstraction levels, which however results in high computation cost. Another way is to propagate features from subsampled points to the original points.
We adopt a hierarchical propagation strategy with distance based interpolation and across level skip links.
In a feature propagation level, we propagate point features from $N_l × (d + C)$ points to $N_{l−1}$ points where $N_{l−1}$ and $N_l$ (with $N_l ≤ N_{l−1}$) are point set size of input and output of set abstraction level $l$. We achieve feature propagation by interpolating feature values f of $N_l$ points at coordinates of the $N_{l−1}$ points. Among the many choices for interpolation, we use inverse distance weighted average based on k nearest neighbors (in default we use $p = 2, k = 3$). The interpolated features on $N_{l−1}$ points are then concatenated with skip linked point features from the set abstraction level. Then the concatenated features are passed through a “unit pointnet” , which is similar to one-by-one convolution in CNNs. A few shared fully connected and ReLU layers are applied to update each point’s feature vector. The process is repeated until we have propagated features to the original set of points.
2.4. Experiments
Code Explanation
In this part, I use the code in Pytorch implementation to explain PointNet++ architecture based on this figure:
PointNetSetAbstraction
Ref to Hierarchical Point Set Feature Learning and Robust Feature Learning under Non-Uniform Sampling Density part
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
class PointNetSetAbstraction(nn.Module):
def __init__(self, npoint, radius, nsample, in_channel, mlp, group_all):
super(PointNetSetAbstraction, self).__init__()
self.npoint = npoint
self.radius = radius
self.nsample = nsample
self.mlp_convs = nn.ModuleList()
self.mlp_bns = nn.ModuleList()
last_channel = in_channel
for out_channel in mlp:
self.mlp_convs.append(nn.Conv2d(last_channel, out_channel, 1))
self.mlp_bns.append(nn.BatchNorm2d(out_channel))
last_channel = out_channel
self.group_all = group_all
def forward(self, xyz, points):
"""
Input:
xyz: input points position data, [B, C, N]
points: input points data, [B, D, N]
Return:
new_xyz: sampled points position data, [B, C, S]
new_points_concat: sample points feature data, [B, D', S]
"""
xyz = xyz.permute(0, 2, 1)
if points is not None:
points = points.permute(0, 2, 1)
if self.group_all:
new_xyz, new_points = sample_and_group_all(xyz, points)
else:
new_xyz, new_points = sample_and_group(self.npoint, self.radius, self.nsample, xyz, points)
# new_xyz: sampled points position data, [B, npoint, C]
# new_points: sampled points data, [B, npoint, nsample, C+D]
new_points = new_points.permute(0, 3, 2, 1) # [B, C+D, nsample,npoint]
for i, conv in enumerate(self.mlp_convs):
bn = self.mlp_bns[i]
new_points = F.relu(bn(conv(new_points)))
new_points = torch.max(new_points, 2)[0]
new_xyz = new_xyz.permute(0, 2, 1)
return new_xyz, new_points
class PointNetSetAbstractionMsg(nn.Module):
def __init__(self, npoint, radius_list, nsample_list, in_channel, mlp_list):
super(PointNetSetAbstractionMsg, self).__init__()
self.npoint = npoint
self.radius_list = radius_list
self.nsample_list = nsample_list
self.conv_blocks = nn.ModuleList()
self.bn_blocks = nn.ModuleList()
for i in range(len(mlp_list)):
convs = nn.ModuleList()
bns = nn.ModuleList()
last_channel = in_channel + 3
for out_channel in mlp_list[i]:
convs.append(nn.Conv2d(last_channel, out_channel, 1))
bns.append(nn.BatchNorm2d(out_channel))
last_channel = out_channel
self.conv_blocks.append(convs)
self.bn_blocks.append(bns)
def forward(self, xyz, points):
"""
Input:
xyz: input points position data, [B, C, N]
points: input points data, [B, D, N]
Return:
new_xyz: sampled points position data, [B, C, S]
new_points_concat: sample points feature data, [B, D', S]
"""
xyz = xyz.permute(0, 2, 1)
if points is not None:
points = points.permute(0, 2, 1)
B, N, C = xyz.shape
S = self.npoint
new_xyz = index_points(xyz, farthest_point_sample(xyz, S))
new_points_list = []
for i, radius in enumerate(self.radius_list):
K = self.nsample_list[i]
group_idx = query_ball_point(radius, K, xyz, new_xyz)
grouped_xyz = index_points(xyz, group_idx)
grouped_xyz -= new_xyz.view(B, S, 1, C)
if points is not None:
grouped_points = index_points(points, group_idx)
grouped_points = torch.cat([grouped_points, grouped_xyz], dim=-1)
else:
grouped_points = grouped_xyz
grouped_points = grouped_points.permute(0, 3, 2, 1) # [B, D, K, S]
for j in range(len(self.conv_blocks[i])):
conv = self.conv_blocks[i][j]
bn = self.bn_blocks[i][j]
grouped_points = F.relu(bn(conv(grouped_points)))
new_points = torch.max(grouped_points, 2)[0] # [B, D', S]
new_points_list.append(new_points)
new_xyz = new_xyz.permute(0, 2, 1)
new_points_concat = torch.cat(new_points_list, dim=1)
return new_xyz, new_points_concat
PointNetFeaturePropagation
Ref to Point Feature Propagation for Set Segmentation part
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
class PointNetFeaturePropagation(nn.Module):
def __init__(self, in_channel, mlp):
super(PointNetFeaturePropagation, self).__init__()
self.mlp_convs = nn.ModuleList()
self.mlp_bns = nn.ModuleList()
last_channel = in_channel
for out_channel in mlp:
self.mlp_convs.append(nn.Conv1d(last_channel, out_channel, 1))
self.mlp_bns.append(nn.BatchNorm1d(out_channel))
last_channel = out_channel
def forward(self, xyz1, xyz2, points1, points2):
"""
Input:
xyz1: input points position data, [B, C, N]
xyz2: sampled input points position data, [B, C, S]
points1: input points data, [B, D, N]
points2: input points data, [B, D, S]
Return:
new_points: upsampled points data, [B, D', N]
"""
xyz1 = xyz1.permute(0, 2, 1)
xyz2 = xyz2.permute(0, 2, 1)
points2 = points2.permute(0, 2, 1)
B, N, C = xyz1.shape
_, S, _ = xyz2.shape
if S == 1:
interpolated_points = points2.repeat(1, N, 1)
else:
dists = square_distance(xyz1, xyz2)
dists, idx = dists.sort(dim=-1)
dists, idx = dists[:, :, :3], idx[:, :, :3] # [B, N, 3]
dist_recip = 1.0 / (dists + 1e-8)
norm = torch.sum(dist_recip, dim=2, keepdim=True)
weight = dist_recip / norm
interpolated_points = torch.sum(index_points(points2, idx) * weight.view(B, N, 3, 1), dim=2)
if points1 is not None:
points1 = points1.permute(0, 2, 1)
new_points = torch.cat([points1, interpolated_points], dim=-1)
else:
new_points = interpolated_points
new_points = new_points.permute(0, 2, 1)
for i, conv in enumerate(self.mlp_convs):
bn = self.mlp_bns[i]
new_points = F.relu(bn(conv(new_points)))
return new_points
PointNet2ClsMSG
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
class get_model(nn.Module):
def __init__(self,num_class,normal_channel=True):
super(get_model, self).__init__()
in_channel = 3 if normal_channel else 0
self.normal_channel = normal_channel
self.sa1 = PointNetSetAbstractionMsg(512, [0.1, 0.2, 0.4], [16, 32, 128], in_channel,[[32, 32, 64], [64, 64, 128], [64, 96, 128]])
self.sa2 = PointNetSetAbstractionMsg(128, [0.2, 0.4, 0.8], [32, 64, 128], 320,[[64, 64, 128], [128, 128, 256], [128, 128, 256]])
self.sa3 = PointNetSetAbstraction(None, None, None, 640 + 3, [256, 512, 1024], True)
self.fc1 = nn.Linear(1024, 512)
self.bn1 = nn.BatchNorm1d(512)
self.drop1 = nn.Dropout(0.4)
self.fc2 = nn.Linear(512, 256)
self.bn2 = nn.BatchNorm1d(256)
self.drop2 = nn.Dropout(0.5)
self.fc3 = nn.Linear(256, num_class)
def forward(self, xyz):
B, _, _ = xyz.shape
if self.normal_channel:
norm = xyz[:, 3:, :]
xyz = xyz[:, :3, :]
else:
norm = None
l1_xyz, l1_points = self.sa1(xyz, norm)
l2_xyz, l2_points = self.sa2(l1_xyz, l1_points)
l3_xyz, l3_points = self.sa3(l2_xyz, l2_points)
x = l3_points.view(B, 1024)
x = self.drop1(F.relu(self.bn1(self.fc1(x))))
x = self.drop2(F.relu(self.bn2(self.fc2(x))))
x = self.fc3(x)
x = F.log_softmax(x, -1)
return x,l3_points
PointNet2Seg
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
class get_model(nn.Module):
def __init__(self, num_classes, normal_channel=False):
super(get_model, self).__init__()
if normal_channel:
additional_channel = 3
else:
additional_channel = 0
self.normal_channel = normal_channel
self.sa1 = PointNetSetAbstractionMsg(512, [0.1, 0.2, 0.4], [32, 64, 128], 3+additional_channel, [[32, 32, 64], [64, 64, 128], [64, 96, 128]])
self.sa2 = PointNetSetAbstractionMsg(128, [0.4,0.8], [64, 128], 128+128+64, [[128, 128, 256], [128, 196, 256]])
self.sa3 = PointNetSetAbstraction(npoint=None, radius=None, nsample=None, in_channel=512 + 3, mlp=[256, 512, 1024], group_all=True)
self.fp3 = PointNetFeaturePropagation(in_channel=1536, mlp=[256, 256])
self.fp2 = PointNetFeaturePropagation(in_channel=576, mlp=[256, 128])
self.fp1 = PointNetFeaturePropagation(in_channel=150+additional_channel, mlp=[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, num_classes, 1)
def forward(self, xyz, cls_label):
# Set Abstraction layers
B,C,N = xyz.shape
if self.normal_channel:
l0_points = xyz
l0_xyz = xyz[:,:3,:]
else:
l0_points = xyz
l0_xyz = xyz
l1_xyz, l1_points = self.sa1(l0_xyz, l0_points)
l2_xyz, l2_points = self.sa2(l1_xyz, l1_points)
l3_xyz, l3_points = self.sa3(l2_xyz, l2_points)
# Feature Propagation layers
l2_points = self.fp3(l2_xyz, l3_xyz, l2_points, l3_points)
l1_points = self.fp2(l1_xyz, l2_xyz, l1_points, l2_points)
cls_label_one_hot = cls_label.view(B,16,1).repeat(1,1,N)
l0_points = self.fp1(l0_xyz, l1_xyz, torch.cat([cls_label_one_hot,l0_xyz,l0_points],1), l1_points)
# FC layers
feat = F.relu(self.bn1(self.conv1(l0_points)))
x = self.drop1(feat)
x = self.conv2(x)
x = F.log_softmax(x, dim=1)
x = x.permute(0, 2, 1)
return x, l3_points
3. Frustum PointNets for 3D Object Detection from RGB-D Data
While PointNets are capable of classifying a whole point cloud or predicting a semantic class for each point in a point cloud , it is unclear how this architecture can be used for instance-level 3D object detection . Towards this goal, we have to address one key challenge: how to efficiently propose possible locations of 3D objects in a 3D space . Imitating the practice in image detection, it is straightforward to enumerate candidate 3D boxes by sliding windows or by 3D region proposal networks. However, the computational complexity of 3D search typically grows cubically with respect to resolution and becomes too expensive for large scenes or real-time applications such as autonomous driving.
Instead, in this work, we reduce the search space following the dimension reduction principle: we take the advantage of mature 2D object detectors. First, we extract the 3D bounding frustum of an object by extruding 2D bounding boxes from image detectors. Then, within the 3D space trimmed by each of the 3D frustums, we consecutively perform 3D object instance segmentation and amodal 3D bounding box regression using two variants of PointNet. The segmentation network predicts the 3D mask of the object of interest (i.e. instance segmentation); and the regression network estimates the amodal 3D bounding box (covering the entire object even if only part of it is visible).
In contrast to previous work that treats RGB-D data as 2D maps for CNNs, our method is more 3D-centric as we lift depth maps to 3D point clouds and process them using 3D tools. This 3D-centric view enables new capabilities for exploring 3D data in a more effective manner. First, in our pipeline, a few transformations are applied successively on 3D coordinates, which align point clouds into a sequence of more constrained and canonical frames. These alignments factor out pose variations in data, and thus make 3D geometry pattern more evident, leading to an easier job of 3D learners. Second, learning in 3D space can better exploits the geometric and topological structure of 3D space. In principle, all objects live in 3D space; therefore, we believe that many geometric structures, such as repetition, planarity, and symmetry, are more naturally parameterized and captured by learners that directly operate in 3D space. The usefulness of this 3D-centric network design philosophy has been supported by much recent experimental evidence.
While PointNets have been applied to single object classification and semantic segmentation, our work explores how to extend the architecture for the purpose of 3D object detection .
3.1. Problem Statement
Given RGB-D data as input, our goal is to classify and localize objects in 3D space. The depth data, obtained from LiDAR or indoor depth sensors, is represented as a point cloud in RGB camera coordinates. The projection matrix is also known so that we can get a 3D frustum from a 2D image region. Each object is represented by a class (one among $k$ predefined classes) and an amodal 3D bounding box. The amodal box bounds the complete object even if part of the object is occluded or truncated. The 3D box is parameterized by its size $h$, $w$, $l$, center $c_x$, $c_y$, $c_z$, and orientation $θ$, $φ$, $ψ$ relative to a predefined canonical pose for each category. In our implementation, we only consider the heading angle $θ$ around the up-axis for orientation.
3.2. 3D Detection with Frustum PointNets
Our system for 3D object detection consists of three modules: frustum proposal, 3D instance segmentation, and 3D amodal bounding box estimation. We will introduce each module in the following subsections. We will focus on the pipeline and functionality of each module, and refer readers to supplementary for specific architectures of the deep networks involved.
3.2.1. Frustum Proposal
The resolution of data produced by most 3D sensors, especially real-time depth sensors, is still lower than RGB images from commodity cameras. Therefore, we leverage mature 2D object detector to propose 2D object regions in RGB images as well as to classify objects. With a known camera projection matrix, a 2D bounding box can be lifted to a frustum (with near and far planes specified by depth sensor range) that defines a 3D search space for the object. We then collect all points within the frustum to form a frustum point cloud . As shown in Fig 4 (a), frustums may orient towards many different directions, which result in large variation in the placement of point clouds. We therefore normalize the frustums by rotating them toward a center view such that the center axis of the frustum is orthogonal to the image plane. This normalization helps improve the rotation-invariance of the algorithm. We call this entire procedure for extracting frustum point clouds from RGB-D data frustum proposal generation .
3.2.2. 3D Instance Segmentation
Given a 2D image region (and its corresponding 3D frustum), several methods might be used to obtain 3D location of the object: One straightforward solution is to directly regress 3D object locations (e.g., by 3D bounding box) from a depth map using 2D CNNs. However, this problem is not easy as occluding objects and background clutter is common in natural scenes (as in Fig. 3), which may severely distract the 3D localization task. Because objects are naturally separated in physical space, segmentation in 3D point cloud is much more natural and easier than that in images where pixels from distant objects can be near-by to each other. Having observed this fact, we propose to segment instances in 3D point cloud instead of in 2D image or depth map. Similar to Mask-RCNN, which achieves instance segmentation by binary classification of pixels in image regions, we realize 3D instance segmentation using a PointNet-based network on point clouds in frustums.
Based on 3D instance segmentation, we are able to achieve residual based 3D localization. That is, rather than regressing the absolute 3D location of the object whose offset from the sensor may vary in large ranges (e.g. from 5m to beyond 50m in KITTI data), we predict the 3D bounding box center in a local coordinate system – 3D mask coordinates as shown in Fig. 4 (c).
3D Instance Segmentation PointNet
The network takes a point cloud in frustum and predicts a probability score for each point that indicates how likely the point belongs to the object of interest. Note that each frustum contains exactly one object of interest. Here those “other” points could be points of non-relevant areas (such as ground, vegetation) or other instances that occlude or are behind the object of interest. Similar to the case in 2D instance segmentation, depending on the position of the frustum, object points in one frustum may become cluttered or occlude points in another. Therefore, our segmentation PointNet is learning the occlusion and clutter patterns as well as recognizing the geometry for the object of a certain category.
In a multi-class detection case, we also leverage the semantics from a 2D detector for better instance segmentation. For example, if we know the object of interest is a pedestrian, then the segmentation network can use this prior to find geometries that look like a person. Specifically, in our architecture we encode the semantic category as a one-hot class vector (k dimensional for the pre-defined k categories) and concatenate the one-hot vector to the intermediate point cloud features.
After 3D instance segmentation, points that are classified as the object of interest are extracted (“masking”).
Having obtained these segmented object points, we further normalize its coordinates to boost the translational invariance of the algorithm, following the same rationale as in the frustum proposal step. In our implementation, we transform the point cloud into a local coordinate by subtracting XYZ values by its centroid. This is illustrated in Fig. 4 (c). Note that we intentionally do not scale the point cloud, because the bounding sphere size of a partial point cloud can be greatly affected by viewpoints and the real size of the point cloud helps the box size estimation.
3.2.3. 3D Amodal Box Estimation
Given the segmented object points (in 3D mask coordinate), this module estimates the object’s amodal oriented 3D bounding box by using a box regression PointNet together with a preprocessing transformer network.
Learning-based 3D Alignment by T-Net: Even though we have aligned segmented object points according to their centroid position, we find that the origin of the mask coordinate frame (Fig. 4 (c)) may still be quite far from the amodal box center. We therefore propose to use a light-weight regression PointNet (T-Net) to estimate the true center of the complete object and then transform the coordinate such that the predicted center becomes the origin (Fig. 4 (d)).
Amodal 3D Box Estimation PointNet: The box estimation network predicts amodal bounding boxes (for entire object even if part of it is unseen) for objects given an object point cloud in 3D object coordinate (Fig. 4 (d)). The network architecture is similar to that for object classification, however the output is no longer object class scores but parameters for a 3D bounding box.
As stated in Sec. 3, we parameterize a 3D bounding box by its center $(c_x, c_y, c_z)$, size $(h, w, l)$ and heading angle $θ$ (along up-axis). We take a “residual” approach for box center estimation. The center residual predicted by the box estimation network is combined with the previous center residual from the T-Net and the masked points’ centroid to recover an absolute center (Eq. 1). For box size and heading angle, we follow previous works and use a hybrid of classification and regression formulations. Specifically we pre-define $N S$ size templates and $N H$ equally split angle bins. Our model will both classify size/heading ($N S$ scores for size, $N H$ scores for heading) to those pre-defined categories as well as predict residual numbers for each category ($3×NS$ residual dimensions for height, width, length, NH residual angles for heading). In the end the net outputs $3 + 4 × NS + 2 × NH$ numbers in total.
3.3. Training with Multi-task Losses
We simultaneously optimize the three nets involved (3D instance segmentation PointNet, T-Net and amodal box estimation PointNet) with multi-task losses (as in Eq. 2). $L_{c1−reg}$ is for T-Net and $L_{c2−reg}$ is for center regression of box estimation net. $L_{h−cls}$ and $L_{h−reg}$ are losses for heading angle prediction while $L_{s−cls}$ and $L_{s−reg}$ are for box size. Softmax is used for all classification tasks and smooth-$l_1$ (huber) loss is used for all regression cases.
Corner Loss for Joint Optimization of Box Parameters: While our 3D bounding box parameterization is compact and complete, learning is not optimized for final 3D box accuracy – center, size and heading have separate loss terms. Imagine cases where center and size are accurately predicted but heading angle is off – the 3D IoU with ground truth box will then be dominated by the angle error. Ideally all three terms (center,size,heading) should be jointly optimized for best 3D box estimation (under IoU metric). To resolve this problem we propose a novel regularization loss, the corner loss:
In essence, the corner loss is the sum of the distances between the eight corners of a predicted box and a ground truth box. Since corner positions are jointly determined by center, size and heading, the corner loss is able to regularize the multi-task training for those parameters.
To compute the corner loss, we firstly construct $NS × NH$ “anchor” boxes from all size templates and heading angle bins. The anchor boxes are then translated to the estimated box center. We denote the anchor box corners as $P^{ij}k$, where $i$, $j$, $k$ are indices for the size class, heading class, and (predefined) corner order, respectively. To avoid large penalty from flipped heading estimation, we further compute distances to corners ($P^{∗∗}_k$) from the flipped ground truth box and use the minimum of the original and flipped cases. $δ{ij}$ , which is one for the ground truth size/heading class and zero else wise, is a two-dimensional mask used to select the distance term we care about.
3.4. Experiments
Code Explanation
Frustum Proposal
This step depend on dataset:
3D Instance Segmentation
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
class PointNetInstanceSeg(nn.Module):
def __init__(self,n_classes=3,n_channel=3):
'''v1 3D Instance Segmentation PointNet
:param n_classes:3
:param one_hot_vec:[bs,n_classes]
'''
super(PointNetInstanceSeg, self).__init__()
self.conv1 = nn.Conv1d(n_channel, 64, 1)
self.conv2 = nn.Conv1d(64, 64, 1)
self.conv3 = nn.Conv1d(64, 64, 1)
self.conv4 = nn.Conv1d(64, 128, 1)
self.conv5 = nn.Conv1d(128, 1024, 1)
self.bn1 = nn.BatchNorm1d(64)
self.bn2 = nn.BatchNorm1d(64)
self.bn3 = nn.BatchNorm1d(64)
self.bn4 = nn.BatchNorm1d(128)
self.bn5 = nn.BatchNorm1d(1024)
self.n_classes = n_classes
self.dconv1 = nn.Conv1d(1088+n_classes, 512, 1)
self.dconv2 = nn.Conv1d(512, 256, 1)
self.dconv3 = nn.Conv1d(256, 128, 1)
self.dconv4 = nn.Conv1d(128, 128, 1)
self.dropout = nn.Dropout(p=0.5)
self.dconv5 = nn.Conv1d(128, 2, 1)
self.dbn1 = nn.BatchNorm1d(512)
self.dbn2 = nn.BatchNorm1d(256)
self.dbn3 = nn.BatchNorm1d(128)
self.dbn4 = nn.BatchNorm1d(128)
def forward(self, pts, one_hot_vec): # bs,4,n
'''
:param pts: [bs,4,n]: x,y,z,intensity
:return: logits: [bs,n,2],scores for bkg/clutter and object
'''
bs = pts.size()[0]
n_pts = pts.size()[2]
out1 = F.relu(self.bn1(self.conv1(pts))) # bs,64,n
out2 = F.relu(self.bn2(self.conv2(out1))) # bs,64,n
out3 = F.relu(self.bn3(self.conv3(out2))) # bs,64,n
out4 = F.relu(self.bn4(self.conv4(out3)))# bs,128,n
out5 = F.relu(self.bn5(self.conv5(out4)))# bs,1024,n
global_feat = torch.max(out5, 2, keepdim=True)[0] #bs,1024,1
expand_one_hot_vec = one_hot_vec.view(bs,-1,1)#bs,3,1
expand_global_feat = torch.cat([global_feat, expand_one_hot_vec],1)#bs,1027,1
expand_global_feat_repeat = expand_global_feat.view(bs,-1,1)\
.repeat(1,1,n_pts)# bs,1027,n
concat_feat = torch.cat([out2,\
expand_global_feat_repeat],1)
# bs, (641024+3)=1091, n
x = F.relu(self.dbn1(self.dconv1(concat_feat)))#bs,512,n
x = F.relu(self.dbn2(self.dconv2(x)))#bs,256,n
x = F.relu(self.dbn3(self.dconv3(x)))#bs,128,n
x = F.relu(self.dbn4(self.dconv4(x)))#bs,128,n
x = self.dropout(x)
x = self.dconv5(x)#bs, 2, n
seg_pred = x.transpose(2,1).contiguous()#bs, n, 2
return seg_pred
Amodal 3D Box Estimation
Learning-based 3D Alignment by T-Net
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
class STNxyz(nn.Module):
def __init__(self,n_classes=3):
super(STNxyz, self).__init__()
self.conv1 = torch.nn.Conv1d(3, 128, 1)
self.conv2 = torch.nn.Conv1d(128, 128, 1)
self.conv3 = torch.nn.Conv1d(128, 256, 1)
#self.conv4 = torch.nn.Conv1d(256, 512, 1)
self.fc1 = nn.Linear(256+n_classes, 256)
self.fc2 = nn.Linear(256, 128)
self.fc3 = nn.Linear(128, 3)
init.zeros_(self.fc3.weight)
init.zeros_(self.fc3.bias)
self.bn1 = nn.BatchNorm1d(128)
self.bn2 = nn.BatchNorm1d(128)
self.bn3 = nn.BatchNorm1d(256)
self.fcbn1 = nn.BatchNorm1d(256)
self.fcbn2 = nn.BatchNorm1d(128)
def forward(self, pts,one_hot_vec):
bs = pts.shape[0]
x = F.relu(self.bn1(self.conv1(pts)))# bs,128,n
x = F.relu(self.bn2(self.conv2(x)))# bs,128,n
x = F.relu(self.bn3(self.conv3(x)))# bs,256,n
x = torch.max(x, 2)[0]# bs,256
expand_one_hot_vec = one_hot_vec.view(bs, -1)# bs,3
x = torch.cat([x, expand_one_hot_vec],1)#bs,259
x = F.relu(self.fcbn1(self.fc1(x)))# bs,256
x = F.relu(self.fcbn2(self.fc2(x)))# bs,128
x = self.fc3(x)# bs,
return x
Amodal 3D Box Estimation PointNet
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
class PointNetEstimation(nn.Module):
def __init__(self,n_classes=2):
'''v1 Amodal 3D Box Estimation Pointnet
:param n_classes:3
:param one_hot_vec:[bs,n_classes]
'''
super(PointNetEstimation, self).__init__()
self.conv1 = nn.Conv1d(3, 128, 1)
self.conv2 = nn.Conv1d(128, 128, 1)
self.conv3 = nn.Conv1d(128, 256, 1)
self.conv4 = nn.Conv1d(256, 512, 1)
self.bn1 = nn.BatchNorm1d(128)
self.bn2 = nn.BatchNorm1d(128)
self.bn3 = nn.BatchNorm1d(256)
self.bn4 = nn.BatchNorm1d(512)
self.n_classes = n_classes
self.fc1 = nn.Linear(512+3, 512)
self.fc2 = nn.Linear(512, 256)
self.fc3 = nn.Linear(256,3+NUM_HEADING_BIN*2+NUM_SIZE_CLUSTER*4)
self.fcbn1 = nn.BatchNorm1d(512)
self.fcbn2 = nn.BatchNorm1d(256)
def forward(self, pts,one_hot_vec): # bs,3,m
'''
:param pts: [bs,3,m]: x,y,z after InstanceSeg
:return: box_pred: [bs,3+NUM_HEADING_BIN*2+NUM_SIZE_CLUSTER*4]
including box centers, heading bin class scores and residual,
and size cluster scores and residual
'''
bs = pts.size()[0]
n_pts = pts.size()[2]
out1 = F.relu(self.bn1(self.conv1(pts))) # bs,128,n
out2 = F.relu(self.bn2(self.conv2(out1))) # bs,128,n
out3 = F.relu(self.bn3(self.conv3(out2))) # bs,256,n
out4 = F.relu(self.bn4(self.conv4(out3)))# bs,512,n
global_feat = torch.max(out4, 2, keepdim=False)[0] #bs,512
expand_one_hot_vec = one_hot_vec.view(bs,-1)#bs,3
expand_global_feat = torch.cat([global_feat, expand_one_hot_vec],1)#bs,515
x = F.relu(self.fcbn1(self.fc1(expand_global_feat)))#bs,512
x = F.relu(self.fcbn2(self.fc2(x))) # bs,256
box_pred = self.fc3(x) # bs,3+NUM_HEADING_BIN*2+NUM_SIZE_CLUSTER*4
return box_pred
Multi-task Losses
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
class FrustumPointNetLoss(nn.Module):
def __init__(self):
super(FrustumPointNetLoss, self).__init__()
def forward(self, logits, mask_label, \
center, center_label, stage1_center, \
heading_scores, heading_residual_normalized, heading_residual, \
heading_class_label, heading_residual_label, \
size_scores,size_residual_normalized,size_residual,
size_class_label,size_residual_label,
corner_loss_weight=10.0,box_loss_weight=1.0):
'''
1.InsSeg
logits: torch.Size([32, 1024, 2]) torch.float32
mask_label: [32, 1024]
2.Center
center: torch.Size([32, 3]) torch.float32
stage1_center: torch.Size([32, 3]) torch.float32
center_label:[32,3]
3.Heading
heading_scores: torch.Size([32, 12]) torch.float32
heading_residual_snormalized: torch.Size([32, 12]) torch.float32
heading_residual: torch.Size([32, 12]) torch.float32
heading_class_label:(32)
heading_residual_label:(32)
4.Size
size_scores: torch.Size([32, 8]) torch.float32
size_residual_normalized: torch.Size([32, 8, 3]) torch.float32
size_residual: torch.Size([32, 8, 3]) torch.float32
size_class_label:(32)
size_residual_label:(32,3)
5.Corner
6.Weight
corner_loss_weight: float scalar
box_loss_weight: float scalar
'''
bs = logits.shape[0]
# 3D Instance Segmentation PointNet Loss
logits = F.log_softmax(logits.view(-1,2),dim=1)#torch.Size([32768, 2])
mask_label = mask_label.view(-1).long()#torch.Size([32768])
mask_loss = F.nll_loss(logits, mask_label)#tensor(0.6361, grad_fn=<NllLossBackward>)
# Center Regression Loss
center_dist = torch.norm(center-center_label,dim=1)#(32,)
center_loss = huber_loss(center_dist, delta=2.0)
stage1_center_dist = torch.norm(center-stage1_center,dim=1)#(32,)
stage1_center_loss = huber_loss(stage1_center_dist, delta=1.0)
# Heading Loss
heading_class_loss = F.nll_loss(F.log_softmax(heading_scores,dim=1), \
heading_class_label.long())#tensor(2.4505, grad_fn=<NllLossBackward>)
hcls_onehot = torch.eye(NUM_HEADING_BIN)[heading_class_label.long()].cuda() # 32,12
heading_residual_normalized_label = \
heading_residual_label / (np.pi / NUM_HEADING_BIN) # 32,
heading_residual_normalized_dist = torch.sum( \
heading_residual_normalized * hcls_onehot.float(), dim=1) # 32,
### Only compute reg loss on gt label
heading_residual_normalized_loss = \
huber_loss(heading_residual_normalized_dist -
heading_residual_normalized_label, delta=1.0)###fix,2020.1.14
# Size loss
size_class_loss = F.nll_loss(F.log_softmax(size_scores,dim=1),\
size_class_label.long())#tensor(2.0240, grad_fn=<NllLossBackward>)
scls_onehot = torch.eye(NUM_SIZE_CLUSTER)[size_class_label.long()].cuda() # 32,8
scls_onehot_repeat = scls_onehot.view(-1, NUM_SIZE_CLUSTER, 1).repeat(1, 1, 3) # 32,8,3
predicted_size_residual_normalized_dist = torch.sum( \
size_residual_normalized * scls_onehot_repeat.cuda(), dim=1)#32,3
mean_size_arr_expand = torch.from_numpy(g_mean_size_arr).float().cuda() \
.view(1, NUM_SIZE_CLUSTER, 3) # 1,8,3
mean_size_label = torch.sum(scls_onehot_repeat * mean_size_arr_expand, dim=1)# 32,3
size_residual_label_normalized = size_residual_label / mean_size_label.cuda()
size_normalized_dist = torch.norm(size_residual_label_normalized-\
predicted_size_residual_normalized_dist,dim=1)#32
size_residual_normalized_loss = huber_loss(size_normalized_dist, delta=1.0)#tensor(11.2784, grad_fn=<MeanBackward0>)
# Corner Loss
corners_3d = get_box3d_corners(center,\
heading_residual,size_residual).cuda()#(bs,NH,NS,8,3)(32, 12, 8, 8, 3)
gt_mask = hcls_onehot.view(bs,NUM_HEADING_BIN,1).repeat(1,1,NUM_SIZE_CLUSTER) * \
scls_onehot.view(bs,1,NUM_SIZE_CLUSTER).repeat(1,NUM_HEADING_BIN,1)# (bs,NH=12,NS=8)
corners_3d_pred = torch.sum(\
gt_mask.view(bs,NUM_HEADING_BIN,NUM_SIZE_CLUSTER,1,1)\
.float().cuda() * corners_3d,\
dim=[1, 2]) # (bs,8,3)
heading_bin_centers = torch.from_numpy(\
np.arange(0, 2 * np.pi, 2 * np.pi / NUM_HEADING_BIN)).float().cuda() # (NH,)
heading_label = heading_residual_label.view(bs,1) + \
heading_bin_centers.view(1,NUM_HEADING_BIN) #(bs,1)+(1,NH)=(bs,NH)
heading_label = torch.sum(hcls_onehot.float() * heading_label, 1)
mean_sizes = torch.from_numpy(g_mean_size_arr)\
.float().view(1,NUM_SIZE_CLUSTER,3).cuda()#(1,NS,3)
size_label = mean_sizes + \
size_residual_label.view(bs,1,3) #(1,NS,3)+(bs,1,3)=(bs,NS,3)
size_label = torch.sum(\
scls_onehot.view(bs,NUM_SIZE_CLUSTER,1).float() * size_label, axis=[1]) # (B,3)
corners_3d_gt = get_box3d_corners_helper( \
center_label, heading_label, size_label) # (B,8,3)
corners_3d_gt_flip = get_box3d_corners_helper( \
center_label, heading_label + np.pi, size_label) # (B,8,3)
corners_dist = torch.min(torch.norm(corners_3d_pred - corners_3d_gt, dim=-1),
torch.norm(corners_3d_pred - corners_3d_gt_flip, dim=-1))
corners_loss = huber_loss(corners_dist, delta=1.0)
# Weighted sum of all losses
total_loss = mask_loss + box_loss_weight * (center_loss + \
heading_class_loss + size_class_loss + \
heading_residual_normalized_loss * 20 + \
size_residual_normalized_loss * 20 + \
stage1_center_loss + \
corner_loss_weight * corners_loss)
###total_loss = mask_loss
#tensor(306.7591, grad_fn=<AddBackward0>)
###if np.isnan(total_loss.item()) or total_loss > 10000.0:
### ipdb.set_trace()
losses = {
'total_loss': total_loss,
'mask_loss': mask_loss,
'mask_loss': box_loss_weight * center_loss,
'heading_class_loss': box_loss_weight * heading_class_loss,
'size_class_loss': box_loss_weight * size_class_loss,
'heading_residual_normalized_loss': box_loss_weight * heading_residual_normalized_loss * 20,
'size_residual_normalized_loss': box_loss_weight * size_residual_normalized_loss * 20,
'stage1_center_loss': box_loss_weight * size_residual_normalized_loss * 20,
'corners_loss': box_loss_weight * corners_loss * corner_loss_weight,
}
return losses
'''
return total_loss, mask_loss, center_loss, heading_class_loss, \
size_class_loss, heading_residual_normalized_loss, \
size_residual_normalized_loss, stage1_center_loss, \
corners_loss
'''
Training
