Commit 155a9f4b authored by Tim Shirley's avatar Tim Shirley

Add a set of utility functions

parent 3f8aa6ac
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import pyntcloud
from pyntcloud import PyntCloud
# the relative location of the data directory
data_dir = "data"
# rendering backend to use
render_backend = "pythreejs"
# known marker size
marker_size = 45 #mm
basis_vectors = [
{
"color": "red",
"vertices": [[0, 0, 0], [10, 0, 0]]
},
{
"color": "green",
"vertices": [[0, 0, 0], [0, 10, 0]]
},
{
"color": "blue",
"vertices": [[0, 0, 0], [0, 0, 10]]
}
]
def load_cloud(dataset, name):
"""load point cloud from storage"""
filename = f"{data_dir}/{dataset}/{name}.ply"
return PyntCloud.from_file(filename)
def store_cloud(cloud_in, dataset, name, mesh=False):
"""store a point cloud"""
filename = f"{data_dir}/{dataset}/{name}.ply"
if mesh:
cloud_in.to_file(filename=filename, as_text=True, mesh=cloud_in.mesh)
else:
cloud_in.to_file(filename=filename, as_text=True)
def store_value(dataset, name, value, overwrite=False):
f = open(f"{data_dir}/{dataset}/measurements.txt", "w+" if overwrite else "a")
f.write(f"{name}: {value}\n")
f.close()
def rename_columns(cloud_in):
"""rename the color columns of a COLMAP point cloud"""
cloud = PyntCloud(cloud_in.points.copy())
column_names = {
"diffuse_red": "red",
"diffuse_green": "green",
"diffuse_blue": "blue"
}
cloud.points = cloud.points.rename(index=int, columns=column_names)
return cloud
def skew_cross(v):
"""create the skew-symmetric cross-product matrix for v"""
# shorten and flatten v if neccessary
if len(v) == 4: v = v[:3]/v[3]
v = v.flatten()
# make a diagonal matrix with v
diag = np.diag(v)
# shift each row of the matrix by one
rolled = np.roll(diag, 1, 1)
# shift each column of the matrix by negative one
skv = np.roll(rolled, -1, 0)
# combine the matrix with its transposed version
return skv - skv.T
def vector_rotation(a, b):
"""calculate a rotation matrix for the rotation from vector a onto b"""
# normalize vectors first
a = a / np.linalg.norm(a)
b = b / np.linalg.norm(b)
# calculate the cross product
v = np.cross(a, b)
# normalize the rotation axis
u = v / np.linalg.norm(v)
# get sine and cosine of the rotation angle
sin_θ = np.linalg.norm(v)
cos_θ = np.dot(a, b)
# create an identity matrix as basis
I = np.identity(3)
# create the cross-product matrix
u_x = skew_cross(u)
# get the tensor product of v with itself
u_tensordot = np.tensordot(u, u, axes=0)
# place the result into a 4 x 4 matrix
R = np.identity(4)
# according to the equation for the rotational matrix
R[:3, :3] = cos_θ*I + sin_θ*u_x + (1-cos_θ)*u_tensordot
return R
def translation(v):
"""create a translation matrix for vector v"""
T = np.identity(4)
# select the last column of the identity matrix
T[:3, -1] = v
return T
def scaling(vec):
"""create a scaling matrix for vector v"""
T = np.identity(4)
# select the last column of the identity matrix
T[:3, :3] = np.diag(vec)
return T
def transform(cloud_in, T):
"""apply a transformation matrix to a point cloud"""
points = cloud_in.points.copy()
# select x, y, z columns as numpy array
data = points[["x", "y", "z"]].values
# create a column of ones
ones = np.ones(data.shape[0])[:,None]
# append column of ones to points, making them 4D vectors
data = np.hstack((data, ones))
# calculate dot product with transformation matrix
data = np.dot(T, data.T).T
# copy transformed points back into dataframe (as 3D vectors)
points[["x", "y", "z"]] = data[:, 0:3]
# reset dtype to float32 (as required by WebGL)
points[["x", "y", "z"]] = points[["x", "y", "z"]] .astype(np.float32)
# check for normals
if "nx" in points.columns:
normals = points[["nx", "ny", "nz"]].values
normals = np.hstack((normals, ones))
normals = np.dot(T, normals.T).T
points[["nx", "ny", "nz"]] = data[:, 0:3]
points[["nx", "ny", "nz"]] = points[["nx", "ny", "nz"]] .astype(np.float32)
return PyntCloud(points)
def fit_ransac_plane(cloud_in):
"""fit a RANSAC plane onto the cloud points"""
from pyntcloud.ransac.fitters import single_fit
from pyntcloud import ransac
# fit a RANSAC plane onto the point cloud and return inliers and the plane model
# note: we use the fitting function directly here, because by default pyntcloud
# doesn't return the model, only the inliers - but it's the model we want
inliers, model = single_fit(cloud_in.xyz,
ransac.RANSAC_MODELS["plane"],
ransac.RANSAC_SAMPLERS["random"],
model_kwargs={"max_dist": 1e-2},
max_iterations=100,
n_inliers_to_stop=10000,
return_model=True)
return inliers, model
def crop_roi(cloud_in, roi_center, roi_area, inliers=True):
"""crop everything outside a region of iterest"""
cloud = PyntCloud(cloud_in.points.copy())
# get the boundaries in x, y and z
min_x = roi_center[0] - roi_area[0] / 2
max_x = roi_center[0] + roi_area[0] / 2
min_y = roi_center[1] - roi_area[1] / 2
max_y = roi_center[1] + roi_area[1] / 2
min_z = roi_center[2] - roi_area[2] / 2
max_z = roi_center[2] + roi_area[2] / 2
# create and apply bounding box filter
bbox = cloud.get_filter("BBOX", min_x=min_x, max_x=max_x,
min_y=min_y, max_y=max_y,
min_z=min_z, max_z=max_z)
# optionally filter outliers
if not inliers:
bbox = ~bbox
cloud.apply_filter(bbox)
return cloud
def plot(cloud_in, mask=None, color=[64, 224, 208], initial_point_size=0.7, polylines=basis_vectors, background="white", backend=render_backend, **kwargs):
"""highlight a subset of points of a cloud"""
_points = cloud_in.points.copy()
if (mask is not None):
_points.loc[mask, ["red", "green", "blue"]] = color
PyntCloud(_points).plot(initial_point_size=initial_point_size, polylines=polylines, background=background, backend=backend, **kwargs)
\ No newline at end of file
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment