Port raytracer to TensorFlow.

diff --git a/tf.py b/tf.py
index 2bc8365..5a10bbb 100755 (executable)
--- a/tf.py
+++ b/tf.py
@@ -1,143 +1,217 @@
 #!/usr/bin/env python
 """
 #!/usr/bin/env python
 """

-The worst raytracer.
-Emil Mikulic <emikulic@gmail.com> was here 2012.
+Raytracer in tensorflow.
+Emil Mikulic <emikulic@gmail.com> was here 2016.

 """
 """

+print 'loading'
+import time
+t0 = time.time()

 import numpy as np
 import numpy as np

-from accidental_complexity import show
-# this file contains mostly essential complexity
-
-def vec(x, y, z):
-  return np.array((x, y, z), dtype=np.float)
-
-def color(r, g, b):
-  return np.reshape(vec(r,g,b), (1,1,3))
+import tensorflow as tf
+import Image
+td = time.time() - t0
+print 'loading took %.3f sec' % td
+
+# width and height
+W, H = 800, 600
+
+def save(tf_array, fn):
+  img = np.asarray(tf_array)
+  a = np.nanmin(img)
+  b = np.nanmax(img)
+  print 'range is', (a, b)
+  #img = (img - a) / (b - a)
+  img = np.round(img).clip(0, 255)
+  Image.fromarray(img.astype(np.uint8)).save(fn)
+
+def square(x):
+  with tf.name_scope('square'):
+    return x * x
+
+def dot_vec(a, b):
+  return tf.reduce_sum(a * b, 0)

 
 def dot(a, b):
 
 def dot(a, b):

-  return (a * b).sum(axis=2)
+  """Dot product of arrays of vectors."""
+  return tf.reduce_sum(a * b, 1)

 
 def mag(a):
 
 def mag(a):

-  return np.sqrt(np.square(a).sum(axis=2))
+  return tf.sqrt(dot(a, a))

 
 def depthwise(a):
 
 def depthwise(a):

-  return np.expand_dims(a, axis=2)
+  return tf.expand_dims(a, 1)

 
 

-def norm(a):
-  return a / depthwise(mag(a))
+def normalize(a):
+  with tf.name_scope('normalize'):
+    return a / depthwise(mag(a))

 
 def reflect(i, n):
   c = n * depthwise(dot(n, -i))
   return 2*c + i
 
 def diffuse(n, l):
 
 def reflect(i, n):
   c = n * depthwise(dot(n, -i))
   return 2*c + i
 
 def diffuse(n, l):

-  return depthwise(dot(n, l).clip(0, 1))
+  return tf.maximum(dot(n, l), 0.)

 
 def specular(r, l, power):
 
 def specular(r, l, power):

-  return pow(diffuse(r, l), power)
+  return tf.pow(diffuse(r, l), power)
+
+def clip(a, x, y):
+  return tf.minimum(tf.maximum(a, x), y)

 
 def lerp(ar, a, b, u, v):
   "[u,v] is mapped to [a,b]"
 
 def lerp(ar, a, b, u, v):
   "[u,v] is mapped to [a,b]"

-  return a + (b-a) * (ar.clip(u,v) - u) / (v-u)
-
-class Scene:
-  def __init__(self):
-    self.light1 = vec(20,10,-10)
-    self.light1_col = color(1,.5,.3)
-    self.light2 = vec(-5,5,5)
-    self.light2_col = color(1,1,1) - self.light1_col
-    self.sky_col = color(.2, .3, .6)
-    self.sph_center = vec(-.5,0,0)
-    self.sph_rad = .5
-    self.sph_col = color(1,1,1)
-    self.camera = vec(0,0,-5)
-    self.ground_col1 = color(.4,.4,.4)
-    self.ground_col2 = color(.8,.8,.8)
-    self.plane_n = vec(0,1,0)
-    self.plane_amb = 0.02
-    self.plane_dif = 0.9
-    self.plane_height = -0.5
-    self.fog_near = 0
-    self.fog_far = 40
-
-  def trace_bg(self, camera, ray):
-    # background
-    plane_depth = (self.plane_height - camera[:,:,1]) / ray[:,:,1]
-    plane_p = camera + depthwise(plane_depth) * ray
-    plane_l1 = norm(self.light1 - plane_p)
-    plane_l2 = norm(self.light2 - plane_p)
-    plane_xpos = plane_p[:,:,0] - np.floor(plane_p[:,:,0])
-    plane_zpos = plane_p[:,:,2] - np.floor(plane_p[:,:,2])
-    checkers = depthwise(np.logical_xor(plane_xpos < 0.5, plane_zpos < 0.5))
-
-    # trace shadows (light 1)
-    ec = plane_p - self.sph_center
-    b = 2 * dot(plane_l1, ec)
-    c = dot(ec, ec) - np.square(self.sph_rad)
-    det = b*b - 4*c
-    depth = (-b - np.sqrt(det)) / 2
-    illum_l1 = np.where(depthwise(depth > 0), 0, self.light1_col)
-
-    # light 2
-    b = 2 * dot(plane_l2, ec)
-    det = b*b - 4*c
-    depth = (-b - np.sqrt(det)) / 2
-    illum_l2 = np.where(depthwise(depth > 0), 0, self.light2_col)
-
-    # shade plane
-    plane_col = np.where(checkers, self.ground_col1, self.ground_col2) * (
-        self.plane_amb + self.plane_dif * (
-          diffuse(self.plane_n, plane_l1) * illum_l1 +
-          diffuse(self.plane_n, plane_l2) * illum_l2))
-
-    fog_dist = mag(plane_p - vec(0, self.plane_height, 0))
-    fog = lerp(fog_dist, 1, 0, self.fog_near, self.fog_far)
-    plane_col *= depthwise(np.square(fog))
-
-    sky = self.sky_col * depthwise(ray[:,:,1])
-    return np.where(depthwise(plane_depth > 0), plane_col, sky)
-
-  def render(self, w, h):
-    # set up scene
-    sph_amb = 0.02
-    sph_diff = 0.05
-    sph_spec = 1.1
-    sph_refl = 0.5
-    sph_shine = 20
-
-    # points on the screen plane (z = 0)
-    aspect = float(w) / float(h)
-    scr_x = np.linspace(-1 + 1./w, 1 - 1./w, w) * aspect
-    scr_y = np.linspace(-1 + 1./h, 1 - 1./h, h) * -1
-    scr = np.zeros((h,w,3), dtype=np.float)
-    scr[:,:,0] = np.reshape(scr_x, (1, w))
-    scr[:,:,1] = np.reshape(scr_y, (h, 1))
-
-    ray = norm(scr - self.camera)
-
-    # sphere
-    ec = self.camera - self.sph_center
-    b = 2 * dot(ray, ec)
-    c = ec.dot(ec) - np.square(self.sph_rad)
-    det = b*b - 4*c
-    sph_depth = (-b - np.sqrt(det)) / 2
-
-    # sphere color
-    sph_p = self.camera + depthwise(sph_depth) * ray
-    sph_l1 = norm(self.light1 - sph_p)
-    sph_l2 = norm(self.light2 - sph_p)
-    sph_n = (sph_p - self.sph_center) / self.sph_rad
-    sph_r = norm(reflect(ray, sph_n))
-    r_col = self.trace_bg(sph_p, sph_r)
-
-    sph = self.sph_col * (sph_amb +
-        sph_diff * diffuse(sph_n, sph_l1) * self.light1_col +
-        sph_diff * diffuse(sph_n, sph_l2) * self.light2_col +
-        sph_spec * specular(sph_r, sph_l1, sph_shine) * self.light1_col +
-        sph_spec * specular(sph_r, sph_l2, sph_shine) * self.light2_col +
+  return a + (b-a) * (clip(ar,u,v) - u) / (v-u)
+
+def img_of(v):
+  return tf.zeros([H*W,3]) + v
+
+def main():
+  print 'building'
+  with tf.Session() as sess:
+    light1 = tf.constant([20.,10.,-10.], name='light1')
+    light1_col = tf.constant([1., .5, .3], name='light1_col')
+    light2 = tf.constant([-5.,5.,5.], name='light2')
+    light2_col = tf.constant([.1, .4, .7], name='light2_col')
+    sky_col = tf.constant([.2,.3,.6], name='sky_col')
+    sph_center = tf.constant([-.5,0,0], name='sph_center')
+    sph_rad = tf.constant(.5, name='sph_rad')
+    sph_col = tf.constant([1.,1.,1.], name='sph_col')
+    ground_col1 = tf.constant([.4,.4,.4], name='ground_col1')
+    ground_col2 = tf.constant([.8,.8,.8], name='ground_col2')
+    plane_n = tf.constant([0.,1.,0.], name='plane_n')
+    plane_amb = tf.constant(0.02, name='plane_amb')
+    plane_dif = tf.constant(0.9, name='plane_dif')
+    plane_height = tf.constant(-.5, name='plane_height')
+    fog_near = tf.constant(0., name='fog_near')
+    fog_far = tf.constant(40., name='fog_far')
+
+    sph_amb = tf.constant(0.02, name='sph_amb')
+    sph_diff = tf.constant(0.05, name='sph_diff')
+    sph_spec = tf.constant(1.1, name='sph_spec')
+    sph_refl = tf.constant(.5, name='sph_refl')
+    sph_shine = tf.constant(20., name='sph_shine')
+
+    def trace_bg(camera, ray):
+      with tf.name_scope('intersect'):
+        camera_y = tf.slice(camera, [0,1], [-1,1])
+        ray_y = tf.slice(ray, [0,1], [-1,1])
+        plane_depth = (plane_height - camera_y) / ray_y
+        plane_p = camera + (plane_depth) * ray
+
+      with tf.name_scope('checkers'):
+        plane_x = tf.slice(plane_p, [0,0], [-1,1])
+        plane_z = tf.slice(plane_p, [0,2], [-1,1])
+        plane_l1 = normalize(light1 - plane_p)
+        plane_l2 = normalize(light2 - plane_p)
+        plane_xpos = plane_x - tf.floor(plane_x)
+        plane_zpos = plane_z - tf.floor(plane_z)
+        checkers = tf.logical_xor(tf.less(plane_xpos, 0.5),
+                       tf.less(plane_zpos, 0.5))
+        checkers = tf.squeeze(checkers, [1]) # (w*h,1) -> (w*h,)
+
+      with tf.name_scope('shadow1'):
+        ec = plane_p - sph_center
+        b = 2. * dot(plane_l1, ec)
+        c = dot(ec, ec) - square(sph_rad)
+        det = b*b - 4.*c
+        depth = (-b - tf.sqrt(det)) / 2.
+        illum_l1 = tf.select(tf.greater(depth, 0.),
+            img_of([0.,0,0]),
+            img_of(light1_col))
+
+      with tf.name_scope('shadow2'):
+        b = 2. * dot(plane_l2, ec)
+        det = b*b - 4.*c
+        depth = (-b - tf.sqrt(det)) / 2.
+        illum_l2 = tf.select(tf.greater(depth, 0.),
+            img_of([0.,0,0]),
+            img_of(light2_col))
+
+      with tf.name_scope('shade_plane'):
+        plane_col = tf.select(checkers, img_of(ground_col1),
+            img_of(ground_col2)) * (
+          plane_amb + plane_dif * (
+            depthwise(diffuse(plane_n, plane_l1)) * illum_l1 +
+            depthwise(diffuse(plane_n, plane_l2)) * illum_l2))
+
+      with tf.name_scope('fog'):
+        fog_dist = mag(plane_p - [0., plane_height, 0.])
+        fog = lerp(fog_dist, 1, 0, fog_near, fog_far)
+        plane_col *= depthwise(square(fog))
+
+      sky = img_of(sky_col) * (ray_y)
+      plane_depth = tf.squeeze(plane_depth, [1])
+      return tf.select(tf.greater(plane_depth, 0), plane_col, sky)
+
+    camera = tf.constant([0.,0,-5], name='camera')
+
+    with tf.name_scope('screen'):
+      # linspace from 0 to N-1, to get X and Y numbers for every row and col.
+      x = tf.linspace(0., W-1, W)
+      y = tf.linspace(0., H-1, H)
+
+      # add half pixel offset (center of pixel), put (0,0) in the center,
+      # scale so Y goes from -1 to +1
+      # invert Y
+      x = (x + .5 - W/2.) / (H/2.)
+      y = (y + .5 - H/2.) / (H/2.) * -1.
+
+      # turn x into array of <x,0,0> vectors
+      x = tf.reshape(x, [W,1]) * tf.reshape([1.,0,0], [1,3])
+      # and y
+      y = tf.reshape(y, [H,1,1]) * tf.reshape([0.,1.,0], [1,1,3])
+
+      # adding them together broadcasts into a 2d array of 3-vectors
+      screen = x + y
+
+      # reshape into 1D array of 3vecs
+      screen = tf.reshape(screen, [W*H, 3])
+
+      # screen coords to normalized rays
+      ray = normalize(screen - camera)
+
+    with tf.name_scope('intersect'):
+      ec = camera - sph_center
+      b = 2. * dot(ray, ec)
+      c = dot_vec(ec, ec) - square(sph_rad)
+      det = b*b - 4*c
+      sph_depth = (-b - tf.sqrt(det)) / 2.
+
+    with tf.name_scope('lights'):
+      sph_p = camera + depthwise(sph_depth) * ray
+      sph_l1 = normalize(light1 - sph_p)
+      sph_l2 = normalize(light2 - sph_p)
+      sph_n = (sph_p - sph_center) / sph_rad
+      sph_r = normalize(reflect(ray, sph_n))
+
+    with tf.name_scope('reflection'):
+      r_col = trace_bg(sph_p, sph_r)
+
+    with tf.name_scope('shade'):
+      sph = sph_col * (sph_amb +
+        sph_diff * depthwise(diffuse(sph_n, sph_l1)) * light1_col +
+        sph_diff * depthwise(diffuse(sph_n, sph_l2)) * light2_col +
+        sph_spec * depthwise(specular(sph_r, sph_l1, sph_shine)) * light1_col +
+        sph_spec * depthwise(specular(sph_r, sph_l2, sph_shine)) * light2_col +

         r_col * sph_refl)
 
         r_col * sph_refl)
 

-    bg = self.trace_bg(np.reshape(self.camera, (1,1,3)), ray)
-    return np.where(depthwise(sph_depth > 0), sph, bg)
+    with tf.name_scope('background'):
+      bg = trace_bg(img_of(camera), ray)
+
+    out = tf.select(tf.less(det, 0.), bg, sph)
+
+    with tf.name_scope('out'):
+      gamma = 2.2
+      out = tf.pow(out, 1. / gamma)
+      out = tf.reshape(out * 255., [H,W,3])
+
+    summary_writer = tf.train.SummaryWriter('log', sess.graph)
+    print 'evaluating'
+    evald = sess.run(out)
+    print 'saving'
+    save(evald, 'out.png')
+
+if __name__ == '__main__':
+  main()

 
 

-show(Scene().render(800, 600))

 # vim:set ts=2 sw=2 et:
 # vim:set ts=2 sw=2 et: