ray.py

import cv2
import sys
import numpy as np
import numpy.linalg
import math
import time

class Primitive(object):

    def __init__(self):
        self.color = (150,50,50)


class Plane(Primitive):

    def __init__(self,point,normal):
        self.point = point
        self.norm = normal
        super(Plane,self).__init__()

    def solve(self,r):
        e = r.e
        d = r.d
        t = e-self.point
        a = d.dot(self.norm)
        b = t.dot(self.norm)
        return [-b / a]

    def getNormalToSurface(self,point):
        return self.norm

class Triangle(Primitive):

    def __init__(self,a,b,c):
        self.a = a
        self.b = b
        self.c = c
        super(Triangle,self).__init__()

    def solve(self,r):
        e1 = r.e
        d1 = r.d

        a = self.a[0] - self.b[0]
        b = self.a[1] - self.b[1]
        c = self.a[2] - self.b[2]
    
        d = self.a[0] - self.c[0]
        e = self.a[1] - self.c[1]
        f = self.a[2] - self.c[2]

        g = d1[0]
        h = d1[1]
        i = d1[2]

        j = self.a[0] - e1[0]
        k = self.a[1] - e1[1]
        l = self.a[2] - e1[2]

        M = a*(e*i-h*f)+b*(g*f-d*i)+c*(d*h-e*g)
        
        t = (-f*(a*k-j*b)+e*(j*c-a*l)+d*(b*l-k*c))/M
        sigma = (i*(a*k-j*b)+h*(j*c-a*l)+g*(b*l-k*c))/M
        beta = (j*(e*i-h*f)+k*(g*f-d*i)+l*(d*h-e*g))/M

        #print t
        if sigma<0 or sigma>1:
            return []
        elif beta<0 or beta>1-sigma:
            return []
        else:
            return [t]
    
    def getNormalToSurface(self,point):
        return np.cross((self.a - self.b),(self.a - point))

class Sphere(Primitive):

    def __init__(self,c,R):
        self.c = c
        self.R = R
        super(Sphere,self).__init__()

    def solve(self,r):
        solutions = []
        e = r.e
        d = r.d
        c = self.c

        a = -d.dot(e-c)
        b = d.dot(d)
        f = math.pow(d.dot(e-c),2) - (d.dot(d))*((e-c).dot(e-c) - math.pow(self.R,2))

        #print a,b,f

        if f>=0:
            solutions.append((a/b + math.sqrt(f)/b))
            solutions.append((a/b - math.sqrt(f)/b))
        return solutions

    def getNormalToSurface(self,point):
        a = (1./self.R)*(point-self.c)
        return a
    
class Scene():

    def __init__(self):
        self.objectList = []
        self.lightList = []
        self.defaultColor = (0,0,0)

    def addObject(self,obj):
        self.objectList.append(obj)

    def addLight(self,light):
        self.lightList.append(light)

    def shadowIntersect(self,ray):
        for obj in self.objectList:
            solutions = obj.solve(ray)
            if len(solutions)>0:
                for s in solutions:
                    if s>0:
                        return True
        return False

    def intersect(self,ray):
        color = self.defaultColor
        minT = float("inf")
        for obj in self.objectList:
            solutions = obj.solve(ray)
            if len(solutions)>0:
                for s in solutions:
                    if s<minT and s>0:
                        intersectionPoint = ray.e+s*ray.d
                        color= np.array([30,30,30]) + sum([light.doPingPongShade(ray,obj,intersectionPoint) for light in self.lightList])
                        minT = s
        return color

class Light():

    def __init__(self,pos):
        self.position = pos
        self.I = np.array([1,1,1])

    def doLambertianShade(self,ray,obj,intersectionPoint):

        n = obj.getNormalToSurface(intersectionPoint)
        l = self.position - intersectionPoint
        v = ray.e - intersectionPoint

        #norm everything, need a better way to do this...
        l = (1./math.sqrt(l.dot(l)))*l
        v = (1./math.sqrt(v.dot(v)))*v
        n = (1./math.sqrt(n.dot(n)))*n
        
        color0 = np.array(obj.color)*self.I*max(0,n.dot(l))

        return color0

    def doPingPongShade(self,ray,obj,intersectionPoint):

        l = self.position - intersectionPoint
        v = ray.e - intersectionPoint
        n = obj.getNormalToSurface(intersectionPoint)
        h = (v+1)/np.linalg.norm((v+l))

        #norm everything, need a better way to do this...
        l = (1./math.sqrt(l.dot(l)))*l
        v = (1./math.sqrt(v.dot(v)))*v
        n = (1./math.sqrt(n.dot(n)))*n
        
        color0 = np.array(obj.color)*self.I*max(0,n.dot(l)) + np.array([20,20,20])*self.I*(max(0,n.dot(h))**100)

        return color0 

class Ray():

    def __init__(self,e,d):
        self.e = e
        self.d = d

    def __str__(self):
        return str(self.e) + " " + str(self.d)
    
class Camera():
    #imagePlane = [top,left,bottom,right]
    #e = camera location vector3
    #w = basis vector (backwards)
    #u = basis vector (rightwards)
    #v = basis vector (upwards)
    #{u,v,w} form a orthonormal basis

    #we should form u,v,w from e,w, and UP

    #d distance
    def __init__(self):
        self.e = np.array([0,0,-3])
        self.d = 2
        self.top = 10
        self.right = 10
        self.bottom = -self.top
        self.left = -self.right
        self.w =  np.array([0,0,-1])
        self.w = (1/math.sqrt(self.w.dot(self.w)))*self.w
        self.u =  np.array([1,0,0])
        self.u = (1/math.sqrt(self.u.dot(self.u)))*self.u
        self.v = np.cross(self.w,self.u)
        self.nx = 1024
        self.ny = 1024
        self.f = 1

        

    def getRay(self,x,y):
        u = self.left + (self.right - self.left)*(x+0.5) / self.nx
        v = self.bottom + (self.top - self.bottom)*(y+0.5) / self.ny

        #print u,v
        d = -self.d*self.w + u*self.u + v*self.v
        e = self.e
        return Ray(e,d)
        
        
def main():
    start = time.time()
    cam = Camera()
    
    s = Sphere(np.array([0,0,0]),1)
    #l = Light(np.array([0,5,0]))
    #l2 = Light(np.array([10,5,0]))
    l3 = Light(np.array([5,5,0]))

    #s2 = Sphere(np.array([0,1,0]),2)
    #s2.color = 85
    
    scene = Scene()
    scene.addObject(s)
    #scene.addLight(l)
    #scene.addLight(l2)
    scene.addLight(l3)
    #scene.addObject(s2)

    #triangle = Triangle(np.array([0,0,0]),np.array([0,50,0]),np.array([0,50,50]))
    #triangle2 = Triangle(np.array([0,0,0]),np.array([0,0,50]),np.array([50,0,50]))
    #scene.addObject(triangle)
    #scene.addObject(triangle2)
    #img = np.zeros((cam.ny,cam.nx))

    plane = Plane(np.array([0,-1,0]),np.array([0,1,0]))
    scene.addObject(plane)


    img = np.zeros((cam.nx,cam.ny,3), np.uint8)
    
    for i in range(cam.ny):
        for j in range(cam.nx):
            r = cam.getRay(j,i)
            img[i][j] = scene.intersect(r)
    print("Finished in %.3f seconds"%(time.time()-start))
            
    cv2.imwrite("test.png",img)
if __name__ == "__main__":
    main()