Robust rays w/ embree 3 in python

embree_rays.py

import embree
import numpy as np


def get_centroids(V, F):
    return V[F].mean(axis=1)


def get_cross_products(V, F):
    V0 = V[F][:, 0, :]
    C = np.cross(V[F][:, 1, :] - V0, V[F][:, 2, :] - V0)
    return C


def get_face_areas(V, F):
    C = get_cross_products(V, F)
    C_norms = np.sqrt(np.sum(C**2, axis=1))
    A = C_norms/2
    return A


def get_surface_normals(V, F):
    C = get_cross_products(V, F)
    C_norms = np.sqrt(np.sum(C**2, axis=1))
    N = C/C_norms.reshape(C.shape[0], 1)
    return N


def get_surface_normals_and_face_areas(V, F):
    C = get_cross_products(V, F)
    C_norms = np.sqrt(np.sum(C**2, axis=1))
    N = C/C_norms.reshape(C.shape[0], 1)
    A = C_norms/2
    return N, A


class TrimeshShapeModel:

    def __init__(self, V, F, P=None, N=None, A=None):
        self.dtype = V.dtype

        self.V = V
        self.F = F

        if N is None and A is None:
            N, A = get_surface_normals_and_face_areas(V, F)
        elif A is None:
            if N.shape[0] != F.shape[0]:
                raise Exception(
                    'must pass same number of surface normals as faces (got ' +
                    '%d faces and %d normals' % (F.shape[0], N.shape[0])
                )
            A = get_face_areas(V, F)
        elif N is None:
            N = get_surface_normals(V, F)

        self.P = get_centroids(V, F)
        self.N = N
        self.A = A

        assert self.P.dtype == self.dtype
        assert self.N.dtype == self.dtype
        assert self.A.dtype == self.dtype

        self._make_scene()

    def _make_scene(self):
        '''Set up an Embree scene. This function allocates some memory that
        Embree manages, and loads vertices and index lists for the
        faces. In Embree parlance, this function creates a "device",
        which manages a "scene", which has one "geometry" in it, which
        is our mesh.

        '''
        device = embree.Device()
        geometry = device.make_geometry(embree.GeometryType.Triangle)
        scene = device.make_scene()
        scene.set_flags(4)
        vertex_buffer = geometry.set_new_buffer(
            embree.BufferType.Vertex, # buf_type
            0, # slot
            embree.Format.Float3, # fmt
            3*np.dtype('float32').itemsize, # byte_stride
            self.V.shape[0], # item_count
        )
        vertex_buffer[:] = self.V[:]
        index_buffer = geometry.set_new_buffer(
            embree.BufferType.Index, # buf_type
            0, # slot
            embree.Format.Uint3, # fmt
            3*np.dtype('uint32').itemsize, # byte_stride,
            self.F.shape[0]
        )
        index_buffer[:] = self.F[:]
        geometry.commit()
        scene.attach_geometry(geometry)
        geometry.release()
        scene.commit()

        # This is the only variable we need to retain a reference to
        # (I think)
        self.scene = scene

    def __reduce__(self):
        return (self.__class__, (self.V, self.F, self.P, self.N, self.A))

    @property
    def num_faces(self):
        return self.P.shape[0]

    def check_vis_1_to_N(self, i, J, eps=None):
        if eps is None:
            eps = 1e3*np.finfo(np.float32).resolution

        D = self.P[J] - self.P[i]
        D /= np.sqrt(np.sum(D**2, axis=1)).reshape(D.shape[0], 1)
        P = self.P[i] + eps*D

        rayhit = embree.RayHit1M(len(J))
        context = embree.IntersectContext()
        rayhit.org[:] = P
        rayhit.dir[:] = D
        rayhit.tnear[:] = 0
        rayhit.tfar[:] = np.inf
        rayhit.flags[:] = 0
        rayhit.geom_id[:] = embree.INVALID_GEOMETRY_ID

        self.scene.intersect1M(context, rayhit)

        return np.logical_and(
            rayhit.geom_id != embree.INVALID_GEOMETRY_ID,
            rayhit.prim_id == J
        )

    def intersect(self, origins, dirs, tnear_override=None):
        m = origins.shape[0]
        if dirs.shape[0] != m:
            raise Exception('origins and dirs need the same number of rows')

        rayhit = embree.RayHit1M(m)
        context = embree.IntersectContext()
        rayhit.org[:] = origins
        rayhit.dir[:] = dirs
        if not tnear_override is None:
            rayhit.tnear[:]=tnear_override[:]
        else:
            rayhit.tnear[:] = 0
        rayhit.tfar[:] = np.inf
        rayhit.flags[:] = 0
        rayhit.geom_id[:] = embree.INVALID_GEOMETRY_ID

        self.scene.intersect1M(context, rayhit)

        I = rayhit.prim_id.copy().astype(np.intp)
        I[rayhit.geom_id == embree.INVALID_GEOMETRY_ID] = -1
        return I

    def get_direct_irradiance(self, F0, dir_sun, eps=None):
        if eps is None:
            eps = 1e3*np.finfo(np.float32).resolution

        # Here, we use Embree directly to find the indices of triangles
        # which are directly illuminated (I_sun) or not (I_shadow).
        ray = embree.Ray1M(self.num_faces)
        ray.org[:] = self.P + eps*self.N
        ray.dir[:] = dir_sun
        ray.tnear[:] = 0
        ray.tfar[:] = np.inf
        ray.flags[:] = 0
        context = embree.IntersectContext()
        self.scene.occluded1M(context, ray)

        # Determine which rays escaped (i.e., can see the sun)
        I = np.isposinf(ray.tfar)

        # Compute the direct irradiance
        E = np.zeros(self.num_faces, dtype=self.dtype)
        E[I] = F0*np.maximum(0, self.N[I]@dir_sun)

        return E