From: Emil Mikulic Date: Mon, 5 Sep 2016 06:17:32 +0000 (+1000) Subject: Port raytracer to TensorFlow. X-Git-Url: https://unix4lyfe.org/gitweb/the-worst-raytracer/commitdiff_plain/454cd24ad1b2ab730e47a186d7d223a7a49370a4 Port raytracer to TensorFlow. --- diff --git a/tf.py b/tf.py index 2bc8365..5a10bbb 100755 --- a/tf.py +++ b/tf.py @@ -1,143 +1,217 @@ #!/usr/bin/env python """ -The worst raytracer. -Emil Mikulic was here 2012. +Raytracer in tensorflow. +Emil Mikulic was here 2016. """ +print 'loading' +import time +t0 = time.time() import numpy as np -from accidental_complexity import show -# this file contains mostly essential complexity - -def vec(x, y, z): - return np.array((x, y, z), dtype=np.float) - -def color(r, g, b): - return np.reshape(vec(r,g,b), (1,1,3)) +import tensorflow as tf +import Image +td = time.time() - t0 +print 'loading took %.3f sec' % td + +# width and height +W, H = 800, 600 + +def save(tf_array, fn): + img = np.asarray(tf_array) + a = np.nanmin(img) + b = np.nanmax(img) + print 'range is', (a, b) + #img = (img - a) / (b - a) + img = np.round(img).clip(0, 255) + Image.fromarray(img.astype(np.uint8)).save(fn) + +def square(x): + with tf.name_scope('square'): + return x * x + +def dot_vec(a, b): + return tf.reduce_sum(a * b, 0) def dot(a, b): - return (a * b).sum(axis=2) + """Dot product of arrays of vectors.""" + return tf.reduce_sum(a * b, 1) def mag(a): - return np.sqrt(np.square(a).sum(axis=2)) + return tf.sqrt(dot(a, a)) def depthwise(a): - return np.expand_dims(a, axis=2) + return tf.expand_dims(a, 1) -def norm(a): - return a / depthwise(mag(a)) +def normalize(a): + with tf.name_scope('normalize'): + return a / depthwise(mag(a)) def reflect(i, n): c = n * depthwise(dot(n, -i)) return 2*c + i def diffuse(n, l): - return depthwise(dot(n, l).clip(0, 1)) + return tf.maximum(dot(n, l), 0.) def specular(r, l, power): - return pow(diffuse(r, l), power) + return tf.pow(diffuse(r, l), power) + +def clip(a, x, y): + return tf.minimum(tf.maximum(a, x), y) def lerp(ar, a, b, u, v): "[u,v] is mapped to [a,b]" - return a + (b-a) * (ar.clip(u,v) - u) / (v-u) - -class Scene: - def __init__(self): - self.light1 = vec(20,10,-10) - self.light1_col = color(1,.5,.3) - self.light2 = vec(-5,5,5) - self.light2_col = color(1,1,1) - self.light1_col - self.sky_col = color(.2, .3, .6) - self.sph_center = vec(-.5,0,0) - self.sph_rad = .5 - self.sph_col = color(1,1,1) - self.camera = vec(0,0,-5) - self.ground_col1 = color(.4,.4,.4) - self.ground_col2 = color(.8,.8,.8) - self.plane_n = vec(0,1,0) - self.plane_amb = 0.02 - self.plane_dif = 0.9 - self.plane_height = -0.5 - self.fog_near = 0 - self.fog_far = 40 - - def trace_bg(self, camera, ray): - # background - plane_depth = (self.plane_height - camera[:,:,1]) / ray[:,:,1] - plane_p = camera + depthwise(plane_depth) * ray - plane_l1 = norm(self.light1 - plane_p) - plane_l2 = norm(self.light2 - plane_p) - plane_xpos = plane_p[:,:,0] - np.floor(plane_p[:,:,0]) - plane_zpos = plane_p[:,:,2] - np.floor(plane_p[:,:,2]) - checkers = depthwise(np.logical_xor(plane_xpos < 0.5, plane_zpos < 0.5)) - - # trace shadows (light 1) - ec = plane_p - self.sph_center - b = 2 * dot(plane_l1, ec) - c = dot(ec, ec) - np.square(self.sph_rad) - det = b*b - 4*c - depth = (-b - np.sqrt(det)) / 2 - illum_l1 = np.where(depthwise(depth > 0), 0, self.light1_col) - - # light 2 - b = 2 * dot(plane_l2, ec) - det = b*b - 4*c - depth = (-b - np.sqrt(det)) / 2 - illum_l2 = np.where(depthwise(depth > 0), 0, self.light2_col) - - # shade plane - plane_col = np.where(checkers, self.ground_col1, self.ground_col2) * ( - self.plane_amb + self.plane_dif * ( - diffuse(self.plane_n, plane_l1) * illum_l1 + - diffuse(self.plane_n, plane_l2) * illum_l2)) - - fog_dist = mag(plane_p - vec(0, self.plane_height, 0)) - fog = lerp(fog_dist, 1, 0, self.fog_near, self.fog_far) - plane_col *= depthwise(np.square(fog)) - - sky = self.sky_col * depthwise(ray[:,:,1]) - return np.where(depthwise(plane_depth > 0), plane_col, sky) - - def render(self, w, h): - # set up scene - sph_amb = 0.02 - sph_diff = 0.05 - sph_spec = 1.1 - sph_refl = 0.5 - sph_shine = 20 - - # points on the screen plane (z = 0) - aspect = float(w) / float(h) - scr_x = np.linspace(-1 + 1./w, 1 - 1./w, w) * aspect - scr_y = np.linspace(-1 + 1./h, 1 - 1./h, h) * -1 - scr = np.zeros((h,w,3), dtype=np.float) - scr[:,:,0] = np.reshape(scr_x, (1, w)) - scr[:,:,1] = np.reshape(scr_y, (h, 1)) - - ray = norm(scr - self.camera) - - # sphere - ec = self.camera - self.sph_center - b = 2 * dot(ray, ec) - c = ec.dot(ec) - np.square(self.sph_rad) - det = b*b - 4*c - sph_depth = (-b - np.sqrt(det)) / 2 - - # sphere color - sph_p = self.camera + depthwise(sph_depth) * ray - sph_l1 = norm(self.light1 - sph_p) - sph_l2 = norm(self.light2 - sph_p) - sph_n = (sph_p - self.sph_center) / self.sph_rad - sph_r = norm(reflect(ray, sph_n)) - r_col = self.trace_bg(sph_p, sph_r) - - sph = self.sph_col * (sph_amb + - sph_diff * diffuse(sph_n, sph_l1) * self.light1_col + - sph_diff * diffuse(sph_n, sph_l2) * self.light2_col + - sph_spec * specular(sph_r, sph_l1, sph_shine) * self.light1_col + - sph_spec * specular(sph_r, sph_l2, sph_shine) * self.light2_col + + return a + (b-a) * (clip(ar,u,v) - u) / (v-u) + +def img_of(v): + return tf.zeros([H*W,3]) + v + +def main(): + print 'building' + with tf.Session() as sess: + light1 = tf.constant([20.,10.,-10.], name='light1') + light1_col = tf.constant([1., .5, .3], name='light1_col') + light2 = tf.constant([-5.,5.,5.], name='light2') + light2_col = tf.constant([.1, .4, .7], name='light2_col') + sky_col = tf.constant([.2,.3,.6], name='sky_col') + sph_center = tf.constant([-.5,0,0], name='sph_center') + sph_rad = tf.constant(.5, name='sph_rad') + sph_col = tf.constant([1.,1.,1.], name='sph_col') + ground_col1 = tf.constant([.4,.4,.4], name='ground_col1') + ground_col2 = tf.constant([.8,.8,.8], name='ground_col2') + plane_n = tf.constant([0.,1.,0.], name='plane_n') + plane_amb = tf.constant(0.02, name='plane_amb') + plane_dif = tf.constant(0.9, name='plane_dif') + plane_height = tf.constant(-.5, name='plane_height') + fog_near = tf.constant(0., name='fog_near') + fog_far = tf.constant(40., name='fog_far') + + sph_amb = tf.constant(0.02, name='sph_amb') + sph_diff = tf.constant(0.05, name='sph_diff') + sph_spec = tf.constant(1.1, name='sph_spec') + sph_refl = tf.constant(.5, name='sph_refl') + sph_shine = tf.constant(20., name='sph_shine') + + def trace_bg(camera, ray): + with tf.name_scope('intersect'): + camera_y = tf.slice(camera, [0,1], [-1,1]) + ray_y = tf.slice(ray, [0,1], [-1,1]) + plane_depth = (plane_height - camera_y) / ray_y + plane_p = camera + (plane_depth) * ray + + with tf.name_scope('checkers'): + plane_x = tf.slice(plane_p, [0,0], [-1,1]) + plane_z = tf.slice(plane_p, [0,2], [-1,1]) + plane_l1 = normalize(light1 - plane_p) + plane_l2 = normalize(light2 - plane_p) + plane_xpos = plane_x - tf.floor(plane_x) + plane_zpos = plane_z - tf.floor(plane_z) + checkers = tf.logical_xor(tf.less(plane_xpos, 0.5), + tf.less(plane_zpos, 0.5)) + checkers = tf.squeeze(checkers, [1]) # (w*h,1) -> (w*h,) + + with tf.name_scope('shadow1'): + ec = plane_p - sph_center + b = 2. * dot(plane_l1, ec) + c = dot(ec, ec) - square(sph_rad) + det = b*b - 4.*c + depth = (-b - tf.sqrt(det)) / 2. + illum_l1 = tf.select(tf.greater(depth, 0.), + img_of([0.,0,0]), + img_of(light1_col)) + + with tf.name_scope('shadow2'): + b = 2. * dot(plane_l2, ec) + det = b*b - 4.*c + depth = (-b - tf.sqrt(det)) / 2. + illum_l2 = tf.select(tf.greater(depth, 0.), + img_of([0.,0,0]), + img_of(light2_col)) + + with tf.name_scope('shade_plane'): + plane_col = tf.select(checkers, img_of(ground_col1), + img_of(ground_col2)) * ( + plane_amb + plane_dif * ( + depthwise(diffuse(plane_n, plane_l1)) * illum_l1 + + depthwise(diffuse(plane_n, plane_l2)) * illum_l2)) + + with tf.name_scope('fog'): + fog_dist = mag(plane_p - [0., plane_height, 0.]) + fog = lerp(fog_dist, 1, 0, fog_near, fog_far) + plane_col *= depthwise(square(fog)) + + sky = img_of(sky_col) * (ray_y) + plane_depth = tf.squeeze(plane_depth, [1]) + return tf.select(tf.greater(plane_depth, 0), plane_col, sky) + + camera = tf.constant([0.,0,-5], name='camera') + + with tf.name_scope('screen'): + # linspace from 0 to N-1, to get X and Y numbers for every row and col. + x = tf.linspace(0., W-1, W) + y = tf.linspace(0., H-1, H) + + # add half pixel offset (center of pixel), put (0,0) in the center, + # scale so Y goes from -1 to +1 + # invert Y + x = (x + .5 - W/2.) / (H/2.) + y = (y + .5 - H/2.) / (H/2.) * -1. + + # turn x into array of vectors + x = tf.reshape(x, [W,1]) * tf.reshape([1.,0,0], [1,3]) + # and y + y = tf.reshape(y, [H,1,1]) * tf.reshape([0.,1.,0], [1,1,3]) + + # adding them together broadcasts into a 2d array of 3-vectors + screen = x + y + + # reshape into 1D array of 3vecs + screen = tf.reshape(screen, [W*H, 3]) + + # screen coords to normalized rays + ray = normalize(screen - camera) + + with tf.name_scope('intersect'): + ec = camera - sph_center + b = 2. * dot(ray, ec) + c = dot_vec(ec, ec) - square(sph_rad) + det = b*b - 4*c + sph_depth = (-b - tf.sqrt(det)) / 2. + + with tf.name_scope('lights'): + sph_p = camera + depthwise(sph_depth) * ray + sph_l1 = normalize(light1 - sph_p) + sph_l2 = normalize(light2 - sph_p) + sph_n = (sph_p - sph_center) / sph_rad + sph_r = normalize(reflect(ray, sph_n)) + + with tf.name_scope('reflection'): + r_col = trace_bg(sph_p, sph_r) + + with tf.name_scope('shade'): + sph = sph_col * (sph_amb + + sph_diff * depthwise(diffuse(sph_n, sph_l1)) * light1_col + + sph_diff * depthwise(diffuse(sph_n, sph_l2)) * light2_col + + sph_spec * depthwise(specular(sph_r, sph_l1, sph_shine)) * light1_col + + sph_spec * depthwise(specular(sph_r, sph_l2, sph_shine)) * light2_col + r_col * sph_refl) - bg = self.trace_bg(np.reshape(self.camera, (1,1,3)), ray) - return np.where(depthwise(sph_depth > 0), sph, bg) + with tf.name_scope('background'): + bg = trace_bg(img_of(camera), ray) + + out = tf.select(tf.less(det, 0.), bg, sph) + + with tf.name_scope('out'): + gamma = 2.2 + out = tf.pow(out, 1. / gamma) + out = tf.reshape(out * 255., [H,W,3]) + + summary_writer = tf.train.SummaryWriter('log', sess.graph) + print 'evaluating' + evald = sess.run(out) + print 'saving' + save(evald, 'out.png') + +if __name__ == '__main__': + main() -show(Scene().render(800, 600)) # vim:set ts=2 sw=2 et: