Merge pull request #1 from tjs20011/raytrace-update

Raytrace update

ExampleMultipleSonarReadings.ipynb

@@ -0,0 +1,103 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from Raytrace.TriangleMesh import Ray\n",
+    "from Raytrace.BVHMesh import BVHMesh\n",
+    "from Raytrace.SideScan import SideScan, ScanReading\n",
+    "import numpy as np\n",
+    "import time"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "stl_path = 'sample_Hull_Mesh.stl'\n",
+    "start_time = time.time()\n",
+    "mesh = BVHMesh(stl_path = stl_path, min_node_size = 100)\n",
+    "end_time = time.time()\n",
+    "\n",
+    "print('Build Time:', end_time - start_time, 'seconds')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "facing = np.array([1, 0, 0])\n",
+    "min_angle = 0\n",
+    "max_angle = -np.pi/2\n",
+    "sample_ray_count = 1000\n",
+    "\n",
+    "origin_start = np.array([5, -1, 1])\n",
+    "origin_end = np.array([9, -1, 1])\n",
+    "readings_count = 100\n",
+    "# liniar interpolation between the two for this simple demo\n",
+    "origins = [origin_start * (1-i) + origin_end * i for i in np.arange(0, 1+1/(readings_count-1)/2, 1/(readings_count-1))]\n",
+    "\n",
+    "orientations = [Ray(facing, origin) for origin in origins]\n",
+    "\n",
+    "rays_list = [SideScan.generate_rays(orientation, min_angle, max_angle, sample_ray_count) for orientation in orientations]\n",
+    "\n",
+    "readings:list[ScanReading] = []\n",
+    "for n,rays in enumerate(rays_list):\n",
+    "    print(n + 1, len(rays_list), sep='/', end='\\r')\n",
+    "    readings.append(SideScan(mesh).scan_rays(rays))\n",
+    "\n",
+    "print()\n",
+    "\n",
+    "# 100 readings at 1000 rays each takes just over 2 minutes on my machine\n",
+    "print('Done')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import plotly.graph_objs as go\n",
+    "from Raytrace.PlotRays import plot_readings_heatmap"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "plot_readings_heatmap(readings)"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.11.5"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

ExampleSingleSonarReading.ipynb

@@ -0,0 +1,114 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from Raytrace.TriangleMesh import Ray\n",
+    "from Raytrace.BVHMesh import BVHMesh\n",
+    "from Raytrace.SideScan import SideScan\n",
+    "import numpy as np\n",
+    "import time"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "stl_path = 'sample_Hull_Mesh.stl'\n",
+    "start_time = time.time()\n",
+    "mesh = BVHMesh(stl_path = stl_path, min_node_size = 100)\n",
+    "end_time = time.time()\n",
+    "\n",
+    "print('Build Time:', end_time - start_time, 'seconds')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "origin = np.array([5, -1, 1])\n",
+    "facing = np.array([1, 0, 0])\n",
+    "min_angle = 0\n",
+    "max_angle = -np.pi/2\n",
+    "sample_ray_count = 1000\n",
+    "\n",
+    "orientation = Ray(facing, origin)\n",
+    "\n",
+    "rays = SideScan.generate_rays(orientation, min_angle, max_angle, sample_ray_count)\n",
+    "\n",
+    "reading = SideScan(mesh).scan_rays(rays)\n",
+    "\n",
+    "print('Triangles:', mesh.triangles.shape[0])\n",
+    "reading.print_summary()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import plotly.graph_objs as go\n",
+    "from PlotRays import plot_mesh, plot_rays, plot_reading"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "mesh_plot = plot_mesh(mesh, color='lightpink', opacity=1)\n",
+    "rays_plot = plot_rays(reading, line={'color':'Blue'})"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "fig = go.Figure(data=[mesh_plot, rays_plot])\n",
+    "fig.update_layout(autosize=False, width=1000, height=500, margin=dict(l=20, r=20, t=20, b=20), scene = {'aspectmode':'data'})\n",
+    "fig.show()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "plot_reading(reading)"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.11.5"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

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