unix4lyfe git - the-worst-raytracer/blob - tf.py

   1 #!/usr/bin/env python
   2 """
   3 Raytracer in tensorflow.
   4 Emil Mikulic <emikulic@gmail.com> was here 2016.
   5 """
   6 print 'loading'
   7 import time
   8 t0 = time.time()
   9 import numpy as np
  10 import tensorflow as tf
  11 import Image
  12 t1 = time.time()
  13 print 'loading took %.3f sec' % (t1 - t0)
  14
  15 # width and height
  16 W, H = 800, 600
  17
  18 # Antialias using AA^2 samples per pixel.
  19 AA = 4
  20
  21 def save(tf_array, fn):
  22   img = np.asarray(tf_array)
  23   a = np.nanmin(img)
  24   b = np.nanmax(img)
  25   print 'range is', (a, b)
  26   #img = (img - a) / (b - a)
  27   img = np.round(img).clip(0, 255)
  28   Image.fromarray(img.astype(np.uint8)).save(fn)
  29
  30 def square(x):
  31   with tf.name_scope('square'):
  32     return x * x
  33
  34 def dot_vec(a, b):
  35   return tf.reduce_sum(a * b, 0)
  36
  37 def dot(a, b):
  38   """Dot product of arrays of vectors."""
  39   return tf.reduce_sum(a * b, 1)
  40
  41 def mag(a):
  42   return tf.sqrt(dot(a, a))
  43
  44 def depthwise(a):
  45   return tf.expand_dims(a, 1)
  46
  47 def normalize(a):
  48   with tf.name_scope('normalize'):
  49     return a / depthwise(mag(a))
  50
  51 def reflect(i, n):
  52   c = n * depthwise(dot(n, -i))
  53   return 2*c + i
  54
  55 def diffuse(n, l):
  56   return tf.maximum(dot(n, l), 0.)
  57
  58 def specular(r, l, power):
  59   return tf.pow(diffuse(r, l), power)
  60
  61 def clip(a, x, y):
  62   return tf.minimum(tf.maximum(a, x), y)
  63
  64 def lerp(ar, a, b, u, v):
  65   "[u,v] is mapped to [a,b]"
  66   return a + (b-a) * (clip(ar,u,v) - u) / (v-u)
  67
  68 def img_of(v):
  69   return tf.zeros([H*W,3]) + v
  70
  71 def main():
  72   print 'building'
  73   with tf.Session() as sess:
  74     light1 = tf.constant([20.,10.,-10.], name='light1')
  75     light1_col = tf.constant([1., .5, .3], name='light1_col')
  76     light2 = tf.constant([-5.,5.,5.], name='light2')
  77     light2_col = tf.constant([.1, .4, .7], name='light2_col')
  78     sky_col = tf.constant([.2,.3,.6], name='sky_col')
  79     sph_center = tf.constant([-.5,0,0], name='sph_center')
  80     sph_rad = tf.constant(.5, name='sph_rad')
  81     sph_col = tf.constant([1.,1.,1.], name='sph_col')
  82     ground_col1 = tf.constant([.4,.4,.4], name='ground_col1')
  83     ground_col2 = tf.constant([.8,.8,.8], name='ground_col2')
  84     plane_n = tf.constant([0.,1.,0.], name='plane_n')
  85     plane_amb = tf.constant(0.02, name='plane_amb')
  86     plane_dif = tf.constant(0.9, name='plane_dif')
  87     plane_height = tf.constant(-.5, name='plane_height')
  88     fog_near = tf.constant(0., name='fog_near')
  89     fog_far = tf.constant(40., name='fog_far')
  90
  91     sph_amb = tf.constant(0.02, name='sph_amb')
  92     sph_diff = tf.constant(0.05, name='sph_diff')
  93     sph_spec = tf.constant(1.1, name='sph_spec')
  94     sph_refl = tf.constant(.5, name='sph_refl')
  95     sph_shine = tf.constant(20., name='sph_shine')
  96
  97     def trace_bg(camera, ray):
  98       with tf.name_scope('intersect'):
  99         camera_y = tf.slice(camera, [0,1], [-1,1])
 100         ray_y = tf.slice(ray, [0,1], [-1,1])
 101         plane_depth = (plane_height - camera_y) / ray_y
 102         plane_p = camera + (plane_depth) * ray
 103
 104       with tf.name_scope('checkers'):
 105         plane_x = tf.slice(plane_p, [0,0], [-1,1])
 106         plane_z = tf.slice(plane_p, [0,2], [-1,1])
 107         plane_l1 = normalize(light1 - plane_p)
 108         plane_l2 = normalize(light2 - plane_p)
 109         plane_xpos = plane_x - tf.floor(plane_x)
 110         plane_zpos = plane_z - tf.floor(plane_z)
 111         checkers = tf.logical_xor(tf.less(plane_xpos, 0.5),
 112                        tf.less(plane_zpos, 0.5))
 113         checkers = tf.squeeze(checkers, [1]) # (w*h,1) -> (w*h,)
 114
 115       with tf.name_scope('shadow1'):
 116         ec = plane_p - sph_center
 117         b = 2. * dot(plane_l1, ec)
 118         c = dot(ec, ec) - square(sph_rad)
 119         det = b*b - 4.*c
 120         depth = (-b - tf.sqrt(det)) / 2.
 121         illum_l1 = tf.select(tf.greater(depth, 0.),
 122             img_of([0.,0,0]),
 123             img_of(light1_col))
 124
 125       with tf.name_scope('shadow2'):
 126         b = 2. * dot(plane_l2, ec)
 127         det = b*b - 4.*c
 128         depth = (-b - tf.sqrt(det)) / 2.
 129         illum_l2 = tf.select(tf.greater(depth, 0.),
 130             img_of([0.,0,0]),
 131             img_of(light2_col))
 132
 133       with tf.name_scope('shade_plane'):
 134         plane_col = tf.select(checkers, img_of(ground_col1),
 135             img_of(ground_col2)) * (
 136           plane_amb + plane_dif * (
 137             depthwise(diffuse(plane_n, plane_l1)) * illum_l1 +
 138             depthwise(diffuse(plane_n, plane_l2)) * illum_l2))
 139
 140       with tf.name_scope('fog'):
 141         fog_dist = mag(plane_p - [0., plane_height, 0.])
 142         fog = lerp(fog_dist, 1, 0, fog_near, fog_far)
 143         plane_col *= depthwise(square(fog))
 144
 145       sky = img_of(sky_col) * (ray_y)
 146       plane_depth = tf.squeeze(plane_depth, [1])
 147       return tf.select(tf.greater(plane_depth, 0), plane_col, sky)
 148
 149     camera = tf.constant([0.,0,-5], name='camera')
 150     aa_x = tf.placeholder(tf.float32, name='aa_x')
 151     aa_y = tf.placeholder(tf.float32, name='aa_y')
 152     accumulator = tf.Variable(tf.zeros([H*W,3]), name='accumulator')
 153
 154     with tf.name_scope('screen'):
 155       # linspace from 0 to N-1, to get X and Y numbers for every row and col.
 156       x = tf.linspace(0., W-1, W)
 157       y = tf.linspace(0., H-1, H)
 158
 159       # antialiasing pixel offsets
 160       ofs_x = (aa_x + .5) / AA
 161       ofs_y = (aa_y + .5) / AA
 162
 163       # add pixel offset (antialiasing), center the image on (0,0)
 164       # scale so Y goes from -1 to +1
 165       # invert Y
 166       x = (x + ofs_x - W/2.) / (H/2.)
 167       y = (y + ofs_y - H/2.) / (H/2.) * -1.
 168
 169       # turn x into array of <x,0,0> vectors
 170       x = tf.reshape(x, [W,1]) * tf.reshape([1.,0,0], [1,3])
 171       # and y
 172       y = tf.reshape(y, [H,1,1]) * tf.reshape([0.,1.,0], [1,1,3])
 173
 174       # adding them together broadcasts into a 2d array of 3-vectors
 175       screen = x + y
 176
 177       # reshape into 1D array of 3vecs
 178       screen = tf.reshape(screen, [W*H, 3])
 179
 180       # screen coords to normalized rays
 181       ray = normalize(screen - camera)
 182
 183     with tf.name_scope('intersect'):
 184       ec = camera - sph_center
 185       b = 2. * dot(ray, ec)
 186       c = dot_vec(ec, ec) - square(sph_rad)
 187       det = b*b - 4*c
 188       sph_depth = (-b - tf.sqrt(det)) / 2.
 189
 190     with tf.name_scope('lights'):
 191       sph_p = camera + depthwise(sph_depth) * ray
 192       sph_l1 = normalize(light1 - sph_p)
 193       sph_l2 = normalize(light2 - sph_p)
 194       sph_n = (sph_p - sph_center) / sph_rad
 195       sph_r = normalize(reflect(ray, sph_n))
 196
 197     with tf.name_scope('reflection'):
 198       r_col = trace_bg(sph_p, sph_r)
 199
 200     with tf.name_scope('shade'):
 201       sph = sph_col * (sph_amb +
 202         sph_diff * depthwise(diffuse(sph_n, sph_l1)) * light1_col +
 203         sph_diff * depthwise(diffuse(sph_n, sph_l2)) * light2_col +
 204         sph_spec * depthwise(specular(sph_r, sph_l1, sph_shine)) * light1_col +
 205         sph_spec * depthwise(specular(sph_r, sph_l2, sph_shine)) * light2_col +
 206         r_col * sph_refl)
 207
 208     with tf.name_scope('background'):
 209       bg = trace_bg(img_of(camera), ray)
 210
 211     # Mask sphere and background.
 212     img = tf.select(tf.less(det, 0.), bg, sph)
 213
 214     # Update accumulator.
 215     new_accum = accumulator + img / AA / AA
 216     update = tf.assign(accumulator, new_accum)
 217
 218     with tf.name_scope('out'):
 219       gamma = 2.2
 220       out = tf.pow(accumulator, 1. / gamma)
 221       out = tf.reshape(out * 255., [H,W,3])
 222
 223     summary_writer = tf.train.SummaryWriter('log', sess.graph)
 224     sess.run(tf.initialize_all_variables())
 225
 226     t2 = time.time()
 227     print 'initializing took %.3f sec' % (t2 - t1)
 228
 229     print 'accumulating'
 230     t3 = t2
 231     for j in range(AA):
 232       for i in range(AA):
 233         sess.run(update, feed_dict={aa_x:i, aa_y:j})
 234         t4 = time.time()
 235         print 'took %.3f sec' % (t4 - t3)
 236         t3 = t4
 237
 238     print 'took %.3f sec to render %d images' % (t3 - t2, AA*AA)
 239
 240     print 'finalizing'
 241     evald = sess.run(out)
 242     print 'saving'
 243     save(evald, 'out.png')
 244
 245 if __name__ == '__main__':
 246   main()
 247
 248 # vim:set ts=2 sw=2 sts=2 et: