camera automatically increases render quality when it is not being moved

1 files changed, 75 insertions, 39 deletions

diff --git a/camera.py b/camera.py
index 57d0a67..a786971 100644
--- a/camera.py
+++ b/camera.py
@@ -3,6 +3,7 @@ from itertools import product, count
 from threading import Thread, Lock
 import time
 import os
+import sys
 
 from transform import rotate
 
@@ -55,6 +56,7 @@ class Camera(Thread):
         pygame.init()
         self.screen = pygame.display.set_mode(size)
         pygame.display.set_caption('')
+        self.clock = pygame.time.Clock()
 
         with self.lock:
             self.context.push()
@@ -68,23 +70,37 @@ class Camera(Thread):
 
         lower_bound, upper_bound = self.geometry.mesh.get_bounds()
         self.scale = np.linalg.norm(upper_bound-lower_bound)
-        self.source_position = [0,0,upper_bound[2]+1.0]
 
         self.nblocks = 64
 
+        self.point = np.array([0, self.scale*1.75, (lower_bound[2]+upper_bound[2])/2])
+        self.axis1 = np.array([1,0,0], float)
+        self.axis2 = np.array([0,0,1], float)
+
+        origins, directions = get_rays(self.point)
+
         with self.lock:
             self.context.push()
-            print 'initializing random number states.'
-            rng_state_count = max(len(geometry.mesh.triangles), self.width*self.height)
-            self.rng_states_gpu = cuda.mem_alloc(rng_state_count*sizeof('curandStateXORWOW', '#include <curand_kernel.h>'))
-            self.init_rng_kernel(np.int32(rng_state_count), self.rng_states_gpu, np.int32(0), np.int32(0), block=(self.nblocks,1,1), grid=(rng_state_count//self.nblocks+1,1))
+            self.origins_gpu = gpuarray.to_gpu(origins.astype(np.float32).view(gpuarray.vec.float3))
+            self.directions_gpu = gpuarray.to_gpu(directions.astype(np.float32).view(gpuarray.vec.float3))
+            self.pixels_gpu = gpuarray.zeros(self.width*self.height, dtype=np.int32)
+            self.context.pop()
+
+        self.render = False
+
+    def build_rgb_lookup(self, source_position=(0,0,0)):
+        with self.lock:
+            self.context.push()
+            print '\ninitializing random number states.'
+            self.rng_states_gpu = cuda.mem_alloc(self.width*self.height*sizeof('curandStateXORWOW', '#include <curand_kernel.h>'))
+            self.init_rng_kernel(np.int32(self.width*self.height), self.rng_states_gpu, np.int32(0), np.int32(0), block=(self.nblocks,1,1), grid=(self.width*self.height//self.nblocks+1,1))
             self.context.synchronize()
             print 'done.'
 
             self.source_positions_gpu = gpuarray.empty(len(geometry.mesh.triangles), dtype=gpuarray.vec.float3)
-            self.source_positions_gpu.fill(gpuarray.vec.make_float3(*self.source_position))
+            self.source_positions_gpu.fill(gpuarray.vec.make_float3(*source_position))
 
-            source_directions = np.mean(self.geometry.mesh[:], axis=1) - self.source_position
+            source_directions = np.mean(self.geometry.mesh[:], axis=1) - source_position
             self.source_directions_gpu = gpuarray.to_gpu(source_directions.astype(np.float32).view(gpuarray.vec.float3))
 
             self.rgb_lookup1_gpu = gpuarray.zeros(self.source_positions_gpu.size, dtype=gpuarray.vec.float3)
@@ -94,7 +110,7 @@ class Camera(Thread):
             rgb_runs = 100
             
             print 'building rgb lookup.'
-            self.build_rgb_lookup_kernel(np.int32(self.source_positions_gpu.size), self.rng_states_gpu, self.source_positions_gpu, self.source_directions_gpu, self.rgb_lookup1_gpu, self.rgb_lookup2_gpu, np.int32(rgb_runs), np.int32(self.max_steps), block=(self.nblocks,1,1), grid=(self.source_positions_gpu.size//self.nblocks+1,1))
+            self.build_rgb_lookup_kernel(np.int32(self.source_positions_gpu.size), self.source_positions_gpu, self.source_directions_gpu, self.rgb_lookup1_gpu, self.rgb_lookup2_gpu, np.uint64(np.random.randint(np.iinfo(np.int64).max)), np.int32(rgb_runs), np.int32(self.max_steps), block=(self.nblocks,1,1), grid=(self.source_positions_gpu.size//self.nblocks+1,1))
             self.context.synchronize()
             print 'done.'
 
@@ -108,26 +124,41 @@ class Camera(Thread):
             rgb_lookup2[rgb_lookup2 > 1.0] = 1.0
             self.rgb_lookup2_gpu.set(rgb_lookup2.view(gpuarray.vec.float3))
 
-            self.nruns = 1
+            self.nimages = 0
+            self.rgb_image = np.zeros((self.pixels_gpu.size, 3), float)
+            self.r_gpu = gpuarray.zeros(self.pixels_gpu.size, dtype=np.float32)
+            self.g_gpu = gpuarray.zeros(self.pixels_gpu.size, dtype=np.float32)
+            self.b_gpu = gpuarray.zeros(self.pixels_gpu.size, dtype=np.float32)
 
             self.context.pop()
 
-        self.point = np.array([0, self.scale*1.75, (lower_bound[2]+upper_bound[2])/2])
-        self.axis1 = np.array([1,0,0], float)
-        self.axis2 = np.array([0,0,1], float)
+    def clear_image(self):
+        try:
+            self.rgb_image.fill(0.0)
+            self.nimages = 0
+        except AttributeError:
+            pass
 
-        origins, directions = get_rays(self.point)
+    def acquire_image(self):
+        if not hasattr(self, 'rng_states_gpu'):
+            self.build_rgb_lookup()
 
         with self.lock:
             self.context.push()
-            self.origins_gpu = gpuarray.to_gpu(origins.astype(np.float32).view(gpuarray.vec.float3))
-            self.directions_gpu = gpuarray.to_gpu(directions.astype(np.float32).view(gpuarray.vec.float3))
-            self.pixels_gpu = gpuarray.zeros(self.width*self.height, dtype=np.int32)
+            self.render_kernel(np.int32(self.pixels_gpu.size), self.rng_states_gpu, self.origins_gpu, self.directions_gpu, self.r_gpu, self.g_gpu, self.b_gpu, self.rgb_lookup1_gpu, self.rgb_lookup2_gpu, np.int32(self.max_steps), block=(self.nblocks,1,1), grid=(self.pixels_gpu.size//self.nblocks+1,1))
+            self.rgb_image[:,0] += self.r_gpu.get()
+            self.rgb_image[:,1] += self.g_gpu.get()
+            self.rgb_image[:,2] += self.b_gpu.get()
+            self.nimages += 1
             self.context.pop()
 
-        self.render = False
+    def render_image(self):
+        scaled_rgb_image = ((self.rgb_image/self.nimages)*255).astype(np.int32)
+
+        image = scaled_rgb_image[:,0] << 16 | scaled_rgb_image[:,1] << 8 | scaled_rgb_image[:,2]
 
-        self.render_run_count = 1
+        pygame.surfarray.blit_array(self.screen, image.reshape(self.size))
+        pygame.display.flip()
 
     def screenshot(self, dir=''):
         root, ext = 'screenshot', 'png'
@@ -161,17 +192,20 @@ class Camera(Thread):
             self.context.pop()
 
     def update(self):
+        if self.render:
+            self.acquire_image()
+            self.render_image()
+            return
+
         with self.lock:
             self.context.push()
             t0 = time.time()
-            if self.render:
-                self.render_kernel(np.int32(self.pixels_gpu.size), self.rng_states_gpu, self.origins_gpu, self.directions_gpu, self.rgb_lookup1_gpu, self.rgb_lookup2_gpu, np.int32(self.render_run_count), self.pixels_gpu, np.int32(self.max_steps), block=(self.nblocks,1,1), grid=(self.pixels_gpu.size//self.nblocks+1,1))
-            else:
-                self.ray_trace_kernel(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))
+            self.ray_trace_kernel(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))
             self.context.synchronize()
             elapsed = time.time() - t0
 
-            print 'elapsed %f sec.' % elapsed
+            print '\relapsed %f sec.' % elapsed,
+            sys.stdout.flush()
 
             pygame.surfarray.blit_array(self.screen, self.pixels_gpu.get().reshape(self.size))
             pygame.display.flip()
@@ -187,20 +221,24 @@ class Camera(Thread):
         current_layer = 0
 
         while not done:
+            self.clock.tick(20)
+
+            if self.render and not clicked and not pygame.event.peek(KEYDOWN):
+                self.acquire_image()
+                self.render_image()
+
             for event in pygame.event.get():
                 if event.type == MOUSEBUTTONDOWN:
                     if event.button == 4:
                         v = self.scale*np.cross(self.axis1,self.axis2)/10.0
-
                         self.translate(v)
-
+                        self.clear_image()
                         self.update()
 
                     if event.button == 5:
                         v = -self.scale*np.cross(self.axis1,self.axis2)/10.0
-
                         self.translate(v)
-
+                        self.clear_image()
                         self.update()
 
                     if event.button == 1:
@@ -224,47 +262,52 @@ class Camera(Thread):
 
                     if shift:
                         v = mouse_direction*self.scale*length/float(self.width)
-
                         self.translate(v)
-
+                        self.clear_image()
                         self.update()
                     else:
                         phi = np.float32(2*np.pi*length/float(self.width))
                         n = rotate(mouse_direction, np.pi/2, -np.cross(self.axis1,self.axis2))
 
                         self.rotate(phi, n)
-
+                        self.clear_image()
                         self.update()
 
                 if event.type == KEYDOWN:
                     if event.key == K_a:
                         v = self.scale*self.axis1/10.0
                         self.translate(v)
+                        self.clear_image()
                         self.update()
 
                     if event.key == K_d:
                         v = -self.scale*self.axis1/10.0
                         self.translate(v)
+                        self.clear_image()
                         self.update()
 
                     if event.key == K_w:
                         v = self.scale*np.cross(self.axis1,self.axis2)/10.0
                         self.translate(v)
+                        self.clear_image()
                         self.update()
 
                     if event.key == K_s:
                         v = -self.scale*np.cross(self.axis1,self.axis2)/10.0
                         self.translate(v)
+                        self.clear_image()
                         self.update()
 
                     if event.key == K_SPACE:
                         v = self.scale*self.axis2/10.0
                         self.translate(v)
+                        self.clear_image()
                         self.update()
 
                     if event.key == K_LCTRL:
                         v = -self.scale*self.axis2/10.0
                         self.translate(v)
+                        self.clear_image()
                         self.update()
 
                     if event.key == K_LSHIFT or event.key == K_RSHIFT:
@@ -307,15 +350,6 @@ class Camera(Thread):
                         self.render = not self.render
                         self.update()
 
-                    if event.key == K_EQUALS:
-                        self.render_run_count += 1
-                        self.update()
-
-                    if event.key == K_MINUS:
-                        if self.render_run_count > 1:
-                            self.render_run_count -= 1
-                        self.update()
-
                 if event.type == KEYUP:
                     if event.key == K_LSHIFT or event.key == K_RSHIFT:
                         shift = False
@@ -324,6 +358,7 @@ class Camera(Thread):
                     done = True
                     break
 
+        print
         pygame.display.quit()
 
 if __name__ == '__main__':
@@ -376,6 +411,7 @@ if __name__ == '__main__':
 
     cuda.init()
     context = make_default_context()
+    print 'device %s' % context.get_device().name()
 
     module = SourceModule(src.kernel, options=['-I' + src.dir], no_extern_c=True, cache_dir=False)
     geometry = build(obj, options.bits)