inference finish
This commit is contained in:
hofee
327
beans/pointcloud.py
Normal file
327
beans/pointcloud.py
Normal file
@@ -0,0 +1,327 @@
|
||||
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()
|
||||
Reference in New Issue
Block a user