mirror of
https://github.com/wassname/PointCNN.git
synced 2026-09-10 11:40:19 +08:00
Significantly clean and document code base. Project approximately complete.
This commit is contained in:
Austin J. Garrett
+243
@@ -0,0 +1,243 @@
|
||||
"""
|
||||
Author: Austin J. Garrett
|
||||
|
||||
PyTorch implementation of the PointCNN paper, as specified in:
|
||||
https://arxiv.org/pdf/1801.07791.pdf
|
||||
Original paper by: Yangyan Li, Rui Bu, Mingchao Sun, Baoquan Chen
|
||||
"""
|
||||
|
||||
# External Modules
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from torch import FloatTensor
|
||||
import numpy as np
|
||||
from typing import Tuple, Callable, Optional
|
||||
|
||||
# Internal Modules
|
||||
from PointCNN.core.util_funcs import UFloatTensor, ULongTensor
|
||||
from PointCNN.core.util_layers import Conv, SepConv, Dense, EndChannels
|
||||
|
||||
class XConv(nn.Module):
|
||||
""" Convolution over a single point and its neighbors. """
|
||||
|
||||
def __init__(self, C_in : int, C_out : int, dims : int, K : int,
|
||||
P : int, C_mid : int, depth_multiplier : int) -> None:
|
||||
"""
|
||||
:param C_in: Input dimension of the points' features.
|
||||
:param C_out: Output dimension of the representative point features.
|
||||
:param dims: Spatial dimensionality of points.
|
||||
:param K: Number of neighbors to convolve over.
|
||||
:param P: Number of representative points.
|
||||
:param C_mid: Dimensionality of lifted point features.
|
||||
:param depth_multiplier: Depth multiplier for internal depthwise separable convolution.
|
||||
"""
|
||||
super(XConv, self).__init__()
|
||||
|
||||
if __debug__:
|
||||
# Only needed for assertions.
|
||||
self.C_in = C_in
|
||||
self.C_mid = C_mid
|
||||
self.dims = dims
|
||||
self.K = K
|
||||
|
||||
self.P = P
|
||||
|
||||
# Additional processing layers
|
||||
# self.pts_layernorm = LayerNorm(2, momentum = 0.9)
|
||||
|
||||
# Main dense linear layers
|
||||
self.dense1 = Dense(dims, C_mid)
|
||||
self.dense2 = Dense(C_mid, C_mid)
|
||||
|
||||
# Layers to generate X
|
||||
self.x_trans = nn.Sequential(
|
||||
EndChannels(Conv(
|
||||
in_channels = dims,
|
||||
out_channels = K*K,
|
||||
kernel_size = (1, K),
|
||||
with_bn = False
|
||||
)),
|
||||
Dense(K*K, K*K, with_bn = False),
|
||||
Dense(K*K, K*K, with_bn = False, activation = None)
|
||||
)
|
||||
|
||||
self.end_conv = EndChannels(SepConv(
|
||||
in_channels = C_mid + C_in,
|
||||
out_channels = C_out,
|
||||
kernel_size = (1, K),
|
||||
depth_multiplier = depth_multiplier
|
||||
)).cuda()
|
||||
|
||||
def forward(self, x : Tuple[UFloatTensor, # (N, P, dims)
|
||||
UFloatTensor, # (N, P, K, dims)
|
||||
Optional[UFloatTensor]] # (N, P, K, C_in)
|
||||
) -> UFloatTensor: # (N, K, C_out)
|
||||
"""
|
||||
Applies XConv to the input data.
|
||||
:param x: (rep_pt, pts, fts) where
|
||||
- rep_pt: Representative point.
|
||||
- pts: Regional point cloud such that fts[:,p_idx,:] is the feature
|
||||
associated with pts[:,p_idx,:].
|
||||
- fts: Regional features such that pts[:,p_idx,:] is the feature
|
||||
associated with fts[:,p_idx,:].
|
||||
:return: Features aggregated into point rep_pt.
|
||||
"""
|
||||
rep_pt, pts, fts = x
|
||||
|
||||
if fts is not None:
|
||||
assert(rep_pt.size()[0] == pts.size()[0] == fts.size()[0]) # Check N is equal.
|
||||
assert(rep_pt.size()[1] == pts.size()[1] == fts.size()[1]) # Check P is equal.
|
||||
assert(pts.size()[2] == fts.size()[2] == self.K) # Check K is equal.
|
||||
assert(fts.size()[3] == self.C_in) # Check C_in is equal.
|
||||
else:
|
||||
assert(rep_pt.size()[0] == pts.size()[0]) # Check N is equal.
|
||||
assert(rep_pt.size()[1] == pts.size()[1]) # Check P is equal.
|
||||
assert(pts.size()[2] == self.K) # Check K is equal.
|
||||
assert(rep_pt.size()[2] == pts.size()[3] == self.dims) # Check dims is equal.
|
||||
|
||||
N = len(pts)
|
||||
P = rep_pt.size()[1] # (N, P, K, dims)
|
||||
p_center = torch.unsqueeze(rep_pt, dim = 2) # (N, P, 1, dims)
|
||||
|
||||
# Move pts to local coordinate system of rep_pt.
|
||||
pts_local = pts - p_center # (N, P, K, dims)
|
||||
# pts_local = self.pts_layernorm(pts - p_center)
|
||||
|
||||
# Individually lift each point into C_mid space.
|
||||
fts_lifted0 = self.dense1(pts_local)
|
||||
fts_lifted = self.dense2(fts_lifted0) # (N, P, K, C_mid)
|
||||
|
||||
if fts is None:
|
||||
fts_cat = fts_lifted
|
||||
else:
|
||||
fts_cat = torch.cat((fts_lifted, fts), -1) # (N, P, K, C_mid + C_in)
|
||||
|
||||
# Learn the (N, K, K) X-transformation matrix.
|
||||
X_shape = (N, P, self.K, self.K)
|
||||
X = self.x_trans(pts_local)
|
||||
X = X.view(*X_shape)
|
||||
|
||||
# Weight and permute fts_cat with the learned X.
|
||||
fts_X = torch.matmul(X, fts_cat)
|
||||
fts_p = self.end_conv(fts_X).squeeze(dim = 2)
|
||||
return fts_p
|
||||
|
||||
class PointCNN(nn.Module):
|
||||
""" Pointwise convolutional model. """
|
||||
|
||||
def __init__(self, C_in : int, C_out : int, dims : int, K : int, D : int, P : int,
|
||||
r_indices_func : Callable[[UFloatTensor, # (N, P, dims)
|
||||
UFloatTensor, # (N, x, dims)
|
||||
int, int],
|
||||
ULongTensor] # (N, P, K)
|
||||
) -> None:
|
||||
"""
|
||||
:param C_in: Input dimension of the points' features.
|
||||
:param C_out: Output dimension of the representative point features.
|
||||
:param dims: Spatial dimensionality of points.
|
||||
:param K: Number of neighbors to convolve over.
|
||||
:param D: "Spread" of neighboring points.
|
||||
:param P: Number of representative points.
|
||||
:param r_indices_func: Selector function of the type,
|
||||
INPUTS
|
||||
rep_pts : Representative points.
|
||||
pts : Point cloud.
|
||||
K : Number of points for each region.
|
||||
D : "Spread" of neighboring points.
|
||||
|
||||
OUTPUT
|
||||
pts_idx : Array of indices into pts such that pts[pts_idx] is the set
|
||||
of points in the "region" around rep_pt.
|
||||
"""
|
||||
super(PointCNN, self).__init__()
|
||||
|
||||
C_mid = C_out // 2 if C_in == 0 else C_out // 4
|
||||
depth_multiplier = min(int(np.ceil(C_out / C_in)), 4)
|
||||
|
||||
self.r_indices_func = lambda rep_pts, pts: r_indices_func(rep_pts, pts, K, D)
|
||||
self.dense = Dense(C_in, C_out // 2) if C_in != 0 else None
|
||||
self.x_conv = XConv(C_out // 2 if C_in != 0 else C_in, C_out, dims, K, P, C_mid, depth_multiplier)
|
||||
self.D = D
|
||||
|
||||
def select_region(self, pts : UFloatTensor, # (N, x, dims)
|
||||
pts_idx : ULongTensor # (N, P, K)
|
||||
) -> UFloatTensor: # (P, K, dims)
|
||||
"""
|
||||
Selects neighborhood points based on output of r_indices_func.
|
||||
:param pts: Point cloud to select regional points from.
|
||||
:param pts_idx: Indices of points in region to be selected.
|
||||
:return: Local neighborhoods around each representative point.
|
||||
"""
|
||||
regions = torch.stack([
|
||||
pts[n][idx,:] for n, idx in enumerate(torch.unbind(pts_idx, dim = 0))
|
||||
], dim = 0)
|
||||
return regions
|
||||
|
||||
def forward(self, x : Tuple[FloatTensor, # (N, P, dims)
|
||||
FloatTensor, # (N, x, dims)
|
||||
FloatTensor] # (N, x, C_in)
|
||||
) -> FloatTensor: # (N, P, C_out)
|
||||
"""
|
||||
Given a set of representative points, a point cloud, and its
|
||||
corresponding features, return a new set of representative points with
|
||||
features projected from the point cloud.
|
||||
:param x: (rep_pts, pts, fts) where
|
||||
- rep_pts: Representative points.
|
||||
- pts: Regional point cloud such that fts[:,p_idx,:] is the
|
||||
feature associated with pts[:,p_idx,:].
|
||||
- fts: Regional features such that pts[:,p_idx,:] is the feature
|
||||
associated with fts[:,p_idx,:].
|
||||
:return: Features aggregated to rep_pts.
|
||||
"""
|
||||
rep_pts, pts, fts = x
|
||||
fts = self.dense(fts) if fts is not None else fts
|
||||
|
||||
# This step takes ~97% of the time. Prime target for optimization: KNN on GPU.
|
||||
pts_idx = self.r_indices_func(rep_pts.cpu(), pts.cpu()).cuda()
|
||||
# -------------------------------------------------------------------------- #
|
||||
|
||||
pts_regional = self.select_region(pts, pts_idx)
|
||||
fts_regional = self.select_region(fts, pts_idx) if fts is not None else fts
|
||||
fts_p = self.x_conv((rep_pts, pts_regional, fts_regional))
|
||||
|
||||
return fts_p
|
||||
|
||||
class RandPointCNN(nn.Module):
|
||||
""" PointCNN with randomly subsampled representative points. """
|
||||
|
||||
def __init__(self, C_in : int, C_out : int, dims : int, K : int, D : int, P : int,
|
||||
r_indices_func : Callable[[UFloatTensor, # (N, P, dims)
|
||||
UFloatTensor, # (N, x, dims)
|
||||
int, int],
|
||||
ULongTensor] # (N, P, K)
|
||||
) -> None:
|
||||
""" See documentation for PointCNN. """
|
||||
super(RandPointCNN, self).__init__()
|
||||
self.pointcnn = PointCNN(C_in, C_out, dims, K, D, P, r_indices_func)
|
||||
self.P = P
|
||||
|
||||
def forward(self, x : Tuple[UFloatTensor, # (N, x, dims)
|
||||
UFloatTensor] # (N, x, dims)
|
||||
) -> Tuple[UFloatTensor, # (N, P, dims)
|
||||
UFloatTensor]: # (N, P, C_out)
|
||||
"""
|
||||
Given a point cloud, and its corresponding features, return a new set
|
||||
of randomly-sampled representative points with features projected from
|
||||
the point cloud.
|
||||
:param x: (pts, fts) where
|
||||
- pts: Regional point cloud such that fts[:,p_idx,:] is the
|
||||
feature associated with pts[:,p_idx,:].
|
||||
- fts: Regional features such that pts[:,p_idx,:] is the feature
|
||||
associated with fts[:,p_idx,:].
|
||||
:return: Randomly subsampled points and their features.
|
||||
"""
|
||||
pts, fts = x
|
||||
if 0 < self.P < pts.size()[1]:
|
||||
# Select random set of indices of subsampled points.
|
||||
idx = np.random.choice(pts.size()[1], self.P, replace = False).tolist()
|
||||
rep_pts = pts[:,idx,:]
|
||||
else:
|
||||
# All input points are representative points.
|
||||
rep_pts = pts
|
||||
rep_pts_fts = self.pointcnn((rep_pts, pts, fts))
|
||||
return rep_pts, rep_pts_fts
|
||||
Reference in New Issue
Block a user