Add a set of utility functions

util.py

+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