Pastebin launched a little side project called VERYVIRAL.com, check it out ;-) Want more features on Pastebin? Sign Up, it's FREE!
Guest

Tiny Tracer

By: a guest on Jan 9th, 2013  |  syntax: Python  |  size: 4.00 KB  |  views: 1,933  |  expires: Never
download  |  raw  |  embed  |  report abuse  |  print
Text below is selected. Please press Ctrl+C to copy to your clipboard. (⌘+C on Mac)
  1. from math import sqrt, pow, pi
  2. import Image
  3.  
  4. class Vector( object ):
  5.        
  6.         def __init__(self,x,y,z):
  7.                 self.x = x
  8.                 self.y = y
  9.                 self.z = z
  10.        
  11.         def dot(self, b):
  12.                 return self.x*b.x + self.y*b.y + self.z*b.z
  13.                
  14.         def cross(self, b):
  15.                 return (self.y*b.z-self.z*b.y, self.z*b.x-self.x*b.z, self.x*b.y-self.y*b.x)
  16.            
  17.         def magnitude(self):
  18.                 return sqrt(self.x**2+self.y**2+self.z**2)
  19.                
  20.         def normal(self):
  21.                 mag = self.magnitude()
  22.                 return Vector(self.x/mag,self.y/mag,self.z/mag)
  23.                
  24.         def __add__(self, b):
  25.                 return Vector(self.x + b.x, self.y+b.y, self.z+b.z)
  26.        
  27.         def __sub__(self, b):
  28.                 return Vector(self.x-b.x, self.y-b.y, self.z-b.z)
  29.                
  30.         def __mul__(self, b):
  31.                 assert type(b) == float or type(b) == int
  32.                 return Vector(self.x*b, self.y*b, self.z*b)            
  33.    
  34. class Sphere( object ):
  35.        
  36.         def __init__(self, center, radius, color):
  37.                 self.c = center
  38.                 self.r = radius
  39.                 self.col = color
  40.                
  41.         def intersection(self, l):
  42.                 q = l.d.dot(l.o - self.c)**2 - (l.o - self.c).dot(l.o - self.c) + self.r**2
  43.                 if q < 0:
  44.                         return Intersection( Vector(0,0,0), -1, Vector(0,0,0), self)
  45.                 else:
  46.                         d = -l.d.dot(l.o - self.c)
  47.                         d1 = d - sqrt(q)
  48.                         d2 = d + sqrt(q)
  49.                         if 0 < d1 and ( d1 < d2 or d2 < 0):
  50.                                 return Intersection(l.o+l.d*d1, d1, self.normal(l.o+l.d*d1), self)
  51.                         elif 0 < d2 and ( d2 < d1 or d1 < 0):
  52.                                 return Intersection(l.o+l.d*d2, d2, self.normal(l.o+l.d*d2), self)
  53.                         else:
  54.                                 return Intersection( Vector(0,0,0), -1, Vector(0,0,0), self)   
  55.                        
  56.         def normal(self, b):
  57.                 return (b - self.c).normal()
  58.                
  59. class Plane( object ):
  60.        
  61.         def __init__(self, point, normal, color):
  62.                 self.n = normal
  63.                 self.p = point
  64.                 self.col = color
  65.                
  66.         def intersection(self, l):
  67.                 d = l.d.dot(self.n)
  68.                 if d == 0:
  69.                         return Intersection( vector(0,0,0), -1, vector(0,0,0), self)
  70.                 else:
  71.                         d = (self.p - l.o).dot(self.n) / d
  72.                         return Intersection(l.o+l.d*d, d, self.n, self)
  73.                
  74. class Ray( object ):
  75.        
  76.         def __init__(self, origin, direction):
  77.                 self.o = origin
  78.                 self.d = direction
  79.                
  80. class Intersection( object ):
  81.        
  82.         def __init__(self, point, distance, normal, obj):
  83.                 self.p = point
  84.                 self.d = distance
  85.                 self.n = normal
  86.                 self.obj = obj
  87.                
  88. def testRay(ray, objects, ignore=None):
  89.         intersect = Intersection( Vector(0,0,0), -1, Vector(0,0,0), None)
  90.        
  91.         for obj in objects:
  92.                 if obj is not ignore:
  93.                         currentIntersect = obj.intersection(ray)
  94.                         if currentIntersect.d > 0 and intersect.d < 0:
  95.                                 intersect = currentIntersect
  96.                         elif 0 < currentIntersect.d < intersect.d:
  97.                                 intersect = currentIntersect
  98.         return intersect
  99.        
  100. def trace(ray, objects, light, maxRecur):
  101.         if maxRecur < 0:
  102.                 return (0,0,0)
  103.         intersect = testRay(ray, objects)              
  104.         if intersect.d == -1:
  105.                 col = vector(AMBIENT,AMBIENT,AMBIENT)
  106.         elif intersect.n.dot(light - intersect.p) < 0:
  107.                 col = intersect.obj.col * AMBIENT
  108.         else:
  109.                 lightRay = Ray(intersect.p, (light-intersect.p).normal())
  110.                 if testRay(lightRay, objects, intersect.obj).d == -1:
  111.                         lightIntensity = 1000.0/(4*pi*(light-intersect.p).magnitude()**2)
  112.                         col = intersect.obj.col * max(intersect.n.normal().dot((light - intersect.p).normal()*lightIntensity), AMBIENT)
  113.                 else:
  114.                         col = intersect.obj.col * AMBIENT
  115.         return col
  116.        
  117. def gammaCorrection(color,factor):
  118.         return (int(pow(color.x/255.0,factor)*255),
  119.                         int(pow(color.y/255.0,factor)*255),
  120.                         int(pow(color.z/255.0,factor)*255))
  121.                        
  122.  
  123. AMBIENT = 0.1
  124. GAMMA_CORRECTION = 1/2.2
  125.                                
  126. objs = []
  127. objs.append(Sphere( Vector(-2,0,-10), 2, Vector(0,255,0)))
  128. objs.append(Sphere( Vector(2,0,-10), 3.5, Vector(255,0,0)))
  129. objs.append(Sphere( Vector(0,-4,-10), 3, Vector(0,0,255)))
  130. objs.append(Plane( Vector(0,0,-12), Vector(0,0,1), Vector(255,255,255)))
  131. lightSource = Vector(0,10,0)
  132. img = Image.new("RGB",(500,500))
  133. cameraPos = Vector(0,0,20)
  134. for x in range(500):
  135.         print x
  136.         for y in range(500):
  137.                 ray = Ray( cameraPos, (Vector(x/50.0-5,y/50.0-5,0)-cameraPos).normal())
  138.                 col = trace(ray, objs, lightSource, 10)
  139.                 img.putpixel((x,499-y),gammaCorrection(col,GAMMA_CORRECTION))
  140. img.save("trace.bmp","BMP")