diff options
Diffstat (limited to 'camera.py')
| -rwxr-xr-x | camera.py | 186 |
1 files changed, 96 insertions, 90 deletions
@@ -9,7 +9,7 @@ from chroma.geometry import Mesh, Solid, Geometry import chroma.make as make import matplotlib.pyplot as plt -from chroma.gpu import GPU, CUDAFuncs, to_float3 +from chroma import gpu from chroma.tools import timeit from chroma.transform import rotate from chroma.optics import vacuum, lambertian_surface @@ -100,22 +100,6 @@ def encode_movie(dir): print 'movie saved to %s.' % filename -def get_rays(position, size = (800, 600), width = 0.035, focal_length=0.018): - height = width*(size[1]/float(size[0])) - - x = np.linspace(-width/2, width/2, size[0]) - z = np.linspace(-height/2, height/2, size[1]) - - grid = np.array(tuple(product(x,[0],z))) - - grid += (0,focal_length,0) - focal_point = np.zeros(3) - - grid += position - focal_point += position - - return grid, focal_point-grid - class Camera(Thread): def __init__(self, geometry, size=(800,600), device_id=None, enable3d=False, green_magenta=False): Thread.__init__(self) @@ -137,11 +121,14 @@ class Camera(Thread): self.spnav = False def init_gpu(self): - self.gpu = GPU(self.device_id) - self.gpu.load_geometry(geometry) + self.gpu_instance = gpu.GPU(self.device_id) + self.gpu_geometry = gpu.GPUGeometry(self.gpu_instance, geometry) + self.gpu_funcs = gpu.GPUFuncs(self.gpu_instance.module) self.width, self.height = size + self.npixels = self.width*self.height + pygame.init() self.screen = pygame.display.set_mode(size) pygame.display.set_caption('') @@ -175,24 +162,26 @@ class Camera(Thread): self.viewing_angle = 0.0 - origins1, directions1 = project.from_film(self.point1, size=self.size, width=self.film_width) - origins2, directions2 = project.from_film(self.point2, size=self.size, width=self.film_width) + pos1, dir1 = project.from_film(self.point1, size=self.size, width=self.film_width) + pos2, dir2 = project.from_film(self.point2, size=self.size, width=self.film_width) + + self.rays1 = gpu.GPURays(pos1, dir1) + self.rays2 = gpu.GPURays(pos2, dir2) - origins = np.concatenate([origins1,origins2]) - directions = np.concatenate([directions1,directions2]) + scope_pos, scope_dir = project.from_film(self.point, size=np.array(self.size)/4.0, width=self.film_width/4.0) - scope_origins, scope_directions = project.from_film(self.point, size=np.array(self.size)/4.0, width=self.film_width/4.0) + self.scope_rays = gpu.GPURays(scope_pos, scope_dir) - self.scope_origins_gpu = ga.to_gpu(to_float3(scope_origins)) - self.scope_directions_gpu = ga.to_gpu(to_float3(scope_directions)) + self.pixels1_gpu = ga.empty(self.width*self.height, dtype=np.int32) + self.pixels2_gpu = ga.empty(self.width*self.height, dtype=np.int32) - self.distances_gpu = ga.empty(self.scope_origins_gpu.size, dtype=np.float32) + self.distances_gpu = ga.empty(self.scope_rays.pos.size, dtype=np.float32) else: - origins, directions = project.from_film(self.point, size=self.size, width=self.film_width) + pos, dir = project.from_film(self.point, size=self.size, width=self.film_width) + + self.rays = gpu.GPURays(pos, dir) - self.origins_gpu = ga.to_gpu(to_float3(origins)) - self.directions_gpu = ga.to_gpu(to_float3(directions)) - self.pixels_gpu = ga.zeros(self.origins_gpu.size, dtype=np.int32) + self.pixels_gpu = ga.empty(self.npixels, dtype=np.int32) self.alpha = True self.movie = False @@ -202,12 +191,15 @@ class Camera(Thread): @timeit def initialize_render(self): - self.rng_states_gpu = cuda.mem_alloc(self.origins_gpu.size*sizeof('curandStateXORWOW', '#include <curand_kernel.h>')) - self.gpu.kernels.init_rng(np.int32(self.origins_gpu.size), self.rng_states_gpu, np.int32(0), np.int32(0), block=(self.nblocks,1,1), grid=(self.origins_gpu.size//self.nblocks+1,1)) + self.rng_states_gpu = gpu.get_rng_states(self.npixels) self.xyz_lookup1_gpu = ga.zeros(len(self.geometry.mesh.triangles), dtype=ga.vec.float3) self.xyz_lookup2_gpu = ga.zeros(len(self.geometry.mesh.triangles), dtype=ga.vec.float3) - self.image_gpu = ga.zeros(self.pixels_gpu.size, dtype=ga.vec.float3) - self.gpu.context.synchronize() + + if self.enable3d: + self.image1_gpu = ga.zeros(self.npixels, dtype=ga.vec.float3) + self.image2_gpu = ga.zeros(self.npixels, dtype=ga.vec.float3) + else: + self.image_gpu = ga.zeros(self.npixels, dtype=ga.vec.float3) self.source_position = self.point @@ -222,33 +214,42 @@ class Camera(Thread): self.nlookup_calls = 0 def update_xyz_lookup(self, source_position): - for i in range(self.xyz_lookup1_gpu.size//(self.origins_gpu.size)+1): - self.gpu.kernels.update_xyz_lookup(np.int32(self.origins_gpu.size), np.int32(self.xyz_lookup1_gpu.size), np.int32(i*self.origins_gpu.size), ga.vec.make_float3(*source_position), self.rng_states_gpu, np.float32(685.0), ga.vec.make_float3(1.0,0.0,0.0), self.xyz_lookup1_gpu, self.xyz_lookup2_gpu, np.int32(self.max_steps), block=(self.nblocks,1,1), grid=(self.origins_gpu.size//self.nblocks+1,1)) - - for i in range(self.xyz_lookup1_gpu.size//(self.origins_gpu.size)+1): - self.gpu.kernels.update_xyz_lookup(np.int32(self.origins_gpu.size), np.int32(self.xyz_lookup1_gpu.size), np.int32(i*self.origins_gpu.size), ga.vec.make_float3(*source_position), self.rng_states_gpu, np.float32(545.0), ga.vec.make_float3(0.0,1.0,0.0), self.xyz_lookup1_gpu, self.xyz_lookup2_gpu, np.int32(self.max_steps), block=(self.nblocks,1,1), grid=(self.origins_gpu.size//self.nblocks+1,1)) - - for i in range(self.xyz_lookup1_gpu.size//(self.origins_gpu.size)+1): - self.gpu.kernels.update_xyz_lookup(np.int32(self.origins_gpu.size), np.int32(self.xyz_lookup1_gpu.size), np.int32(i*self.origins_gpu.size), ga.vec.make_float3(*source_position), self.rng_states_gpu, np.float32(445.0), ga.vec.make_float3(0.0,0.0,1.0), self.xyz_lookup1_gpu, self.xyz_lookup2_gpu, np.int32(self.max_steps), block=(self.nblocks,1,1), grid=(self.origins_gpu.size//self.nblocks+1,1)) + for wavelength, rgb_tuple in \ + zip([685.0, 545.0, 445.0],[(1,0,0),(0,1,0),(0,0,1)]): + for i in range(self.xyz_lookup1_gpu.size//(self.npixels)+1): + self.gpu_funcs.update_xyz_lookup(np.int32(self.npixels), np.int32(self.xyz_lookup1_gpu.size), np.int32(i*self.npixels), ga.vec.make_float3(*source_position), self.rng_states_gpu, np.float32(wavelength), ga.vec.make_float3(*rgb_tuple), self.xyz_lookup1_gpu, self.xyz_lookup2_gpu, np.int32(self.max_steps), block=(self.nblocks,1,1), grid=(self.npixels//self.nblocks+1,1)) self.nlookup_calls += 1 def clear_image(self): - self.image_gpu.fill(ga.vec.make_float3(0.0,0.0,0.0)) + if self.enable3d: + self.image1_gpu.fill(ga.vec.make_float3(0.0,0.0,0.0)) + self.image2_gpu.fill(ga.vec.make_float3(0.0,0.0,0.0)) + else: + self.image_gpu.fill(ga.vec.make_float3(0.0,0.0,0.0)) self.nimages = 0 - def update_image(self): - self.gpu.kernels.update_xyz_image(np.int32(self.origins_gpu.size), self.rng_states_gpu, self.origins_gpu, self.directions_gpu, np.float32(685.0), ga.vec.make_float3(1.0,0.0,0.0), self.xyz_lookup1_gpu, self.xyz_lookup2_gpu, self.image_gpu, np.int32(self.nlookup_calls), np.int32(self.max_steps), block=(self.nblocks,1,1), grid=(self.origins_gpu.size//self.nblocks+1,1)) + def update_image_from_rays(self, image, rays): + for wavelength, rgb_tuple in \ + zip([685.0, 545.0, 445.0],[(1,0,0),(0,1,0),(0,0,1)]): + self.gpu_funcs.update_xyz_image(np.int32(rays.pos.size), self.rng_states_gpu, rays.pos, rays.dir, np.float32(wavelength), ga.vec.make_float3(*rgb_tuple), self.xyz_lookup1_gpu, self.xyz_lookup2_gpu, image, np.int32(self.nlookup_calls), np.int32(self.max_steps), block=(self.nblocks,1,1), grid=(rays.pos.size//self.nblocks+1,1)) - self.gpu.kernels.update_xyz_image(np.int32(self.origins_gpu.size), self.rng_states_gpu, self.origins_gpu, self.directions_gpu, np.float32(545.0), ga.vec.make_float3(0.0,1.0,0.0), self.xyz_lookup1_gpu, self.xyz_lookup2_gpu, self.image_gpu, np.int32(self.nlookup_calls), np.int32(self.max_steps), block=(self.nblocks,1,1), grid=(self.origins_gpu.size//self.nblocks+1,1)) - - self.gpu.kernels.update_xyz_image(np.int32(self.origins_gpu.size), self.rng_states_gpu, self.origins_gpu, self.directions_gpu, np.float32(445.0), ga.vec.make_float3(0.0,0.0,1.0), self.xyz_lookup1_gpu, self.xyz_lookup2_gpu, self.image_gpu, np.int32(self.nlookup_calls), np.int32(self.max_steps), block=(self.nblocks,1,1), grid=(self.origins_gpu.size//self.nblocks+1,1)) + def update_image(self): + if self.enable3d: + self.update_image_from_rays(self.image1_gpu, self.rays1) + self.update_image_from_rays(self.image2_gpu, self.rays2) + else: + self.update_image_from_rays(self.image_gpu, self.rays) self.nimages += 1 def process_image(self): - self.gpu.kernels.process_image(np.int32(self.pixels_gpu.size), self.image_gpu, self.pixels_gpu, np.int32(self.nimages), block=(self.nblocks,1,1), grid=((self.pixels_gpu.size)//self.nblocks+1,1)) + if self.enable3d: + self.gpu_funcs.process_image(np.int32(self.pixels1_gpu.size), self.image1_gpu, self.pixels1_gpu, np.int32(self.nimages), block=(self.nblocks,1,1), grid=((self.pixels1_gpu.size)//self.nblocks+1,1)) + self.gpu_funcs.process_image(np.int32(self.pixels2_gpu.size), self.image2_gpu, self.pixels2_gpu, np.int32(self.nimages), block=(self.nblocks,1,1), grid=((self.pixels2_gpu.size)//self.nblocks+1,1)) + else: + self.gpu_funcs.process_image(np.int32(self.pixels_gpu.size), self.image_gpu, self.pixels_gpu, np.int32(self.nimages), block=(self.nblocks,1,1), grid=((self.pixels_gpu.size)//self.nblocks+1,1)) def screenshot(self, dir='', start=0): root, ext = 'screenshot', 'png' @@ -263,19 +264,19 @@ class Camera(Thread): print 'image saved to %s' % filename def rotate(self, phi, n): - self.gpu.kernels.rotate(np.int32(self.pixels_gpu.size), self.origins_gpu, np.float32(phi), ga.vec.make_float3(*n), block=(self.nblocks,1,1), grid=(self.pixels_gpu.size//self.nblocks+1,1)) - self.gpu.kernels.rotate(np.int32(self.pixels_gpu.size), self.directions_gpu, np.float32(phi), ga.vec.make_float3(*n), block=(self.nblocks,1,1), grid=(self.pixels_gpu.size//self.nblocks+1,1)) - - self.point = rotate(self.point, phi, n) - self.axis1 = rotate(self.axis1, phi, n) - self.axis2 = rotate(self.axis2, phi, n) - if self.enable3d: - self.gpu.kernels.rotate(np.int32(self.scope_origins_gpu.size), self.scope_origins_gpu, np.float32(phi), ga.vec.make_float3(*n), block=(self.nblocks,1,1), grid=(self.pixels_gpu.size//self.nblocks+1,1)) - self.gpu.kernels.rotate(np.int32(self.scope_directions_gpu.size), self.scope_directions_gpu, np.float32(phi), ga.vec.make_float3(*n), block=(self.nblocks,1,1), grid=(self.pixels_gpu.size//self.nblocks+1,1)) + self.rays1.rotate(phi, n) + self.rays2.rotate(phi, n) + self.scope_rays.rotate(phi, n) self.point1 = rotate(self.point1, phi, n) self.point2 = rotate(self.point2, phi, n) + else: + self.rays.rotate(phi, n) + + self.point = rotate(self.point, phi, n) + self.axis1 = rotate(self.axis1, phi, n) + self.axis2 = rotate(self.axis2, phi, n) if self.render: self.clear_image() @@ -283,15 +284,15 @@ class Camera(Thread): self.update() def rotate_around_point(self, phi, n, point, redraw=True): - self.gpu.kernels.rotate_around_point(np.int32(self.origins_gpu.size), self.origins_gpu, np.float32(phi), ga.vec.make_float3(*n), ga.vec.make_float3(*point), block=(self.nblocks,1,1), grid=(self.origins_gpu.size//self.nblocks+1,1)) - self.gpu.kernels.rotate(np.int32(self.directions_gpu.size), self.directions_gpu, np.float32(phi), ga.vec.make_float3(*n), block=(self.nblocks,1,1), grid=(self.directions_gpu.size//self.nblocks+1,1)) - self.axis1 = rotate(self.axis1, phi, n) self.axis2 = rotate(self.axis2, phi, n) if self.enable3d: - self.gpu.kernels.rotate_around_point(np.int32(self.scope_origins_gpu.size), self.scope_origins_gpu, np.float32(phi), ga.vec.make_float3(*n), ga.vec.make_float3(*point), block=(self.nblocks,1,1), grid=(self.origins_gpu.size//self.nblocks+1,1)) - self.gpu.kernels.rotate(np.int32(self.scope_directions_gpu.size), self.scope_directions_gpu, np.float32(phi), ga.vec.make_float3(*n), block=(self.nblocks,1,1), grid=(self.directions_gpu.size//self.nblocks+1,1)) + self.rays1.rotate_around_point(phi, n, point) + self.rays2.rotate_around_point(phi, n, point) + self.scope_rays.rotate_around_point(phi, n, point) + else: + self.rays.rotate_around_point(phi, n, point) if redraw: if self.render: @@ -300,15 +301,17 @@ class Camera(Thread): self.update() def translate(self, v, redraw=True): - self.gpu.kernels.translate(np.int32(self.pixels_gpu.size), self.origins_gpu, ga.vec.make_float3(*v), block=(self.nblocks,1,1), grid=(self.pixels_gpu.size//self.nblocks,1)) - self.point += v if self.enable3d: - self.gpu.kernels.translate(np.int32(self.scope_origins_gpu.size), self.scope_origins_gpu, ga.vec.make_float3(*v), block=(self.nblocks,1,1), grid=(self.pixels_gpu.size//self.nblocks,1)) + self.rays1.translate(v) + self.rays2.translate(v) + self.scope_rays.translate(v) self.point1 += v self.point2 += v + else: + self.rays.translate(v) if redraw: if self.render: @@ -324,13 +327,21 @@ class Camera(Thread): self.process_image() else: if self.alpha: - self.gpu.kernels.ray_trace_alpha(np.int32(self.pixels_gpu.size), self.origins_gpu, self.directions_gpu, self.pixels_gpu, block=(self.nblocks,1,1), grid=(self.pixels_gpu.size//self.nblocks+1,1)) + if self.enable3d: + self.gpu_funcs.ray_trace_alpha(np.int32(self.rays1.pos.size), self.rays1.pos, self.rays1.dir, self.pixels1_gpu, block=(self.nblocks,1,1), grid=(self.rays2.pos.size//self.nblocks+1,1)) + self.gpu_funcs.ray_trace_alpha(np.int32(self.rays2.pos.size), self.rays2.pos, self.rays2.dir, self.pixels2_gpu, block=(self.nblocks,1,1), grid=(self.rays2.pos.size//self.nblocks+1,1)) + else: + self.gpu_funcs.ray_trace_alpha(np.int32(self.rays.pos.size), self.rays.pos, self.rays.dir, self.pixels_gpu, block=(self.nblocks,1,1), grid=(self.rays.pos.size//self.nblocks+1,1)) else: - self.gpu.kernels.ray_trace(np.int32(self.pixels_gpu.size), self.origins_gpu, self.directions_gpu, self.pixels_gpu, block=(self.nblocks,1,1), grid=(self.pixels_gpu.size//self.nblocks+1,1)) + if self.enable3d: + self.gpu_funcs.ray_trace(np.int32(self.rays1.pos.size), self.rays1.pos, self.rays1.dir, self.pixels1_gpu, block=(self.nblocks,1,1), grid=(self.rays1.pos.size//self.nblocks+1,1)) + self.gpu_funcs.ray_trace(np.int32(self.rays2.pos.size), self.rays2.pos, self.rays2.dir, self.pixels2_gpu, block=(self.nblocks,1,1), grid=(self.rays2.pos.size//self.nblocks+1,1)) + else: + self.gpu_funcs.ray_trace(np.int32(self.rays.pos.size), self.rays.pos, self.rays.dir, self.pixels_gpu, block=(self.nblocks,1,1), grid=(self.rays.pos.size//self.nblocks+1,1)) def update(self): if self.enable3d: - self.gpu.kernels.distance_to_mesh(np.int32(self.scope_origins_gpu.size), self.scope_origins_gpu, self.scope_directions_gpu, self.distances_gpu, block=(self.nblocks,1,1), grid=(self.scope_origins_gpu.size//self.nblocks,1)) + self.gpu_funcs.distance_to_mesh(np.int32(self.scope_rays.pos.size), self.scope_rays.pos, self.scope_rays.dir, self.distances_gpu, block=(self.nblocks,1,1), grid=(self.scope_rays.pos.size//self.nblocks,1)) baseline = ga.min(self.distances_gpu).get().item() @@ -339,7 +350,7 @@ class Camera(Thread): v1 = d1/np.linalg.norm(d1) v1 *= baseline/60 - np.linalg.norm(d1) - self.gpu.kernels.translate(np.int32(self.pixels_gpu.size//2), self.origins_gpu[:self.pixels_gpu.size//2], ga.vec.make_float3(*v1), block=(self.nblocks,1,1), grid=((self.pixels_gpu.size//2)//self.nblocks+1,1)) + self.rays1.translate(v1) self.point1 += v1 @@ -347,7 +358,7 @@ class Camera(Thread): v2 = d2/np.linalg.norm(d2) v2 *= baseline/60 - np.linalg.norm(d2) - self.gpu.kernels.translate(np.int32(self.pixels_gpu.size//2), self.origins_gpu[self.pixels_gpu.size//2:], ga.vec.make_float3(*v2), block=(self.nblocks,1,1), grid=((self.pixels_gpu.size//2)//self.nblocks+1,1)) + self.rays2.translate(v2) self.point2 += v2 @@ -360,25 +371,23 @@ class Camera(Thread): phi = new_viewing_angle - self.viewing_angle - self.gpu.kernels.rotate_around_point(np.int32(self.pixels_gpu.size//2), self.origins_gpu[:self.pixels_gpu.size//2], np.float32(phi), ga.vec.make_float3(*self.axis1), ga.vec.make_float3(*self.point1), block=(self.nblocks,1,1), grid=((self.pixels_gpu.size//2)//self.nblocks+1,1)) - self.gpu.kernels.rotate_around_point(np.int32(self.pixels_gpu.size//2), self.origins_gpu[self.pixels_gpu.size//2:], np.float32(-phi), ga.vec.make_float3(*self.axis1), ga.vec.make_float3(*self.point2), block=(self.nblocks,1,1), grid=((self.pixels_gpu.size//2)//self.nblocks+1,1)) - self.gpu.kernels.rotate(np.int32(self.pixels_gpu.size//2), self.directions_gpu[:self.pixels_gpu.size//2], np.float32(phi), ga.vec.make_float3(*self.axis1), block=(self.nblocks,1,1), grid=((self.pixels_gpu.size//2)//self.nblocks+1,1)) - self.gpu.kernels.rotate(np.int32(self.pixels_gpu.size//2), self.directions_gpu[self.pixels_gpu.size//2:], np.float32(-phi), ga.vec.make_float3(*self.axis1), block=(self.nblocks,1,1), grid=((self.pixels_gpu.size//2)//self.nblocks+1,1)) + self.rays1.rotate_around_point(phi, self.axis1, self.point1) + self.rays2.rotate_around_point(-phi, self.axis1, self.point2) self.viewing_angle = new_viewing_angle self.update_pixels() - pixels = self.pixels_gpu.get() - if self.enable3d: - pixels1 = pixels[:len(pixels)//2] - pixels2 = pixels[len(pixels)//2:] + pixels1 = self.pixels1_gpu.get() + pixels2 = self.pixels2_gpu.get() if self.green_magenta: - pixels = (pixels1 & 0xff00) | (pixels2 & 0xff00ff) + pixels = (pixels1 & 0x00ff00) | (pixels2 & 0xff00ff) else: - pixels = (pixels1 & 0xff0000) | (pixels2 & 0xffff) + pixels = (pixels1 & 0xff0000) | (pixels2 & 0x00ffff) + else: + pixels = self.pixels_gpu.get() pygame.surfarray.blit_array(self.screen, pixels.reshape(self.size)) if self.doom_mode: @@ -392,11 +401,11 @@ class Camera(Thread): def loadlayer(self, layer): try: layergeometry = self.bvh_layers[layer] + layergeometry.activate() except KeyError: layergeometry = build(bvh_mesh(self.geometry, layer),8) - self.bvh_layers[layer] = layergeometry + self.bvh_layers[layer] = gpu.GPUGeometry(self.gpu, layergeometry) - self.gpu.load_geometry(layergeometry) self.update() def process_event(self, event): @@ -593,9 +602,6 @@ class Camera(Thread): self.update() - if self.enable3d: - self.update() - self.done = False self.clicked = False @@ -616,7 +622,7 @@ class Camera(Thread): if self.spnav: self.spnav_module.spnav_close() - del self.gpu + del self.gpu_instance class EventViewer(Camera): def __init__(self, geometry, filename, **kwargs): @@ -624,7 +630,7 @@ class EventViewer(Camera): self.rr = RootReader(filename) def color_hit_pmts(self): - self.gpu.reset_colors() + self.gpu_geometry.reset_colors() hit = self.ev.channels.hit t = self.ev.channels.t @@ -632,7 +638,7 @@ class EventViewer(Camera): # Important: Compute range only with HIT channels solid_colors = map_to_color(q, range=(q[hit].min(),q[hit].max())) - self.gpu.color_solids(hit, solid_colors) + self.gpu_geometry.color_solids(hit, solid_colors) self.update() def process_event(self, event): |