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