5a10bbbe557e630e0980d0ca405f7cf2fc555f92
3 Raytracer in tensorflow.
4 Emil Mikulic <emikulic@gmail.com> was here 2016.
10 import tensorflow
as tf
13 print 'loading took %.3f sec' % td
18 def save(tf_array
, fn
):
19 img
= np
.asarray(tf_array
)
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
)
28 with tf
.name_scope('square'):
32 return tf
.reduce_sum(a
* b
, 0)
35 """Dot product of arrays of vectors."""
36 return tf
.reduce_sum(a
* b
, 1)
39 return tf
.sqrt(dot(a
, a
))
42 return tf
.expand_dims(a
, 1)
45 with tf
.name_scope('normalize'):
46 return a
/ depthwise(mag(a
))
49 c
= n
* depthwise(dot(n
, -i
))
53 return tf
.maximum(dot(n
, l
), 0.)
55 def specular(r
, l
, power
):
56 return tf
.pow(diffuse(r
, l
), power
)
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
)
66 return tf
.zeros([H
*W
,3]) + v
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')
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,)
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
)
117 depth
= (-b
- tf
.sqrt(det
)) / 2.
118 illum_l1
= tf
.select(tf
.greater(depth
, 0.),
122 with tf
.name_scope('shadow2'):
123 b
= 2. * dot(plane_l2
, ec
)
125 depth
= (-b
- tf
.sqrt(det
)) / 2.
126 illum_l2
= tf
.select(tf
.greater(depth
, 0.),
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
))
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
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])
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
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
)
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
)
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
+
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'):
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
)
210 evald
= sess
.run(out
)
212 save(evald
, 'out.png')
214 if __name__
== '__main__':
217 # vim:set ts=2 sw=2 et: