79 lines
2.9 KiB
Python
79 lines
2.9 KiB
Python
|
import numpy as np
|
||
|
|
import torch
|
||
|
|
from scipy.spatial.distance import cdist
|
||
|
|
|
||
|
|
|
||
|
|
class PclUtil:
|
||
|
|
CHAMFER = 1
|
||
|
|
|
||
|
|
@staticmethod
|
||
|
|
def transform(pts, pose=np.eye(4), scale=np.ones(3), inverse=False):
|
||
|
|
aug_scale = np.ones(4)
|
||
|
|
aug_scale[:3] = scale
|
||
|
|
aug_scale_mat = np.diag(aug_scale)
|
||
|
|
scale_pose = pose @ aug_scale_mat
|
||
|
|
aug_pts = np.hstack((pts, np.ones((pts.shape[0], 1))))
|
||
|
|
if inverse:
|
||
|
|
scale_pose = np.linalg.inv(scale_pose)
|
||
|
|
transformed_pts = scale_pose @ aug_pts.T
|
||
|
|
return transformed_pts.T[:, :3]
|
||
|
|
|
||
|
|
@staticmethod
|
||
|
|
def cam2canonical(cam_pts, cam2canonical_pose):
|
||
|
|
aug_pts = np.hstack((cam_pts, np.ones((cam_pts.shape[0], 1))))
|
||
|
|
transformed_pts = cam2canonical_pose @ aug_pts.T
|
||
|
|
return transformed_pts.T[:, :3]
|
||
|
|
|
||
|
|
@staticmethod
|
||
|
|
def transform_batch(pts, pose, scale, inverse=False):
|
||
|
|
batch_size = pts.shape[0]
|
||
|
|
aug_scale_mat = torch.eye(4).unsqueeze(0).repeat(batch_size, 1, 1)
|
||
|
|
for i in range(3):
|
||
|
|
aug_scale_mat[..., i, i] = scale[..., i]
|
||
|
|
scale_pose = pose @ aug_scale_mat
|
||
|
|
aug_pts = torch.cat((pts, torch.ones_like(pts[..., :1])), dim=-1)
|
||
|
|
if inverse:
|
||
|
|
scale_pose = torch.inverse(scale_pose)
|
||
|
|
transformers_pts = scale_pose @ aug_pts.transpose(1, 2)
|
||
|
|
return transformers_pts.transpose(1, 2)[..., :3]
|
||
|
|
|
||
|
|
@staticmethod
|
||
|
|
def transform_n_batch(pts, pose, scale=None, inverse=False):
|
||
|
|
transformed_pts_shape = (pts.shape[0], pose.shape[1], pts.shape[1], pts.shape[2])
|
||
|
|
transformed_pts = np.zeros(transformed_pts_shape)
|
||
|
|
batch_size = pose.shape[0]
|
||
|
|
n = pose.shape[1]
|
||
|
|
if scale is None:
|
||
|
|
scale = np.ones((batch_size, n, 3))
|
||
|
|
for batch_i in range(batch_size):
|
||
|
|
for i in range(n):
|
||
|
|
transformed_pts[batch_i, i, :, :] = PclUtil.transform(pts[batch_i], pose[batch_i, i],
|
||
|
|
scale[batch_i, i], inverse)
|
||
|
|
return transformed_pts
|
||
|
|
|
||
|
|
@staticmethod
|
||
|
|
def chamfer_distance(pts1, pts2):
|
||
|
|
dist_matrix1 = cdist(pts1, pts2, 'euclidean')
|
||
|
|
dist_matrix2 = cdist(pts2, pts1, 'euclidean')
|
||
|
|
chamfer_dist = np.mean(np.min(dist_matrix1, axis=1)) + np.mean(np.min(dist_matrix2, axis=1))
|
||
|
|
return chamfer_dist
|
||
|
|
|
||
|
|
@staticmethod
|
||
|
|
def distance(pts1, pts2, eval_type=1):
|
||
|
|
if eval_type == PclUtil.CHAMFER:
|
||
|
|
return PclUtil.chamfer_distance(pts1, pts2)
|
||
|
|
else:
|
||
|
|
raise ValueError('Unknown evaluation type:', eval_type)
|
||
|
|
|
||
|
|
@staticmethod
|
||
|
|
def sample_pcl(pcl, n_pts=1024):
|
||
|
|
indices = np.random.choice(pcl.shape[0], n_pts, replace=pcl.shape[0] < n_pts)
|
||
|
|
return pcl[indices, :]
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == '__main__':
|
||
|
|
batch_pts = np.random.random((2, 16, 3))
|
||
|
|
batch_n_pose = np.random.random((2, 3, 4, 4))
|
||
|
|
batch_n_scale = np.random.random((2, 3, 3))
|
||
|
|
poses = PclUtil.transform_n_batch(batch_pts, batch_n_pose, batch_n_scale)
|