327 lines
12 KiB
Python
327 lines
12 KiB
Python
import numpy as np
|
|
import torch
|
|
|
|
|
|
def rotation_6d_to_matrix_numpy(d6):
|
|
a1, a2 = d6[:3], d6[3:]
|
|
b1 = a1 / np.linalg.norm(a1)
|
|
b2 = a2 - np.dot(b1, a2) * b1
|
|
b2 = b2 / np.linalg.norm(b2)
|
|
b3 = np.cross(b1, b2)
|
|
return np.stack((b1, b2, b3), axis=-2)
|
|
|
|
class Pointcloud:
|
|
def __init__(self, points, camera=None, name="unknown", parent=None, from_parent_index=None):
|
|
self.points = points # Nx3 numpy array
|
|
if camera is not None and camera.shape == (9,):
|
|
rotation_matrix = rotation_6d_to_matrix_numpy(camera)
|
|
translation = camera[6:]
|
|
camera = np.eye(4)
|
|
camera[:3, :3] = rotation_matrix
|
|
camera[:3, 3] = translation
|
|
self.camera = camera # 4x4 numpy array
|
|
self.current_transform = np.eye(4)
|
|
self.transform_history = {}
|
|
self.original_points = points.copy()
|
|
self.name = name
|
|
self.parent = parent
|
|
self.from_parent_index = from_parent_index
|
|
|
|
# --------- basic ---------
|
|
def __getitem__(self, item):
|
|
return self.points[item]
|
|
|
|
def __len__(self):
|
|
return len(self.points)
|
|
|
|
def __repr__(self):
|
|
return f"Pointcloud({self.points})"
|
|
|
|
def __str__(self):
|
|
return f"Pointcloud with {len(self.points)} points. \n\tCenter: {np.mean(self.points, axis=0)} \n\tTransform history: {list(self.transform_history.keys())}"
|
|
|
|
def __eq__(self, other):
|
|
return np.array_equal(self.points, other.points)
|
|
|
|
def __ne__(self, other):
|
|
return not np.array_equal(self.points, other.points)
|
|
|
|
def __contains__(self, item):
|
|
return item in self.points
|
|
|
|
def __iter__(self):
|
|
return iter
|
|
|
|
def concat(self, other):
|
|
return Pointcloud(np.concatenate([self.points, other.points], axis=0))
|
|
|
|
# --------- downsample ---------
|
|
def voxel_downsample(self, voxel_size):
|
|
voxel_indices = np.floor(self.points / voxel_size).astype(np.int32)
|
|
_, inverse, counts = np.unique(voxel_indices, axis=0, return_inverse=True, return_counts=True)
|
|
idx_sort = np.argsort(inverse)
|
|
idx_unique = idx_sort[np.cumsum(counts)-counts]
|
|
downsampled_points = self.points[idx_unique]
|
|
return Pointcloud(
|
|
downsampled_points,
|
|
name=self.name+'(voxel downsampled with voxel_size='+str(voxel_size)+')',
|
|
parent=self.parent,
|
|
from_parent_index=idx_unique
|
|
)
|
|
|
|
def random_downsample(self, num_points):
|
|
if self.points.shape[0] == 0:
|
|
downsampled_points = self.points
|
|
idx = np.array([])
|
|
else:
|
|
idx = np.random.choice(len(self.points), num_points, replace=True)
|
|
downsampled_points = self.points[idx]
|
|
return Pointcloud(
|
|
downsampled_points,
|
|
name=self.name+'(random downsampled with num_points='+str(num_points)+')',
|
|
parent=self.parent,
|
|
from_parent_index=idx
|
|
)
|
|
|
|
def fps_downsample(self, num_points):
|
|
N = self.points.shape[0]
|
|
mask = np.zeros(N, dtype=bool)
|
|
|
|
sampled_indices = np.zeros(num_points, dtype=int)
|
|
sampled_indices[0] = np.random.randint(0, N)
|
|
distances = np.linalg.norm(self.points - self.points[sampled_indices[0]], axis=1)
|
|
for i in range(1, num_points):
|
|
farthest_index = np.argmax(distances)
|
|
sampled_indices[i] = farthest_index
|
|
mask[farthest_index] = True
|
|
|
|
new_distances = np.linalg.norm(self.points - self.points[farthest_index], axis=1)
|
|
distances = np.minimum(distances, new_distances)
|
|
|
|
sampled_points = self.points[sampled_indices]
|
|
return Pointcloud(
|
|
sampled_points,
|
|
name=self.name+'(fps downsampled with num_points='+str(num_points)+')',
|
|
parent=self.parent,
|
|
from_parent_index=sampled_indices
|
|
)
|
|
# --------- transform ---------
|
|
def transform(self, transform_matrix, name=None):
|
|
self.points = np.dot(self.points, transform_matrix[:3, :3].T) + transform_matrix[:3, 3]
|
|
self.current_transform = np.dot(self.current_transform, transform_matrix)
|
|
if name is None:
|
|
name = f"transform_{len(self.transform_history)}"
|
|
self.transform_history[name] = transform_matrix
|
|
return self
|
|
|
|
def translate(self, translation, name=None):
|
|
transform_matrix = np.eye(4)
|
|
transform_matrix[:3, 3] = translation
|
|
self.transform(transform_matrix, name)
|
|
return self
|
|
|
|
def rotate(self, rotation, name=None):
|
|
transform_matrix = np.eye(4)
|
|
if rotation.shape == (3, 3):
|
|
rotation_matrix = rotation
|
|
elif rotation.shape == (6,):
|
|
rotation_matrix = rotation_6d_to_matrix_numpy(rotation)
|
|
transform_matrix[:3, :3] = rotation_matrix
|
|
self.transform(transform_matrix, name)
|
|
return self
|
|
|
|
def scale(self, scale_factor, name=None):
|
|
transform_matrix = np.eye(4)
|
|
transform_matrix[:3, :3] = scale_factor * np.eye(3)
|
|
self.transform(transform_matrix, name)
|
|
return self
|
|
|
|
def same_transform(self, other):
|
|
return np.allclose(self.current_transform, other.current_transform)
|
|
|
|
def print_transform_history(self):
|
|
print(f"Transform history of {self.name}:")
|
|
for name, transform_matrix in self.transform_history.items():
|
|
print(f"\t-{name}:")
|
|
for i in range(4):
|
|
print(f"\t\t{transform_matrix[i]}")
|
|
|
|
# --------- tensor ---------
|
|
def get_batchlized_points(self, batch_size=1):
|
|
return torch.tensor(self.points).unsqueeze(0).repeat(batch_size, 1, 1)
|
|
|
|
# --------- visualize ---------
|
|
def visualize(self, point_size=1, color=None):
|
|
import plotly.graph_objects as go
|
|
fig = go.Figure()
|
|
if color is None:
|
|
if self.name is not None:
|
|
hash_value = hash(self.name)
|
|
r = (hash_value & 0xFF) / 255.0
|
|
g = ((hash_value >> 8) & 0xFF) / 255.0
|
|
b = ((hash_value >> 16) & 0xFF) / 255.0
|
|
color = f'rgb({int(r*255)},{int(g*255)},{int(b*255)})'
|
|
else:
|
|
color = "gray"
|
|
|
|
|
|
if self.points is not None:
|
|
fig.add_trace(go.Scatter3d(
|
|
x=self.points[:, 0], y=self.points[:, 1], z=self.points[:, 2],
|
|
mode='markers', marker=dict(size=point_size, color=color, opacity=0.5), name=self.name
|
|
))
|
|
|
|
if self.camera is not None:
|
|
origin = self.camera[:3, 3]
|
|
z_axis = self.camera[:3, 2]
|
|
fig.add_trace(go.Cone(
|
|
x=[origin[0]], y=[origin[1]], z=[origin[2]],
|
|
u=[z_axis[0]], v=[z_axis[1]], w=[z_axis[2]],
|
|
colorscale="blues",
|
|
sizemode="absolute", sizeref=0.05, anchor="tail", showscale=False
|
|
))
|
|
|
|
title = self.name
|
|
fig.update_layout(
|
|
title=title,
|
|
scene=dict(
|
|
xaxis_title='X',
|
|
yaxis_title='Y',
|
|
zaxis_title='Z'
|
|
),
|
|
margin=dict(l=0, r=0, b=0, t=40),
|
|
scene_camera=dict(eye=dict(x=1.25, y=1.25, z=1.25))
|
|
)
|
|
fig.show()
|
|
|
|
# --------- save and load ---------
|
|
def save(self, file_path):
|
|
np.save(file_path, self.points)
|
|
|
|
def savetxt(self, file_path):
|
|
np.savetxt(file_path, self.points)
|
|
|
|
def load(self, file_path):
|
|
self.points = np.load(file_path)
|
|
|
|
def loadtxt(self, file_path):
|
|
self.points = np.loadtxt(file_path)
|
|
|
|
|
|
class PointcloudGroup:
|
|
def __init__(self, pointclouds: list[Pointcloud] = [], name="unknown"):
|
|
self.pointclouds = pointclouds
|
|
self.name = name
|
|
|
|
# --------- basic ---------
|
|
def __getitem__(self, item):
|
|
return self.pointclouds[item]
|
|
|
|
def __len__(self):
|
|
return len(self.pointclouds)
|
|
|
|
def __repr__(self):
|
|
return f"PointcloudGroup({self.pointclouds})"
|
|
|
|
def __str__(self):
|
|
return f"PointcloudGroup with {len(self.pointclouds)} pointclouds."
|
|
|
|
def __eq__(self, other):
|
|
return np.array_equal(self.pointclouds, other.pointclouds)
|
|
|
|
def __ne__(self, other):
|
|
return not np.array_equal(self.pointclouds, other.pointclouds)
|
|
|
|
def __contains__(self, item):
|
|
return item in self.pointclouds
|
|
|
|
def __iter__(self):
|
|
return iter
|
|
|
|
def __add__(self, pointcloud: Pointcloud):
|
|
new_group = PointcloudGroup(self.name)
|
|
new_group.pointclouds = self.pointclouds.copy()
|
|
new_group.pointclouds.append(pointcloud)
|
|
return new_group
|
|
|
|
def __iadd__(self, pointcloud: Pointcloud):
|
|
self.pointclouds.append(pointcloud)
|
|
return self
|
|
|
|
def add(self, pointcloud: Pointcloud):
|
|
self.pointclouds.append(pointcloud)
|
|
|
|
def concat(self, other):
|
|
new_group = PointcloudGroup(self.name)
|
|
new_group.pointclouds = self.pointclouds.copy()
|
|
new_group.pointclouds.extend(other.pointclouds)
|
|
return new_group
|
|
|
|
# --------- merge ---------
|
|
def merge_pointclouds(self, name="unknown"):
|
|
points = np.concatenate([pointcloud.points for pointcloud in self.pointclouds], axis=0)
|
|
return Pointcloud(points, name=name)
|
|
|
|
# --------- transform ---------
|
|
def transform(self, transform_matrix, name="unknown_transform"):
|
|
for pointcloud in self.pointclouds:
|
|
pointcloud.transform(transform_matrix, name)
|
|
|
|
def translate(self, translation, name="unknown_translate"):
|
|
transform_matrix = np.eye(4)
|
|
transform_matrix[:3, 3] = translation
|
|
self.transform(transform_matrix, name)
|
|
|
|
def rotate(self, rotation_matrix, name="unknown_ratate"):
|
|
transform_matrix = np.eye(4)
|
|
transform_matrix[:3, :3] = rotation_matrix
|
|
self.transform(transform_matrix, name)
|
|
|
|
def scale(self, scale_factor, name="unknown_scale"):
|
|
transform_matrix = np.eye(4)
|
|
transform_matrix[:3, :3] = scale_factor * np.eye(3)
|
|
self.transform(transform_matrix, name)
|
|
|
|
# --------- visualize ---------
|
|
def visualize(self, point_size=1, color=None):
|
|
import plotly.graph_objects as go
|
|
fig = go.Figure()
|
|
for pointcloud in self.pointclouds:
|
|
if color is None:
|
|
if pointcloud.name is not None:
|
|
hash_value = hash(pointcloud.name)
|
|
r = (hash_value & 0xFF) / 255.0
|
|
g = ((hash_value >> 8) & 0xFF) / 255.0
|
|
b = ((hash_value >> 16) & 0xFF) / 255.0
|
|
color = f'rgb({int(r*255)},{int(g*255)},{int(b*255)})'
|
|
else:
|
|
color = "gray"
|
|
|
|
if pointcloud.points is not None:
|
|
fig.add_trace(go.Scatter3d(
|
|
x=pointcloud.points[:, 0], y=pointcloud.points[:, 1], z=pointcloud.points[:, 2],
|
|
mode='markers', marker=dict(size=point_size, color=color, opacity=0.5), name=pointcloud.name
|
|
))
|
|
|
|
if pointcloud.camera is not None:
|
|
origin = pointcloud.camera[:3, 3]
|
|
z_axis = pointcloud.camera[:3, 2]
|
|
fig.add_trace(go.Cone(
|
|
x=[origin[0]], y=[origin[1]], z=[origin[2]],
|
|
u=[z_axis[0]], v=[z_axis[1]], w=[z_axis[2]],
|
|
colorscale="blues",
|
|
sizemode="absolute", sizeref=0.05, anchor="tail", showscale=False
|
|
))
|
|
|
|
title = self.name
|
|
fig.update_layout(
|
|
title=title,
|
|
scene=dict(
|
|
xaxis_title='X',
|
|
yaxis_title='Y',
|
|
zaxis_title='Z'
|
|
),
|
|
margin=dict(l=0, r=0, b=0, t=40),
|
|
scene_camera=dict(eye=dict(x=1.25, y=1.25, z=1.25))
|
|
)
|
|
fig.show() |