Add antialiasing.
[the-worst-raytracer] / 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 for j in range(AA):
231 for i in range(AA):
232 sess.run(update, feed_dict={aa_x:i, aa_y:j})
233
234 t3 = time.time()
235 print 'took %.3f sec to render %d images' % (t3 - t2, AA*AA)
236
237 print 'finalizing'
238 evald = sess.run(out)
239 print 'saving'
240 save(evald, 'out.png')
241
242 if __name__ == '__main__':
243 main()
244
245 # vim:set ts=2 sw=2 sts=2 et: