5a10bbbe557e630e0980d0ca405f7cf2fc555f92
1 #!/usr/bin/env python
2 """
3 Raytracer in tensorflow.
4 Emil Mikulic <emikulic@gmail.com> was here 2016.
5 """
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
15 # width and height
16 W, H = 800, 600
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)
27 def square(x):
28 with tf.name_scope('square'):
29 return x * x
31 def dot_vec(a, b):
32 return tf.reduce_sum(a * b, 0)
34 def dot(a, b):
35 """Dot product of arrays of vectors."""
36 return tf.reduce_sum(a * b, 1)
38 def mag(a):
39 return tf.sqrt(dot(a, a))
41 def depthwise(a):
42 return tf.expand_dims(a, 1)
44 def normalize(a):
45 with tf.name_scope('normalize'):
46 return a / depthwise(mag(a))
48 def reflect(i, n):
49 c = n * depthwise(dot(n, -i))
50 return 2*c + i
52 def diffuse(n, l):
53 return tf.maximum(dot(n, l), 0.)
55 def specular(r, l, power):
56 return tf.pow(diffuse(r, l), power)
58 def clip(a, x, y):
59 return tf.minimum(tf.maximum(a, x), y)
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)
65 def img_of(v):
66 return tf.zeros([H*W,3]) + v
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')
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')
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')
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
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,)
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))
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))
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))
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))
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)
146 camera = tf.constant([0.,0,-5], name='camera')
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)
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.
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])
164 # adding them together broadcasts into a 2d array of 3-vectors
165 screen = x + y
167 # reshape into 1D array of 3vecs
168 screen = tf.reshape(screen, [W*H, 3])
170 # screen coords to normalized rays
171 ray = normalize(screen - camera)
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.
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))
187 with tf.name_scope('reflection'):
188 r_col = trace_bg(sph_p, sph_r)
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)
198 with tf.name_scope('background'):
199 bg = trace_bg(img_of(camera), ray)
201 out = tf.select(tf.less(det, 0.), bg, sph)
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])
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')
214 if __name__ == '__main__':
215 main()
217 # vim:set ts=2 sw=2 et: