diff --git a/Raytrace/BVHMesh.py b/Raytrace/BVHMesh.py new file mode 100644 index 0000000..6edb52e --- /dev/null +++ b/Raytrace/BVHMesh.py @@ -0,0 +1,178 @@ +from typing import Self, Literal +import numpy as np +from Raytrace.TriangleMesh import TriangleMesh, Ray + +# Adapted from https://jacco.ompf2.com/2022/04/13/how-to-build-a-bvh-part-1-basics/ +class BVHAABB: + min_pos:np.ndarray + max_pos:np.ndarray + def __init__(self): + self.min_pos = np.array([np.nan]*3) + self.max_pos = np.array([np.nan]*3) + @staticmethod + def from_array(triangles:np.ndarray) -> 'BVHAABB': + out = BVHAABB() + if triangles.size == 0: return out + triangles = triangles.reshape((-1, 3)) + out.min_pos = triangles.min(axis = 0) + out.max_pos = triangles.max(axis = 0) + return out + def raytrace(self, ray:Ray) -> float: + old_error_state = np.seterr(divide='ignore') + t1 = (self.min_pos - ray.origin) * ray.inverse_direction + t2 = (self.max_pos - ray.origin) * ray.inverse_direction + np.seterr(**old_error_state) + t = np.stack([t1, t2]) + tmin:float = t.min(axis=0).max() + tmax:float = t.max(axis=0).min() + return tmin if tmax >= tmin and tmax > 0 else np.inf + def grow(self, point:np.ndarray): + self.min_pos = np.nanmin(np.stack([point, self.min_pos]), axis=0) + self.max_pos = np.nanmax(np.stack([point, self.max_pos]), axis=0) + def grow_aabb(self, other:Self): + self.grow(other.min_pos) + self.grow(other.max_pos) + def area(self) -> float: + e = self.max_pos - self.min_pos + return e[0] * e[1] + e[1] * e[2] + e[2] * e[0] + def __str__(self): + return f'({float(self.min_pos[0])}, {float(self.min_pos[1])}, {float(self.min_pos[2])})#({float(self.max_pos[0])}, {float(self.max_pos[1])}, {float(self.max_pos[2])})' +class BVHNode: + aabb:BVHAABB + left:Self + right:Self + start_index:int + tri_count:int + def get_subarray(self, triangles:np.ndarray) -> np.ndarray: + return triangles[self.start_index:self.start_index + self.tri_count] + def update_bounds(self, triangles:np.ndarray): + self.aabb = BVHAABB.from_array(self.get_subarray(triangles)) + def is_leaf(self) -> bool: + return self.tri_count > 0 +class BVHMesh(TriangleMesh): + centroids:np.ndarray + root:BVHNode + node_count:int + def __init__(self, *args, min_node_size = 2, **kwargs) -> None: + super().__init__(*args, **kwargs) + self.build_BVH(min_node_size) + def build_BVH(self, min_node_size = 100) -> None: + # calculate triangle centroids for partitioning + self.centroids = self.triangles.sum(axis = 1) / 3 + # assign all triangles to root node + self.node_count = 1 + self.root = root = BVHNode() + root.start_index = 0 + root.tri_count = self.triangles.shape[0] + if root.tri_count == 0: return + root.update_bounds(self.triangles) + # subdivide recursively + self.subdivide(root, min_node_size) + def subdivide(self, node:BVHNode, min_node_size = 100) -> None: + # determine split axis using SAH + print(self.node_count, node.tri_count,end=' \r') + if node.tri_count < min_node_size: return + best_cost, axis, split_pos = self.find_best_split(node) + if best_cost >= node.tri_count * node.aabb.area(): return + # in-place partition + i = node.start_index + j = i + node.tri_count - 1 + while i <= j: + if self.centroids[i, axis] < split_pos: i += 1 + else: + self.triangles[[i,j]] = self.triangles[[j,i]] + self.centroids[[i,j]] = self.centroids[[j,i]] + j -= 1 + # abort split if one of the sides is empty + left_count = i - node.start_index + if left_count == 0 or left_count == node.tri_count: return + self.node_count += 2 + # create child nodes + node.left = left = BVHNode() + left.start_index = node.start_index + left.tri_count = left_count + + node.right = right = BVHNode() + right.start_index = i + right.tri_count = node.tri_count - left_count + + node.tri_count = 0 + left.update_bounds(self.triangles) + right.update_bounds(self.triangles) + # recurse + self.subdivide(left) + self.subdivide(right) + + def raytrace(self, ray:Ray) -> float: + return self.BVH_raytrace(ray) + + def BVH_raytrace(self, ray:Ray) -> float: + node:BVHNode = self.root + stack:list[BVHNode] = [] + best:float = np.inf + while 1: + if node.is_leaf(): + intersections = self.batch_triangle_ray_intersection(node.get_subarray(self.triangles), ray) + intersections[intersections < 0] = np.inf + best = min(best, np.min(intersections)) + if len(stack) == 0: break + node = stack.pop() + continue + child1, child2 = node.left, node.right + dist1 = node.left.aabb.raytrace(ray) + dist2 = node.right.aabb.raytrace(ray) + if dist1 > dist2: + dist1, dist2 = dist2, dist1 + child1, child2 = child2, child1 + if not np.isfinite(dist1): + if len(stack) == 0: break + node = stack.pop() + else: + node = child1 + if np.isfinite(dist2): + stack.append(child2) + return best + def find_best_split(self, node:BVHNode) -> tuple[float, int, float]: + BINS = 8 + axis:int = -1 + split_pos:float = 0 + best_cost:float = np.inf + triangles:np.ndarray[tuple[int, Literal[3], Literal[3]], np.dtype[np.float32]] = node.get_subarray(self.triangles) + centroids:np.ndarray[tuple[int, Literal[3]], np.dtype[np.float32]] = node.get_subarray(self.centroids) + for axis_i in range(3): + bounds_min = float(np.min(centroids[:,axis_i])) + bounds_max = float(np.max(centroids[:,axis_i])) + + if bounds_min == bounds_max: continue + # populate the bins + scale:float = BINS / (bounds_max - bounds_min) + bin_idx = ((centroids[:,axis_i] - bounds_min) * scale).astype(int) + bin_idx[bin_idx > BINS - 1] = BINS - 1 + bin_bounds = [BVHAABB.from_array(triangles[bin_idx == n]) for n in range(BINS)] + bin_counts = [np.count_nonzero(bin_idx == n) for n in range(BINS)] + # gather data for the 7 planes between the 8 bins + left_area = np.zeros(BINS - 1, float) + right_area = np.zeros(BINS - 1, float) + left_count = np.zeros(BINS - 1, int) + right_count = np.zeros(BINS - 1, int) + left_aabb = BVHAABB() + right_aabb = BVHAABB() + left_sum = right_sum = 0 + for i in range(BINS - 1): + left_sum += bin_counts[i] + left_count[i] = left_sum + left_aabb.grow_aabb(bin_bounds[i]) + left_area[i] = left_aabb.area() + right_sum += bin_counts[BINS - 1 - i] + right_count[BINS - 2 - i] = right_sum + right_aabb.grow_aabb(bin_bounds[BINS - 1 - i]) + right_area[BINS - 2 - i] = right_aabb.area() + # calculate SAH cost for the 7 planes + scale = 1. / scale + for i in range(BINS - 1): + cost:float = left_count[i] * left_area[i] + right_count[i] * right_area[i] + if cost < best_cost: + split_pos = bounds_min + scale * (i + 1) + axis = axis_i + best_cost = cost + return best_cost, axis, split_pos \ No newline at end of file