learning3d 0.0.1__py3-none-any.whl

learning3d/models/pointnet.py

@@ -0,0 +1,108 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from .pooling import Pooling
+
+
+class PointNet(torch.nn.Module):
+	def __init__(self, emb_dims=1024, input_shape="bnc", use_bn=False, global_feat=True):
+		# emb_dims:			Embedding Dimensions for PointNet.
+		# input_shape:		Shape of Input Point Cloud (b: batch, n: no of points, c: channels)
+		super(PointNet, self).__init__()
+		if input_shape not in ["bcn", "bnc"]:
+			raise ValueError("Allowed shapes are 'bcn' (batch * channels * num_in_points), 'bnc' ")
+		self.input_shape = input_shape
+		self.emb_dims = emb_dims
+		self.use_bn = use_bn
+		self.global_feat = global_feat
+		if not self.global_feat: self.pooling = Pooling('max')
+
+		self.layers = self.create_structure()
+
+	def create_structure(self):
+		self.conv1 = torch.nn.Conv1d(3, 64, 1)
+		self.conv2 = torch.nn.Conv1d(64, 64, 1)
+		self.conv3 = torch.nn.Conv1d(64, 64, 1)
+		self.conv4 = torch.nn.Conv1d(64, 128, 1)
+		self.conv5 = torch.nn.Conv1d(128, self.emb_dims, 1)
+		self.relu = torch.nn.ReLU()
+
+		if self.use_bn:
+			self.bn1 = torch.nn.BatchNorm1d(64)
+			self.bn2 = torch.nn.BatchNorm1d(64)
+			self.bn3 = torch.nn.BatchNorm1d(64)
+			self.bn4 = torch.nn.BatchNorm1d(128)
+			self.bn5 = torch.nn.BatchNorm1d(self.emb_dims)
+
+		if self.use_bn:
+			layers = [self.conv1, self.bn1, self.relu,
+					  self.conv2, self.bn2, self.relu,
+					  self.conv3, self.bn3, self.relu,
+					  self.conv4, self.bn4, self.relu,
+					  self.conv5, self.bn5, self.relu]
+		else:
+			layers = [self.conv1, self.relu,
+					  self.conv2, self.relu, 
+					  self.conv3, self.relu,
+					  self.conv4, self.relu,
+					  self.conv5, self.relu]
+		return layers
+
+
+	def forward(self, input_data):
+		# input_data: 		Point Cloud having shape input_shape.
+		# output:			PointNet features (Batch x emb_dims)
+		if self.input_shape == "bnc":
+			num_points = input_data.shape[1]
+			input_data = input_data.permute(0, 2, 1)
+		else:
+			num_points = input_data.shape[2]
+		if input_data.shape[1] != 3:
+			raise RuntimeError("shape of x must be of [Batch x 3 x NumInPoints]")
+
+		output = input_data
+		for idx, layer in enumerate(self.layers):
+			output = layer(output)
+			if idx == 1 and not self.global_feat: point_feature = output
+
+		if self.global_feat:
+			return output
+		else:
+			output = self.pooling(output)
+			output = output.view(-1, self.emb_dims, 1).repeat(1, 1, num_points)
+			return torch.cat([output, point_feature], 1)
+
+
+if __name__ == '__main__':
+	# Test the code.
+	x = torch.rand((10,1024,3))
+
+	pn = PointNet(use_bn=True)
+	y = pn(x)
+	print("Network Architecture: ")
+	print(pn)
+	print("Input Shape of PointNet: ", x.shape, "\nOutput Shape of PointNet: ", y.shape)
+
+	class PointNet_modified(PointNet):
+		def __init__(self):
+			super().__init__()
+
+		def create_structure(self):
+			self.conv1 = torch.nn.Conv1d(3, 64, 1)
+			self.conv2 = torch.nn.Conv1d(64, 128, 1)
+			self.conv3 = torch.nn.Conv1d(128, self.emb_dims, 1)
+			self.relu = torch.nn.ReLU()
+
+			layers = [self.conv1, self.relu,
+					  self.conv2, self.relu,
+					  self.conv3, self.relu]
+			return layers
+
+	pn = PointNet_modified()
+	y = pn(x)
+	print("\n\n\nModified Network Architecture: ")
+	print(pn)
+	print("Input Shape of PointNet: ", x.shape, "\nOutput Shape of PointNet: ", y.shape)
+
+
+

learning3d/models/pointnetlk.py

@@ -0,0 +1,173 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from .pointnet import PointNet
+from .pooling import Pooling
+from .. ops import data_utils
+from .. ops import se3, so3, invmat
+
+
+class PointNetLK(nn.Module):
+	def __init__(self, feature_model=PointNet(), delta=1.0e-2, learn_delta=False, xtol=1.0e-7, p0_zero_mean=True, p1_zero_mean=True, pooling='max'):
+		super().__init__()
+		self.feature_model = feature_model
+		self.pooling = Pooling(pooling)
+		self.inverse = invmat.InvMatrix.apply
+		self.exp = se3.Exp # [B, 6] -> [B, 4, 4]
+		self.transform = se3.transform # [B, 1, 4, 4] x [B, N, 3] -> [B, N, 3]
+
+		w1, w2, w3, v1, v2, v3 = delta, delta, delta, delta, delta, delta
+		twist = torch.Tensor([w1, w2, w3, v1, v2, v3])
+		self.dt = torch.nn.Parameter(twist.view(1, 6), requires_grad=learn_delta)
+
+		# results
+		self.last_err = None
+		self.g_series = None # for debug purpose
+		self.prev_r = None
+		self.g = None # estimation result
+		self.itr = 0
+		self.xtol = xtol
+		self.p0_zero_mean = p0_zero_mean
+		self.p1_zero_mean = p1_zero_mean
+
+	def forward(self, template, source, maxiter=10):
+		template, source, template_mean, source_mean = data_utils.mean_shift(template, source, 
+																			 self.p0_zero_mean, self.p1_zero_mean)
+
+		result = self.iclk(template, source, maxiter)
+		result = data_utils.postprocess_data(result, template, source, template_mean, source_mean, 
+											 self.p0_zero_mean, self.p1_zero_mean)
+		return result
+
+	def iclk(self, template, source, maxiter):
+		batch_size = template.size(0)
+
+		est_T0 = torch.eye(4).to(template).view(1, 4, 4).expand(template.size(0), 4, 4).contiguous()
+		est_T = est_T0
+		self.est_T_series = torch.zeros(maxiter+1, *est_T0.size(), dtype=est_T0.dtype)
+		self.est_T_series[0] = est_T0.clone()
+
+		training = self.handle_batchNorm(template, source)
+
+		# re-calc. with current modules
+		template_features = self.pooling(self.feature_model(template)) # [B, N, 3] -> [B, K]
+
+		# approx. J by finite difference
+		dt = self.dt.to(template).expand(batch_size, 6)
+		J = self.approx_Jic(template, template_features, dt)
+
+		self.last_err = None
+		pinv = self.compute_inverse_jacobian(J, template_features, source)
+		if pinv == {}:
+			result = {'est_R': est_T[:,0:3,0:3],
+					  'est_t': est_T[:,0:3,3],
+					  'est_T': est_T,
+					  'r': None,
+					  'transformed_source': self.transform(est_T.unsqueeze(1), source),
+					  'itr': 1,
+					  'est_T_series': self.est_T_series}
+			return result
+
+		itr = 0
+		r = None
+		for itr in range(maxiter):
+			self.prev_r = r
+			transformed_source = self.transform(est_T.unsqueeze(1), source) # [B, 1, 4, 4] x [B, N, 3] -> [B, N, 3]
+			source_features = self.pooling(self.feature_model(transformed_source)) # [B, N, 3] -> [B, K]
+			r = source_features - template_features
+
+			pose = -pinv.bmm(r.unsqueeze(-1)).view(batch_size, 6)
+
+			check = pose.norm(p=2, dim=1, keepdim=True).max()
+			if float(check) < self.xtol:
+				if itr == 0:
+					self.last_err = 0 # no update.
+				break
+
+			est_T = self.update(est_T, pose)
+			self.est_T_series[itr+1] = est_T.clone()
+
+		rep = len(range(itr, maxiter))
+		self.est_T_series[(itr+1):] = est_T.clone().unsqueeze(0).repeat(rep, 1, 1, 1)
+
+		self.feature_model.train(training)
+		self.est_T = est_T
+
+		result = {'est_R': est_T[:,0:3,0:3],
+				  'est_t': est_T[:,0:3,3],
+				  'est_T': est_T,
+				  'r': r,
+				  'transformed_source': self.transform(est_T.unsqueeze(1), source),
+				  'itr': itr+1,
+				  'est_T_series': self.est_T_series}
+		
+		return result
+
+	def update(self, g, dx):
+		# [B, 4, 4] x [B, 6] -> [B, 4, 4]
+		dg = self.exp(dx)
+		return dg.matmul(g)
+
+	def approx_Jic(self, template, template_features, dt):
+		# p0: [B, N, 3], Variable
+		# f0: [B, K], corresponding feature vector
+		# dt: [B, 6], Variable
+		# Jk = (feature_model(p(-delta[k], p0)) - f0) / delta[k]
+
+		batch_size = template.size(0)
+		num_points = template.size(1)
+
+		# compute transforms
+		transf = torch.zeros(batch_size, 6, 4, 4).to(template)
+		for b in range(template.size(0)):
+			d = torch.diag(dt[b, :]) # [6, 6]
+			D = self.exp(-d) # [6, 4, 4]
+			transf[b, :, :, :] = D[:, :, :]
+		transf = transf.unsqueeze(2).contiguous()  #   [B, 6, 1, 4, 4]
+		p = self.transform(transf, template.unsqueeze(1)) # x [B, 1, N, 3] -> [B, 6, N, 3]
+
+		#f0 = self.feature_model(p0).unsqueeze(-1) # [B, K, 1]
+		template_features = template_features.unsqueeze(-1) # [B, K, 1]
+		f = self.pooling(self.feature_model(p.view(-1, num_points, 3))).view(batch_size, 6, -1).transpose(1, 2) # [B, K, 6]
+
+		df = template_features - f # [B, K, 6]
+		J = df / dt.unsqueeze(1)
+
+		return J
+
+	def compute_inverse_jacobian(self, J, template_features, source):
+		# compute pinv(J) to solve J*x = -r
+		try:
+			Jt = J.transpose(1, 2) # [B, 6, K]
+			H = Jt.bmm(J) # [B, 6, 6]
+			B = self.inverse(H)
+			pinv = B.bmm(Jt) # [B, 6, K]
+			return pinv
+		except RuntimeError as err:
+			# singular...?
+			self.last_err = err
+			g = torch.eye(4).to(source).view(1, 4, 4).expand(source.size(0), 4, 4).contiguous()
+			#print(err)
+			# Perhaps we can use MP-inverse, but,...
+			# probably, self.dt is way too small...
+			source_features = self.pooling(self.feature_model(source)) # [B, N, 3] -> [B, K]
+			r = source_features - template_features
+			self.feature_model.train(self.feature_model.training)
+			return {}
+
+	def handle_batchNorm(self, template, source):
+		training = self.feature_model.training
+		if training:
+			# first, update BatchNorm modules
+			template_features, source_features = self.pooling(self.feature_model(template)), self.pooling(self.feature_model(source))
+		self.feature_model.eval()	# and fix them.
+		return training
+
+
+if __name__ == '__main__':
+	template, source = torch.rand(10,1024,3), torch.rand(10,1024,3)
+	pn = PointNet()
+
+	net = PointNetLK(pn)
+	result = net(template, source)
+	import ipdb; ipdb.set_trace()

learning3d/models/pooling.py

@@ -0,0 +1,15 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class Pooling(torch.nn.Module):
+	def __init__(self, pool_type='max'):
+		self.pool_type = pool_type
+		super(Pooling, self).__init__()
+
+	def forward(self, input):
+		if self.pool_type == 'max':
+			return torch.max(input, 2)[0].contiguous()
+		elif self.pool_type == 'avg' or self.pool_type == 'average':
+			return torch.mean(input, 2).contiguous()

learning3d/models/ppfnet.py

@@ -0,0 +1,102 @@
+"""Feature Extraction and Parameter Prediction networks
+"""
+import logging
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .. utils import sample_and_group_multi
+
+_raw_features_sizes = {'xyz': 3, 'dxyz': 3, 'ppf': 4}
+_raw_features_order = {'xyz': 0, 'dxyz': 1, 'ppf': 2}
+
+
+def get_prepool(in_dim, out_dim):
+	"""Shared FC part in PointNet before max pooling"""
+	net = nn.Sequential(
+		nn.Conv2d(in_dim, out_dim // 2, 1),
+		nn.GroupNorm(8, out_dim // 2),
+		nn.ReLU(),
+		nn.Conv2d(out_dim // 2, out_dim // 2, 1),
+		nn.GroupNorm(8, out_dim // 2),
+		nn.ReLU(),
+		nn.Conv2d(out_dim // 2, out_dim, 1),
+		nn.GroupNorm(8, out_dim),
+		nn.ReLU(),
+	)
+	return net
+
+
+def get_postpool(in_dim, out_dim):
+	"""Linear layers in PointNet after max pooling
+
+	Args:
+		in_dim: Number of input channels
+		out_dim: Number of output channels. Typically smaller than in_dim
+
+	"""
+	net = nn.Sequential(
+		nn.Conv1d(in_dim, in_dim, 1),
+		nn.GroupNorm(8, in_dim),
+		nn.ReLU(),
+		nn.Conv1d(in_dim, out_dim, 1),
+		nn.GroupNorm(8, out_dim),
+		nn.ReLU(),
+		nn.Conv1d(out_dim, out_dim, 1),
+	)
+
+	return net
+
+
+class PPFNet(nn.Module):
+	"""Feature extraction Module that extracts hybrid features"""
+	def __init__(self, features=['ppf', 'dxyz', 'xyz'], emb_dims=96, radius=0.3, num_neighbors=64):
+		super().__init__()
+
+		self._logger = logging.getLogger(self.__class__.__name__)
+		self._logger.info('Using early fusion, feature dim = {}'.format(emb_dims))
+		self.radius = radius
+		self.n_sample = num_neighbors
+
+		self.features = sorted(features, key=lambda f: _raw_features_order[f])
+		self._logger.info('Feature extraction using features {}'.format(', '.join(self.features)))
+
+		# Layers
+		raw_dim = np.sum([_raw_features_sizes[f] for f in self.features])  # number of channels after concat
+		self.prepool = get_prepool(raw_dim, emb_dims * 2)
+		self.postpool = get_postpool(emb_dims * 2, emb_dims)
+
+	def forward(self, xyz, normals):
+		"""Forward pass of the feature extraction network
+
+		Args:
+			xyz: (B, N, 3)
+			normals: (B, N, 3)
+
+		Returns:
+			cluster features (B, N, C)
+
+		"""
+		features = sample_and_group_multi(-1, self.radius, self.n_sample, xyz, normals)
+		features['xyz'] = features['xyz'][:, :, None, :]
+
+		# Gate and concat
+		concat = []
+		for i in range(len(self.features)):
+			f = self.features[i]
+			expanded = (features[f]).expand(-1, -1, self.n_sample, -1)
+			concat.append(expanded)
+		fused_input_feat = torch.cat(concat, -1)
+
+		# Prepool_FC, pool, postpool-FC
+		new_feat = fused_input_feat.permute(0, 3, 2, 1)  # [B, 10, n_sample, N]
+		new_feat = self.prepool(new_feat)
+
+		pooled_feat = torch.max(new_feat, 2)[0]  # Max pooling (B, C, N)
+
+		post_feat = self.postpool(pooled_feat)  # Post pooling dense layers
+		cluster_feat = post_feat.permute(0, 2, 1)
+		cluster_feat = cluster_feat / torch.norm(cluster_feat, dim=-1, keepdim=True)
+
+		return cluster_feat  # (B, N, C)