tutorial

prime.py

import pyoptixprime as prime
import tinyobjloader as tol
import Camera
import numpy as np
import cv2

def BufferDesc(ctx, bformat, arr):
    buffer = ctx.createBufferDesc(bformat, prime.RTP_BUFFER_TYPE_HOST, arr)
    # calculate the number of elements
    print(arr.shape)
    buffer.setRange(0, arr.size//arr.shape[-1])
    return buffer

# ('o': 3 x np.float32, 'd': 3 x np.float32)
def RaysDesc(ctx, rays):
    rays = BufferDesc(ctx, prime.RTP_BUFFER_FORMAT_RAY_ORIGIN_DIRECTION, rays)
    return rays

# ('t': np.float32)
def HitsDesc(ctx, rays, hformat=prime.RTP_BUFFER_FORMAT_HIT_T):
    hits = BufferDesc(ctx, hformat, rays)
    return hits

class Shape(object):
    def __init__(self, ctx, varr, iarr):
        # actually we can share vbufferdesc
        self._vbdesc = BufferDesc(ctx, prime.RTP_BUFFER_FORMAT_VERTEX_FLOAT3, varr)
        self._ibdesc = BufferDesc(ctx, prime.RTP_BUFFER_FORMAT_INDICES_INT3, iarr)

        self._mdl = ctx.createModel()
        self._mdl.setTriangles2(self._vbdesc, self._ibdesc)
        self._mdl.update(prime.RTP_MODEL_HINT_ASYNC)

    def createQuery(self, rays, hits, hit_type=prime.RTP_BUFFER_FORMAT_HIT_T):
        query = self._mdl.createQuery(prime.RTP_QUERY_TYPE_CLOSEST)
        query.setRays(rays)
        query.setHits(hits)
        return query

    def finish(self):
        self._mdl.finish()

"""
To build a shape:
1) create model from context. ctx.createModel
2) Set index and vertex buffer descriptors with mdl.setTriangles2
3) mdl.update()
4) make sure update is finish with mdl.finish()

To create a query:
1) mdl.createQuery(closest or any hit)
2) set rays descriptor with setRays
3) set hits descriptor with setHits

For more complex applications, you can change hit format.
"""

def main():
    ctx = prime.Context(prime.RTP_CONTEXT_TYPE_CUDA)
    scene = tol.LoadObj("monkey.obj")

    varr = np.asarray(scene['attribs']['vertices'], dtype=np.float32).reshape(-1, 3)

    shapes = []
    for name, shape in scene['shapes'].items():
        # an index is a tuple of three index (vertex_id, tex_id, normal_id)
        # we only need the first
        # copy() is necessary here, since the api requires continous memory
        iarr = np.asarray(shape['indices'], dtype=np.int32).reshape(-1, 3)[:,0].copy()
        # Three indices form a triangle
        new_one = Shape(ctx, varr, iarr.reshape(-1, 3))
        shapes.append(new_one)
    
    cam = Camera.Camera(extent=(800, 800))
    cam.lookAt((0, 0, 5), (0, 0, 0), (0, 1, 0))
    # (800, 800, 2, 3) -> (800, 800, 6)
    rays = cam.genRays()
    rays_desc = RaysDesc(ctx, rays.reshape(800, 800, 6))

    hits = np.zeros(shape=(800, 800, 1), dtype=np.float32)
    hits_desc = HitsDesc(ctx, hits)

    # now we create query from each shape
    querys = []
    for shape in shapes:
        querys.append(shape.createQuery(rays_desc, hits_desc))

    # Now make sure all shape is ready
    for shape in shapes:
        shape.finish()

    # query on the first shape
    querys[0].execute()

    frame = np.zeros_like(hits)
    frame[hits>0] = 1
    
    cv2.imshow('Result', frame)
    cv2.waitKey(0)

if __name__ == '__main__':
    main()