3 Raytracer in tensorflow.
4 Emil Mikulic <emikulic@gmail.com> was here 2016.
10 import tensorflow
as tf
13 print 'loading took %.3f sec' % (t1
- t0
)
18 # Antialias using AA^2 samples per pixel.
21 def save(tf_array
, fn
):
22 img
= np
.asarray(tf_array
)
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
)
31 with tf
.name_scope('square'):
35 return tf
.reduce_sum(a
* b
, 0)
38 """Dot product of arrays of vectors."""
39 return tf
.reduce_sum(a
* b
, 1)
42 return tf
.sqrt(dot(a
, a
))
45 return tf
.expand_dims(a
, 1)
48 with tf
.name_scope('normalize'):
49 return a
/ depthwise(mag(a
))
52 c
= n
* depthwise(dot(n
, -i
))
56 return tf
.maximum(dot(n
, l
), 0.)
58 def specular(r
, l
, power
):
59 return tf
.pow(diffuse(r
, l
), power
)
62 return tf
.minimum(tf
.maximum(a
, x
), y
)
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
)
69 return tf
.zeros([H
*W
,3]) + v
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')
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')
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
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,)
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
)
120 depth
= (-b
- tf
.sqrt(det
)) / 2.
121 illum_l1
= tf
.select(tf
.greater(depth
, 0.),
125 with tf
.name_scope('shadow2'):
126 b
= 2. * dot(plane_l2
, ec
)
128 depth
= (-b
- tf
.sqrt(det
)) / 2.
129 illum_l2
= tf
.select(tf
.greater(depth
, 0.),
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
))
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
))
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
)
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')
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
)
159 # antialiasing pixel offsets
160 ofs_x
= (aa_x
+ .5) / AA
161 ofs_y
= (aa_y
+ .5) / AA
163 # add pixel offset (antialiasing), center the image on (0,0)
164 # scale so Y goes from -1 to +1
166 x
= (x
+ ofs_x
- W
/2.) / (H
/2.)
167 y
= (y
+ ofs_y
- H
/2.) / (H
/2.) * -1.
169 # turn x into array of <x,0,0> vectors
170 x
= tf
.reshape(x
, [W
,1]) * tf
.reshape([1.,0,0], [1,3])
172 y
= tf
.reshape(y
, [H
,1,1]) * tf
.reshape([0.,1.,0], [1,1,3])
174 # adding them together broadcasts into a 2d array of 3-vectors
177 # reshape into 1D array of 3vecs
178 screen
= tf
.reshape(screen
, [W
*H
, 3])
180 # screen coords to normalized rays
181 ray
= normalize(screen
- camera
)
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
)
188 sph_depth
= (-b
- tf
.sqrt(det
)) / 2.
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
))
197 with tf
.name_scope('reflection'):
198 r_col
= trace_bg(sph_p
, sph_r
)
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
+
208 with tf
.name_scope('background'):
209 bg
= trace_bg(img_of(camera
), ray
)
211 # Mask sphere and background.
212 img
= tf
.select(tf
.less(det
, 0.), bg
, sph
)
214 # Update accumulator.
215 new_accum
= accumulator
+ img
/ AA
/ AA
216 update
= tf
.assign(accumulator
, new_accum
)
218 with tf
.name_scope('out'):
220 out
= tf
.pow(accumulator
, 1. / gamma
)
221 out
= tf
.reshape(out
* 255., [H
,W
,3])
223 summary_writer
= tf
.train
.SummaryWriter('log', sess
.graph
)
224 sess
.run(tf
.initialize_all_variables())
227 print 'initializing took %.3f sec' % (t2
- t1
)
232 sess
.run(update
, feed_dict
={aa_x
:i
, aa_y
:j
})
235 print 'took %.3f sec to render %d images' % (t3
- t2
, AA
*AA
)
238 evald
= sess
.run(out
)
240 save(evald
, 'out.png')
242 if __name__
== '__main__':
245 # vim:set ts=2 sw=2 sts=2 et: