class BoundingBox():
"""
Simple axis-aligned bounding box class
"""
def __init__(self, corners, boundary="[][][]"):
"""
Initializes bounding box
Corners is list of length 2 with numpy arrays of shape (3,)
Boundary is denoting closed/open intervals in each axis although it doesn't end up being relevant
"""
self.corners = corners
self.boundary = boundary
def width(self):
"""
Returns real number, the length of the longest side
"""
return np.amax(self.corners[1]-self.corners[0])
def center(self):
"""
Returns numpy array of shape (3,)
"""
return 0.5*(self.corners[0]+self.corners[1])
def meshgrid(self, x, y, z):
"""
x, y, z are integers
Returns coordinates forming regular grid on bounding box where x, y, z are the number of cells in each axis
Coordinates xyz is a numpy array of shape (x, y, z, 3)
"""
xi, yi, zi = np.meshgrid(np.linspace(self.corners[0][0], self.corners[1][0], x), np.linspace(self.corners[0][1], self.corners[1][1], y), np.linspace(self.corners[0][1], self.corners[1][1], z))
xyz = np.concatenate([np.expand_dims(xi, -1), np.expand_dims(yi, -1), np.expand_dims(zi, -1)], axis=-1)
return xyz
def halve(self, axis):
"""
Subdivides bounding box into two along axis
"""
left_corner = np.array(self.corners[1])
left_corner[axis] = (self.corners[0][axis] + self.corners[1][axis]) / 2
left_boundary = list(self.boundary)
left_boundary[2*axis+1] = ")"
left_boundary = "".join(left_boundary)
left_bbox = BoundingBox([self.corners[0], left_corner], boundary=left_boundary)
right_corner = np.array(self.corners[0])
right_corner[axis] = (self.corners[0][axis] + self.corners[1][axis]) / 2
right_boundary = list(self.boundary)
right_boundary[2*axis+0] = "["
right_boundary = "".join(right_boundary)
right_bbox = BoundingBox([right_corner, self.corners[1]], boundary=right_boundary)
return left_bbox, right_bbox
def get_filter(self, pcd):
"""
Returns numpy boolean array representing mask on points
"""
points = np.array(pcd.points)
lower_mask = np.ones(points.shape[0], dtype=bool)
for i in range(3):
if self.boundary[2*i+0] == "(":
lower_mask = lower_mask & (points[:,i] > self.corners[0][i])
else:
lower_mask = lower_mask & (points[:,i] >= self.corners[0][i])
upper_mask = np.ones(points.shape[0], dtype=bool)
for i in range(3):
if self.boundary[2*i+1] == ")":
upper_mask = upper_mask & (points[:,i] < self.corners[1][i])
else:
upper_mask = upper_mask & (points[:,i] <= self.corners[1][i])
return lower_mask & upper_mask
def filter_point_cloud(self, pcd):
"""
Returns a new point cloud formed by only taking points in the bounding box
"""
points = np.array(pcd.points)
normals = np.array(pcd.normals)
groups = np.array(pcd.groups)
ids = np.array(pcd.ids)
lower_mask = np.ones(points.shape[0], dtype=bool)
for i in range(3):
if self.boundary[2*i+0] == "(":
lower_mask = lower_mask & (points[:,i] > self.corners[0][i])
else:
lower_mask = lower_mask & (points[:,i] >= self.corners[0][i])
upper_mask = np.ones(points.shape[0], dtype=bool)
for i in range(3):
if self.boundary[2*i+1] == ")":
upper_mask = upper_mask & (points[:,i] < self.corners[1][i])
else:
upper_mask = upper_mask & (points[:,i] <= self.corners[1][i])
return PointCloud(points[lower_mask & upper_mask], normals[lower_mask & upper_mask], groups[lower_mask & upper_mask], ids[lower_mask & upper_mask])
class PointCloud():
"""
Simple point cloud class
"""
def __init__(self, points, normals, groups=None, ids=None):
"""
Initializes point cloud with oriented normals
Points and normals should be numpy arrays of shape (n, 3)
Normals doesn't actually have to be normals, it's fine to set it to np.zeros_like(points) if they're not available
Groups and ids are metadata for each point used in our Barnes-Hut implementation
"""
self.points = points
self.normals = normals
self.area = 1
self.areas = np.ones_like(points[:,0])
self.n = self.points.shape[0]
if groups is None:
self.groups = np.zeros((self.n,), dtype=np.int32)
elif isinstance(groups, int):
self.groups = groups + np.zeros((self.n,), dtype=np.int32)
else:
self.groups = groups
if ids is None:
self.ids = np.arange(0, self.n, dtype=np.int32)
else:
self.ids = ids
def __len__(self):
"""
Returns length
"""
return self.n
def compute_weights(self, area, weight_normals=True):
"""
Note that the area must be set explicitly by this function, otherwise it will be assumed to be 1 in the constructor
Computes weights used for discretizing the integration
Lazy implementation so I'm just multiplying the normals by the weights so that they can be used in the Barnes-Hut algorithm implicitly
"""
# Library implementation should be O(nlogn)
nbrs = NearestNeighbors(n_neighbors=10).fit(self.points)
self.area = area
distances, indices = nbrs.kneighbors(self.points)
self.weights = np.mean(distances, axis=-1)**2
self.weights *= (area / np.sum(self.weights))
if weight_normals:
self.normals = self.weights[:,None] * self.normals
def get_bounding_box(self):
"""
Returns bounding box slightly (1.2 x) larger than the smallest bounding box
"""
bbox = BoundingBox([np.min(self.points, axis=0), np.max(self.points, axis=0)], boundary="[][][]")
center = bbox.center()
width = bbox.width()
new_bbox = BoundingBox([center - 0.6 * width, center + 0.6 * width], boundary="[][][]")
return new_bbox
def merge(self, o):
"""
Returns a merged point cloud of self and o
"""
points = np.concatenate([self.points, o.points], axis=0)
normals = np.concatenate([self.normals, o.normals], axis=0)
groups = np.concatenate([self.groups, o.groups], axis=0)
ids = np.concatenate([self.ids, o.ids], axis=0)
return PointCloud(points, normals, groups, ids)
class Octree():
"""
Simple octree class
"""
def __init__(self, pcd, bbox=None, max_height=None, parent=None):
"""
Creates octree from a set of points
"""
self.max_height = max_height
self.parent = parent
self.children = [[[None, None], [None, None]], [[None, None], [None, None]]]
if bbox is None:
self.bbox = pcd.get_bounding_box()
else:
self.bbox = bbox
self.pcd = None
if len(pcd) > 1 and self.max_height != 0:
for i1 in [0, 1]:
for i2 in [0, 1]:
for i3 in [0, 1]:
child_bbox = self.bbox.halve(axis=0)[i1]
child_pcd = child_bbox.filter_point_cloud(pcd)
child_bbox = child_bbox.halve(axis=1)[i2]
child_pcd = child_bbox.filter_point_cloud(child_pcd)
child_bbox = child_bbox.halve(axis=2)[i3]
child_pcd = child_bbox.filter_point_cloud(child_pcd)
self.children[i1][i2][i3] = Octree(child_pcd, child_bbox, None if max_height is None else max_height-1)
else:
self.pcd = pcd
def get_centers(self):
"""
Returns centers of nodes in the octree
Used for choosing evaluation targets
"""
centers = np.array(self.bbox.center())[None,:]
for i1 in [0, 1]:
for i2 in [0, 1]:
for i3 in [0, 1]:
if self.children[i1][i2][i3] is not None:
centers = np.concatenate((centers, self.children[i1][i2][i3].get_centers()), axis=0)
return centers
def fill_grid(self, grid, indicator, targets, xyz, mask):
"""
Assuming indicator has been computed for sparse points (targets), fills denser grid (grid) using octree structure
grid is mutable numpy array of shape (m,) and used as the output
xyz is the coordinates underlying grid so it has shape (m, 3)
indicator has shape (n,)
targets has shape (n, 3)
mask is a numpy boolean array of shape (m,), just set to all True when calling
"""
target_pcd = PointCloud(targets, np.zeros_like(targets))
target_filt = self.bbox.get_filter(target_pcd)
if np.sum(target_filt) == 0:
return
val = np.mean(indicator[target_filt])
targets = targets[target_filt,:]
indicator = indicator[target_filt]
xyz_pcd = PointCloud(xyz[mask], np.zeros_like(xyz[mask]))
filt = self.bbox.get_filter(xyz_pcd)
ind = np.flatnonzero(mask)[filt]
grid[ind] = val
new_mask = np.zeros_like(mask)
new_mask[ind] = 1
# print(mask.shape, mask.sum(), new_mask.sum())
for i1 in [0, 1]:
for i2 in [0, 1]:
for i3 in [0, 1]:
if self.children[i1][i2][i3] is not None:
self.children[i1][i2][i3].fill_grid(grid, np.array(indicator), np.array(targets), xyz, np.array(new_mask))
