Depth of Field

[WoodNeck]

#define FLT_MAX 3.402823466e+38
precision mediump float;
uniform vec4 uScreenSize;
uniform float uTime;

const int NUM_OBJECTS = 5;

const int LAMBERTIAN = 0;
const int METAL = 1;
const int DIELECTRIC = 2;

struct ray {
    vec3 o, d;
};
struct record {
    float t;
    vec3 p, n;
};
struct material {
    int type;
    vec3 albedo;
};
struct hitable {
    vec3 c;
    float r;
    material mat;
};
struct camprop {
    vec3 lower_left_corner;
    vec3 horizontal;
    vec3 vertical;
    vec3 u;
    vec3 v;
    vec3 w;
    float lens_radius;
};
struct camera {
    vec3 o;
    vec3 d;
    vec3 up;
    camprop prop;
};

highp float rand(vec2 co) {
    highp float a = 12.9898;
    highp float b = 78.233;
    highp float c = 43758.5453;
    highp float dt= dot(co.xy ,vec2(a,b));
    highp float sn= mod(dt,3.14);
    return fract(sin(sn) * c);
}

float schlick(float cosine, float ref_idx) {
    float r0 = (1. - ref_idx) / (1. + ref_idx);
    r0 = r0 * r0;
    return r0 + (1. - r0) * pow((1. - cosine), 5.);
}

vec3 randomInUnitSphere(float sd) {
    return vec3(
        rand(vec2(sd += .1, sd += .1)),
        rand(vec2(sd += .1, sd += .1)),
        rand(vec2(sd += .1, sd += .1))
    );
}

vec3 randomInUnitDisk(float sd) {
    return 2.0 * vec3(rand(vec2(sd += .1, sd += .1)), rand(vec2(sd += .1, sd += .1)), 0) - vec3(1, 1, 0);
}

camprop getCameraProp(vec3 lookFrom, vec3 lookAt, vec3 up, float fov, float aspect, float aperture, float focusDist) {
    float half_height = tan(radians(fov * .5));
    float half_width = aspect * half_height;

    vec3 w = normalize(lookFrom - lookAt);
    vec3 u = normalize(cross(up, w));
    vec3 v = cross(w, u);
    vec3 lower_left_corner = lookFrom - half_width * focusDist * u - half_height * focusDist * v - focusDist * w;
    vec3 horizontal = 2. * half_width * focusDist * u;
    vec3 vertical = 2. * half_height * focusDist * v;
    return camprop(lower_left_corner, horizontal, vertical, u, v, w, aperture / 2.);
}

ray getRay(camera cam, vec2 coord, float sd) {
    vec3 d = cam.prop.lens_radius * randomInUnitDisk(sd);
    vec3 offset = cam.prop.u * d.x + cam.prop.v * d.y;
    vec3 rd = normalize(cam.prop.lower_left_corner
        + coord.x * cam.prop.horizontal
        + coord.y * cam.prop.vertical
        - cam.o - offset);
    return ray(cam.o + offset, rd);
}

vec3 rayDirection(float fov, vec2 size, vec2 fragCoord) {
    vec2 xy = fragCoord - size * .5;
    float z = size.y / tan(radians(fov) * .5);

    return normalize(vec3(xy, -z));
}

bool refract(inout ray r, in vec3 n, float ior, out vec3 refracted) {
    float dt = dot(r.d, n);
    float d = 1.0 - ior * ior * (1.0 - dt * dt);
    if (d > 0.) {
        refracted = ior * (r.d - n * dt) - n * sqrt(d);
        return true;
    }
    return false;
}

bool sphere(hitable sp, ray r, float tMin, inout record hr, out material mat) {
    vec3 oc = r.o - sp.c;
    float a = dot(r.d, r.d);
    float b = dot(oc, r.d);
    float c = dot(oc, oc) - sp.r * sp.r;

    float d = b * b - a * c;
    if (d < 0.) return false;

    float s = sqrt(d);

    float val = (-b - s) / a;
    if (val > tMin && val < hr.t) {
        hr.p = r.o + val * r.d;
        hr.n = (hr.p - sp.c) / sp.r;
        hr.t = val;
        mat = sp.mat;
        return true;
    }
    val = (-b + s) / a;
    if (val > tMin && val < hr.t) {
        hr.p = r.o + val * r.d;
        hr.n = (hr.p - sp.c) / sp.r;
        hr.t = val;
        mat = sp.mat;
        return true;
    }

    return false;
}

bool worldHit(ray r, float tMin, float tMax, hitable[NUM_OBJECTS] sp, out record rec, out material mat) {
    bool hit = false;
    rec.t = tMax;

    for (int i = 0; i < NUM_OBJECTS; i++) {
        hit = sphere(sp[i], r, tMin, rec, mat) || hit;
    }

    return hit;
}

bool scatLamb(inout ray r, record rec, float sd) {
    vec3 t = normalize(rec.n + randomInUnitSphere(sd));
    r.o = rec.p;
    r.d = t;
    return true;
}

bool scatMetal(inout ray r, record rec) {
    vec3 reflected = reflect(normalize(r.d), rec.n);
    r.o = rec.p;
    r.d = reflected;
    return (dot(reflected, rec.n) > 0.);
}

bool scatDielectric(inout ray r, record rec, int rayIdx) {
    float ri = 1.5;
    float ior = ri;
    vec3 reflected = reflect(normalize(r.d), rec.n);
    vec3 outward_normal = -rec.n;
    vec3 refracted;

    float reflect_prob;
    float cosine = dot(r.d, rec.n);
    cosine = sqrt(1. - ri*ri*(1.0-cosine*cosine));

    if (dot(r.d, rec.n) <= 0.) {
        outward_normal = rec.n;
        ior = 1. / ior;
        cosine = ri * -dot(r.d, rec.n);
    }

    if (refract(r, outward_normal, ior, refracted)) {
        reflect_prob = schlick(cosine, ri);
    } else {
        reflect_prob = 1.0;
    }

    if (rand(rec.p.xy + float(rayIdx)) < reflect_prob) {
        r.o = rec.p;
        r.d = reflected;
    } else {
        r.o = rec.p;
        r.d = refracted;
    }

    return true;
}

vec3 color(ray r, hitable[NUM_OBJECTS] sp, float sd, int rayIdx) {
    const int MAX_RECURSION = 50;
    float tMax = FLT_MAX;
    vec3 col = vec3(1);

    record hr;

    for (int i = 0; i < MAX_RECURSION; i++) {
        material mat;
        if (worldHit(r, 0.001, FLT_MAX, sp, hr, mat)) {
            bool scat = mat.type == LAMBERTIAN
                ? scatLamb(r, hr, sd)
                : mat.type == METAL
                    ? scatMetal(r, hr)
                    : scatDielectric(r, hr, rayIdx);
            if (scat) {
                col *= mat.albedo;
            } else {
                return vec3(0, 0, 0);
            }
        } else {
            float t = .5 * r.d.y + .5;
            col *= mix(vec3(1), vec3(.5, .7, 1.), t);
            return col;
        }
    }
    return col;
}

void main() {
    // set camera
    vec3 lookFrom = vec3(3, 3, 2);
    vec3 lookAt = vec3(0, 0, -1);
    vec3 up = vec3(0, 1, 0);
    float vFov = 60.0;
    float aspect = uScreenSize.x / uScreenSize.y;
    float dist_to_focus = length(lookFrom - lookAt);
    float aperture = 2.0;

    camprop prop = getCameraProp(lookFrom, lookAt, up, vFov, aspect, aperture, dist_to_focus);
    camera cam = camera(lookFrom, lookAt, up, prop);

    const int ns = 100;

    hitable spheres[NUM_OBJECTS];
    spheres[0] = hitable(vec3(0, 0, -1), .5, material(LAMBERTIAN, vec3(.1, .2, .5)));
    spheres[1] = hitable(vec3(0, -100.5, -1), 100., material(LAMBERTIAN, vec3(.8, .8, 0)));
    spheres[2] = hitable(vec3(1, 0, -1), .5, material(METAL, vec3(.8, .6, .2)));
    spheres[3] = hitable(vec3(-1, 0, -1), .5, material(DIELECTRIC, vec3(1., 1., 1.)));
    spheres[4] = hitable(vec3(-1, 0, -1), -.49, material(DIELECTRIC, vec3(1., 1., 1.)));

    float sd = rand(gl_FragCoord.xy);
    vec3 col = vec3(0, 0, 0);
    float ins = 1. / float(ns);
    for (int i = 0; i < ns; i++) {
        float rx = 2. * rand(gl_FragCoord.xx + float(i) * ins) - 1.;
        float ry = 2. * rand(gl_FragCoord.yy + float(i) * ins) - 1.;
        ray r = getRay(cam, (gl_FragCoord.xy + vec2(rx, ry)) / uScreenSize.xy, sd);
        col += color(r, spheres, sd, i);
        sd += 0.1;
        sd = fract(sd);
    }
    col = col * ins;
    col = vec3(sqrt(col[0]), sqrt(col[1]), sqrt(col[2]));
    gl_FragColor = vec4(col, 1);
}