raytracer.py

"""
Requires
pillow
numpy
"""

from PIL import Image
from itertools import product
from math import sin, cos, sqrt, inf
from numpy import dot, subtract, multiply, add
from numpy.linalg import norm

WHITE = (255, 255, 255)
BLACK = (0, 0, 0)
MAX_BOUNCE = 3
IMG_SIZE = 800, 600
BG_COLOR = WHITE

class Sphere:
    def __init__(self, p, r, c):
        self.p = p
        self.r2 = r * r
        self.c = c

    def intersect(self, ray):
        o, d = ray
        L = subtract(o, self.p)
        a = dot(d, d)
        b = 2 * dot(d, L)
        c = dot(L, L) - self.r2

        discr = b * b - 4 * a * c
        if discr < 0:
            return False

        if discr == 0:
            x0 = x1 = -0.5 * b / a
            return x0
        
        q = -0.5 * (b + sqrt(discr)) if b > 0 else -0.5 * (b - sqrt(discr))
        x0 = q / a
        x1 = c / q

        if x0 > x1 : x0, x1 = x1, x0
        
        if x0 < 0 :
            x0 = x1
            if x0 < 0: return False

        return x0

    def normal(self, i):
        n = subtract(i, self.p)
        return n / norm(n)

    def phong(self, l, i, v):
        ambient = WHITE

        v = subtract(v, i)
        v = v / norm(v)

        n = subtract(i, self.p)
        n = n / norm(n)

        l = subtract(l, i)
        l = l / norm(l)

        h = add(l, v)
        h = h / norm(h)

        diffuse = multiply(self.c, dot(n, l))

        s = dot(n, h)**64
        specular = multiply(WHITE, s)

        return multiply(specular, 0.5) + multiply(diffuse, 0.5) # + multiply(ambient, 0.3)

def nearest_intersect(s, r):
    gmin, tmin = None, inf
    for g in scene:
        if g == s : continue
        t = g.intersect(r)
        # if not t is False and t > 0:
        if not t is False and t < tmin and t > 0:
            gmin, tmin = g, t
    return gmin, tmin

def reflect(s, r, m=0):
    o, d = r
    c = (127,127,127)
    g, t = nearest_intersect(s, r)

    if not g is None:
        i = add(o, multiply(d, t))

        n = g.normal(i)
        rv = subtract(d, multiply(n, 2.0 * dot(d, n)))

        c = g.phong(light, i, o)
        if m > 0:
            rc = reflect(g, (i, rv), m - 1)
            c = multiply(add(c, rc), 0.5)
        c = tuple(int(x) for x in c)
        # c = BLACK
    return c

def raytracer(size):
    w, h = size
    w2, h2 = w / 2, h / 2

    for x, y in product(range(w), range(h)):
        c = WHITE
        d = (x - w2, y - h2, h2)
        d = d / norm(d)
        r = camera, d
        g, t = nearest_intersect(None, r)
        if not g is None:
            i = add(camera, multiply(d, t))
            n = g.normal(i)
            rv = subtract(d, multiply(n, 2.0 * dot(d, n)))

            c = g.phong(light, i, camera)
            rc = reflect(g, (i, rv), MAX_BOUNCE)
            c = multiply(add(c, rc), 0.5)
            # c = rc
            c = tuple(int(x) for x in c)

        yield (x, y), c
    

light = (-300.0, -300.0, 000.0)

scene = [
    Sphere((100.0, 200.0, 300.0), 100, (100, 100, 200)),
    Sphere((-100.0, 200.0, 300.0), 100, (100, 200, 100)),
    Sphere((0.0, 0.0, 800.0), 400, (200, 100, 100)),
]

camera = (0, 0, 0)

im = Image.new("RGB", IMG_SIZE, BG_COLOR)

for p, c in raytracer(IMG_SIZE):
    im.putpixel(p, c)

im.show()