you can rotate just the camera by holding the control key and take movies by pressing the m key.

1 files changed, 82 insertions, 12 deletions

diff --git a/camera.py b/camera.py
index 148ddf9..a39060d 100644
--- a/camera.py
+++ b/camera.py
@@ -16,6 +16,10 @@ from pycuda import gpuarray
 from pycuda.characterize import sizeof
 import pycuda.driver as cuda
 
+from subprocess import call
+import shutil
+import tempfile
+
 def timeit(func):
     def f(*args, **kwargs):
         t0 = time.time()
@@ -24,6 +28,20 @@ def timeit(func):
         print '%s elapsed in %s().' % (datetime.timedelta(seconds=elapsed), func.__name__)
     return f
 
+def encode_movie(dir):
+    root, ext = 'movie', 'avi'
+    for i in count():
+        filename = '.'.join([root + str(i).zfill(5), ext])
+
+        if not os.path.exists(filename):
+            break
+
+    call(['mencoder', 'mf://' + dir + '/*.png', '-mf', 'fps=10', '-o', filename, '-ovc', 'xvid', '-xvidencopts', 'bitrate=3000'])
+
+    shutil.rmtree(dir)
+
+    print 'movie saved to %s.' % filename
+
 def get_rays(position, size = (800, 600), film_size = (0.035, 0.024), focal_length=0.05):
     """
     Generate ray positions and directions from a pinhole camera facing the negative y direction.
@@ -72,6 +90,7 @@ class Camera(Thread):
             self.context.push()
             self.ray_trace_kernel = self.module.get_function('ray_trace')
             self.rotate_kernel = self.module.get_function('rotate')
+            self.rotate_around_point_kernel = self.module.get_function('rotate_around_point')
             self.translate_kernel = self.module.get_function('translate')
             self.update_xyz_lookup_kernel = self.module.get_function('update_xyz_lookup')
             self.update_xyz_image_kernel = self.module.get_function('update_xyz_image')
@@ -97,6 +116,9 @@ class Camera(Thread):
             self.pixels_gpu = gpuarray.zeros(self.width*self.height, dtype=np.int32)
             self.context.pop()
 
+        self.movie = False
+        self.movie_index = 0
+        self.movie_dir = None
         self.render = False
 
     @timeit
@@ -167,11 +189,11 @@ class Camera(Thread):
             self.process_image_kernel(np.int32(self.width*self.height), self.image_gpu, self.pixels_gpu, np.int32(self.nimages), block=(self.nblocks,1,1), grid=((self.width*self.height)//self.nblocks+1,1))
             self.context.pop()
 
-    def screenshot(self, dir=''):
+    def screenshot(self, dir='', start=0):
         root, ext = 'screenshot', 'png'
 
-        for i in count():
-            filename = os.path.join(dir, '.'.join([root + str(i).zfill(4), ext]))
+        for i in count(start):
+            filename = os.path.join(dir, '.'.join([root + str(i).zfill(5), ext]))
 
             if not os.path.exists(filename):
                 break
@@ -182,8 +204,8 @@ class Camera(Thread):
     def rotate(self, phi, n):
         with self.lock:
             self.context.push()
-            self.rotate_kernel(np.int32(self.pixels_gpu.size), self.origins_gpu, phi, gpuarray.vec.make_float3(*n), block=(self.nblocks,1,1), grid=(self.pixels_gpu.size//self.nblocks+1,1))
-            self.rotate_kernel(np.int32(self.pixels_gpu.size), self.directions_gpu, phi, gpuarray.vec.make_float3(*n), block=(self.nblocks,1,1), grid=(self.pixels_gpu.size//self.nblocks+1,1))
+            self.rotate_kernel(np.int32(self.pixels_gpu.size), self.origins_gpu, np.float32(phi), gpuarray.vec.make_float3(*n), block=(self.nblocks,1,1), grid=(self.pixels_gpu.size//self.nblocks+1,1))
+            self.rotate_kernel(np.int32(self.pixels_gpu.size), self.directions_gpu, np.float32(phi), gpuarray.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)
@@ -195,6 +217,21 @@ class Camera(Thread):
 
         self.update()
 
+    def rotate_around_point(self, phi, n, point):
+        with self.lock:
+            self.context.push()
+            self.rotate_around_point_kernel(np.int32(self.origins_gpu.size), self.origins_gpu, np.float32(phi), gpuarray.vec.make_float3(*n), gpuarray.vec.make_float3(*point), block=(self.nblocks,1,1), grid=(self.origins_gpu.size//self.nblocks+1,1))
+            self.rotate_kernel(np.int32(self.directions_gpu.size), self.directions_gpu, np.float32(phi), gpuarray.vec.make_float3(*n), block=(self.nblocks,1,1), grid=(self.directions_gpu.size//self.nblocks+1,1))
+            self.context.pop()
+
+        self.axis1 = rotate(self.axis1, phi, n)
+        self.axis2 = rotate(self.axis2, phi, n)
+
+        if self.render:
+            self.clear_image()
+
+        self.update()
+
     def translate(self, v):
         with self.lock:
             self.context.push()
@@ -210,7 +247,7 @@ class Camera(Thread):
 
     def update(self):
         if self.render:
-            while self.nlookup_calls < 10:
+            while self.nlookup_calls < 100:
                 self.update_xyz_lookup(self.source_position)
             self.update_image()
             self.process_image()
@@ -226,12 +263,17 @@ class Camera(Thread):
             pygame.display.flip()
             self.context.pop()
 
+        if self.movie:
+            self.screenshot(self.movie_dir, self.movie_index)
+            self.movie_index += 1
+
     def run(self):
         self.update()
         
         done = False
         clicked = False
         shift = False
+        ctrl = False
 
         #current_layer = 0
 
@@ -277,7 +319,10 @@ class Camera(Thread):
                         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)
+                        if ctrl:
+                            self.rotate_around_point(phi, n, self.point)
+                        else:
+                            self.rotate(phi, n)
 
                 if event.type == KEYDOWN:
                     if event.key == K_a:
@@ -300,18 +345,26 @@ class Camera(Thread):
                         v = self.scale*self.axis2/10.0
                         self.translate(v)
 
-                    if event.key == K_LCTRL:
-                        v = -self.scale*self.axis2/10.0
-                        self.translate(v)
+                    # if event.key == K_LCTRL:
+                    #     v = -self.scale*self.axis2/10.0
+                    #     self.translate(v)
 
-                    if event.key == K_p:
+                    if event.key == K_F6:
                         self.clear_xyz_lookup()
-                        self.update_xyz_lookup(self.point)
+                        self.clear_image()
+                        self.source_position = self.point
+
+                    if event.key == K_p:
+                        for i in range(100):
+                            self.update_xyz_lookup(self.point)
                         self.source_position = self.point
 
                     if event.key == K_LSHIFT or event.key == K_RSHIFT:
                         shift = True
 
+                    if event.key == K_LCTRL or event.key == K_RCTRL:
+                        ctrl = True
+
                     if event.key == K_ESCAPE:
                         done = True
                         break
@@ -350,16 +403,33 @@ class Camera(Thread):
                             self.initialize_render()
 
                         self.render = not self.render
+                        self.clear_image()
                         self.update()
 
+                    if event.key == K_m:
+                        if self.movie:
+                            encode_movie(self.movie_dir)
+                            self.movie_dir = None
+                            self.movie = False
+                        else:
+                            self.movie_index = 0
+                            self.movie_dir = tempfile.mkdtemp()
+                            self.movie = True
+
                 if event.type == KEYUP:
                     if event.key == K_LSHIFT or event.key == K_RSHIFT:
                         shift = False
 
+                    if event.key == K_LCTRL or event.key == K_RCTRL:
+                        ctrl = False
+
                 if event.type == pygame.QUIT:
                     done = True
                     break
 
+        if self.movie:
+            encode_movie(self.movie_dir)
+
         pygame.display.quit()
 
 if __name__ == '__main__':