Python ray-tracer for a triangle soup

souptracer.py

import numpy as np
import numba as nb

"""
A small ray-tracer for rendering a triangle soup. It gives all intersections
along the ray (in no particular order) and does not do back-face culling.
Uses a BVH for algorithmic acceleration and numba for low-level optimization.

Usage:
```
import souptracer

# Create an array of triangle vertices e.g. read from file
# tri_verts.shape = (n_tris, 3, 3)

scene = souptracer.Scene()
scene.add_triangles(tri_verts)
scene.build_bvh()

for x in range(...):
    for y in range(...):
        ray_o, ray_d = get_ray(x, y)
        hits = scene.intersect(ray_o, ray_d)
        # Use hits to render e.g. binary mask, depth, thickness
```
"""

@nb.jit(nopython=True, parallel=True, cache=True, nogil=True)
def _triangle_intersect(ray_o, ray_d, a, b, c, eps=0.00001):
    # Triangle-ray intersection test (Moeller-Trumbore): see https://tinyurl.com/db5e3khr
    v0v1 = b - a
    v0v2 = c - a
    pvec = np.cross(ray_d, v0v2)
    det = np.dot(v0v1, pvec)
    # No culling! Our soup has two sided triangles!
    if np.abs(det) < eps:
        return None
    invdet = 1. / det
    tvec = ray_o - a
    # Solve for first barycentric coord
    u = np.dot(tvec, pvec) * invdet
    if u < 0. or u > 1.:
        return None
    qvec = np.cross(tvec, v0v1)
    # Solve for second barycentric coord
    v = np.dot(ray_d, qvec) * invdet
    if v < 0. or (u + v) > 1.:
        return None
    t = np.dot(v0v2, qvec) * invdet
    if t < 0.:
        return None
    return t

@nb.jit(nopython=True, parallel=True, cache=True, nogil=True)
def _aabb_intersect(ray_o, ray_d, minp, maxp):
    # AABB-ray intersection test: see https://tinyurl.com/4es6u7hf
    invdir = 1. / ray_d
    t1 = (minp - ray_o) * invdir
    t2 = (maxp - ray_o) * invdir
    # Compute x positions on ray
    tmin = (t1[0] if invdir[0] >= 0. else t2[0])
    tmax = (t2[0] if invdir[0] >= 0. else t1[0])
    # Compute y positions on ray
    tymin = (t1[1] if invdir[1] >= 0. else t2[1])
    tymax = (t2[1] if invdir[1] >= 0. else t1[1])
    # Check x-y miss
    if tmin > tymax or tymin > tmax:
        return None
    if tymin > tmin: tmin = tymin
    if tymax < tmax: tmax = tymax
    # Compute z positions on ray
    tzmin = (t1[2] if invdir[2] >= 0. else t2[2])
    tzmax = (t2[2] if invdir[2] >= 0. else t1[2])
    # Check xy-z miss
    if tmin > tzmax or tzmin > tmax:
        return None
    if tzmin > tmin: tmin = tzmin
    return tmin

class Primitive:
    def intersect(self):
        return None

class Triangle(Primitive):
    def __init__(self, a, b, c):
        self.a = a; self.b = b; self.c = c
        self.centroid = (a + b + c) / 3.
        verts = np.vstack((a, b, c))
        minp = np.min(verts, axis=0)
        maxp = np.max(verts, axis=0)
        self.aabb = AABB(minp, maxp, False)

    def intersect(self, ray_o, ray_d):
        return _triangle_intersect(ray_o, ray_d, self.a, self.b, self.c)

class AABB(Primitive):
    # Default AABB is empty and needs extending
    # TODO better way to handle default empty constructor?
    def __init__(self, minp=None, maxp=None, empty=True):
        self.empty = empty
        if not empty:
            # Sanitize the input corners
            self.minp = np.minimum(minp, maxp)
            self.maxp = np.maximum(minp, maxp)

    def intersect(self, ray_o, ray_d):
        if self.empty:
            return None
        return _aabb_intersect(ray_o, ray_d, self.minp, self.maxp)

    def extend(self, aabb):
        if not aabb.empty:
            self.minp = (aabb.minp if self.empty else np.minimum(self.minp, aabb.minp))
            self.maxp = (aabb.maxp if self.empty else np.maximum(self.maxp, aabb.maxp))
            self.empty = False

class Node:
    def __init__(self):
        self.left = None; self.right = None
        self.primitives = []
        self.aabb = AABB()

class Scene:
    """The virtual scene holding all primitives and having some extent."""

    def __init__(self):
        """Creates an empty scene."""
        self.primitives = []
        self.extents = AABB()
        self.use_bvh = False

    def add_triangles(self, tri_verts):
        """Adds the provided list of triangle vertices to the scene."""
        for a,b,c in tri_verts:
            tri = Triangle(a, b, c)
            self.primitives.append(tri)
            self.extents.extend(tri.aabb)
        # New primitives require (re)building the BVH hierarchy
        self.use_bvh = False

    @staticmethod
    def _intersect_bvh_node_recursive(ray_o, ray_d, node):
        # Check if node is hit
        hit = node.aabb.intersect(ray_o, ray_d)
        if hit == None:
            return []
        if len(node.primitives) > 0:
            # Leaf: collect any primitive intersections
            hits = []
            for p in node.primitives:
                hit = p.intersect(ray_o, ray_d)
                if hit != None:
                    hits.append(hit)
            return hits
        else:
            # Interior: recurse and combine lower hits
            hits1 = Scene._intersect_bvh_node_recursive(ray_o, ray_d, node.left)
            hits2 = Scene._intersect_bvh_node_recursive(ray_o, ray_d, node.right)
            return hits1 + hits2

    @staticmethod
    def _add_bvh_node_recursive(primitives, leaf_size=6):
        node = Node()
        if len(primitives) < leaf_size:
            # Leaf node: we hold primitives
            for p in primitives:
                node.primitives.append(p)
                node.aabb.extend(p.aabb)
            return node
        else:
            # Interior node: we hold children
            for p in primitives:
                node.aabb.extend(p.aabb)
            # Simple naive median split
            extents = node.aabb.maxp - node.aabb.minp
            dim = np.argmax(extents)
            primitives.sort(key=lambda x: x.centroid[dim])
            node.left = Scene._add_bvh_node_recursive(
                primitives[:len(primitives)//2],
                leaf_size
            )
            node.right = Scene._add_bvh_node_recursive(
                primitives[len(primitives)//2:],
                leaf_size
            )
            return node

    def build_bvh(self):
        """Builds the BVH for the scene's primitives."""
        self.root = Scene._add_bvh_node_recursive(self.primitives)
        self.use_bvh = True

    def intersect(self, ray_o, ray_d):
        """Intersect the provided ray with the scene's primitives."""
        if self.use_bvh:
            return Scene._intersect_bvh_node_recursive(ray_o, ray_d, self.root)
        else:
            # Best acceleration without BVH is check against extents first
            hits = []
            if self.extents.intersect(ray_o, ray_d) != None:
                for i in range(len(self.primitives)):
                    hit = self.primitives[i].intersect(ray_o, ray_d)
                    if hit != None:
                        hits.append(hit)
            return hits