469 lines
16 KiB
Python
469 lines
16 KiB
Python
import numpy as np
|
|
from PIL import Image
|
|
import math
|
|
import time
|
|
|
|
EPSILON = 1e-6
|
|
PI = math.pi
|
|
|
|
|
|
def normalize(v):
|
|
n = np.linalg.norm(v)
|
|
return v / n if n > 1e-12 else v
|
|
|
|
|
|
def reflect(d, n):
|
|
return d - 2.0 * np.dot(d, n) * n
|
|
|
|
|
|
def refract(incident, normal, eta):
|
|
cos_i = min(-np.dot(incident, normal), 1.0)
|
|
sin_t_sq = eta * eta * (1.0 - cos_i * cos_i)
|
|
if sin_t_sq > 1.0:
|
|
return None
|
|
cos_t = math.sqrt(max(0.0, 1.0 - sin_t_sq))
|
|
return eta * incident + (eta * cos_i - cos_t) * normal
|
|
|
|
|
|
def schlick(cos, ref_idx):
|
|
r0 = (1.0 - ref_idx) / (1.0 + ref_idx)
|
|
r0 *= r0
|
|
return r0 + (1.0 - r0) * math.pow(1.0 - cos, 5.0)
|
|
|
|
|
|
class Ray:
|
|
__slots__ = ('origin', 'direction')
|
|
|
|
def __init__(self, origin, direction):
|
|
self.origin = np.asarray(origin, dtype=np.float64)
|
|
self.direction = normalize(np.asarray(direction, dtype=np.float64))
|
|
|
|
|
|
class HitRecord:
|
|
__slots__ = ('t', 'point', 'normal', 'front_face', 'material')
|
|
|
|
def __init__(self):
|
|
self.t = float('inf')
|
|
self.point = np.zeros(3)
|
|
self.normal = np.zeros(3)
|
|
self.front_face = True
|
|
self.material = None
|
|
|
|
|
|
class Hittable:
|
|
def hit(self, r, t_min=EPSILON, t_max=float('inf')):
|
|
raise NotImplementedError
|
|
|
|
|
|
class Sphere(Hittable):
|
|
__slots__ = ('center', 'radius', 'material')
|
|
|
|
def __init__(self, center, radius, material):
|
|
self.center = np.asarray(center, dtype=np.float64)
|
|
self.radius = float(radius)
|
|
self.material = material
|
|
|
|
def hit(self, r, t_min=EPSILON, t_max=float('inf')):
|
|
oc = r.origin - self.center
|
|
a = np.dot(r.direction, r.direction)
|
|
b = np.dot(oc, r.direction)
|
|
c = np.dot(oc, oc) - self.radius * self.radius
|
|
disc = b * b - a * c
|
|
if disc < 0:
|
|
return None
|
|
sq = math.sqrt(disc)
|
|
root = (-b - sq) / a
|
|
if root < t_min or root > t_max:
|
|
root = (-b + sq) / a
|
|
if root < t_min or root > t_max:
|
|
return None
|
|
rec = HitRecord()
|
|
rec.t = root
|
|
rec.point = r.origin + root * r.direction
|
|
out = normalize(rec.point - self.center)
|
|
if np.dot(r.direction, out) < 0:
|
|
rec.normal = out
|
|
rec.front_face = True
|
|
else:
|
|
rec.normal = -out
|
|
rec.front_face = False
|
|
rec.material = self.material
|
|
return rec
|
|
|
|
|
|
class Plane(Hittable):
|
|
__slots__ = ('point', 'normal', 'material')
|
|
|
|
def __init__(self, point, normal, material):
|
|
self.point = np.asarray(point, dtype=np.float64)
|
|
self.normal = normalize(np.asarray(normal, dtype=np.float64))
|
|
self.material = material
|
|
|
|
def hit(self, r, t_min=EPSILON, t_max=float('inf')):
|
|
denom = np.dot(r.direction, self.normal)
|
|
if abs(denom) < 1e-12:
|
|
return None
|
|
t = np.dot(self.point - r.origin, self.normal) / denom
|
|
if t < t_min or t > t_max:
|
|
return None
|
|
rec = HitRecord()
|
|
rec.t = t
|
|
rec.point = r.origin + t * r.direction
|
|
if np.dot(r.direction, self.normal) < 0:
|
|
rec.normal = self.normal
|
|
rec.front_face = True
|
|
else:
|
|
rec.normal = -self.normal
|
|
rec.front_face = False
|
|
rec.material = self.material
|
|
return rec
|
|
|
|
|
|
class Box(Hittable):
|
|
__slots__ = ('min_pt', 'max_pt', 'material')
|
|
|
|
def __init__(self, min_pt, max_pt, material):
|
|
self.min_pt = np.asarray(min_pt, dtype=np.float64)
|
|
self.max_pt = np.asarray(max_pt, dtype=np.float64)
|
|
self.material = material
|
|
|
|
def hit(self, r, t_min=EPSILON, t_max=float('inf')):
|
|
t_lo = t_min
|
|
t_hi = t_max
|
|
hit_axis = -1
|
|
inv_ds = [0.0, 0.0, 0.0]
|
|
for axis in range(3):
|
|
inv_d = 1.0 / r.direction[axis] if abs(r.direction[axis]) > 1e-12 else 1e12
|
|
inv_ds[axis] = inv_d
|
|
t0 = (self.min_pt[axis] - r.origin[axis]) * inv_d
|
|
t1 = (self.max_pt[axis] - r.origin[axis]) * inv_d
|
|
if inv_d < 0:
|
|
t0, t1 = t1, t0
|
|
if t0 > t_lo:
|
|
t_lo = t0
|
|
hit_axis = axis
|
|
t_hi = min(t_hi, t1)
|
|
if t_hi < t_lo:
|
|
return None
|
|
if t_lo >= t_max or t_lo <= t_min:
|
|
return None
|
|
rec = HitRecord()
|
|
n = np.zeros(3)
|
|
n[hit_axis] = 1.0 if inv_ds[hit_axis] > 0 else -1.0
|
|
if np.dot(r.direction, n) < 0:
|
|
rec.normal = n
|
|
rec.front_face = True
|
|
else:
|
|
rec.normal = -n
|
|
rec.front_face = False
|
|
rec.t = t_lo
|
|
rec.point = r.origin + t_lo * r.direction
|
|
rec.material = self.material
|
|
return rec
|
|
|
|
|
|
class Scene:
|
|
__slots__ = ('objects', 'lights')
|
|
|
|
def __init__(self):
|
|
self.objects = []
|
|
self.lights = []
|
|
|
|
def add(self, obj):
|
|
self.objects.append(obj)
|
|
|
|
def add_light(self, obj):
|
|
self.lights.append(obj)
|
|
self.objects.append(obj)
|
|
|
|
def hit(self, r, t_min=EPSILON, t_max=float('inf')):
|
|
closest = None
|
|
min_t = t_max
|
|
for obj in self.objects:
|
|
rec = obj.hit(r, t_min, min_t)
|
|
if rec is not None and rec.t < min_t:
|
|
closest = rec
|
|
min_t = rec.t
|
|
return closest
|
|
|
|
def blocked(self, origin, direction, t_max):
|
|
r = Ray(origin, direction)
|
|
for obj in self.objects:
|
|
if obj.hit(r, EPSILON, t_max) is not None:
|
|
return True
|
|
return False
|
|
|
|
|
|
class Material:
|
|
def scatter(self, ray_in, rec, rng):
|
|
raise NotImplementedError
|
|
|
|
def get_emission(self):
|
|
return np.zeros(3)
|
|
|
|
|
|
class Diffuse(Material):
|
|
__slots__ = ('albedo',)
|
|
|
|
def __init__(self, albedo):
|
|
self.albedo = np.asarray(albedo, dtype=np.float64)
|
|
|
|
def scatter(self, ray_in, rec, rng):
|
|
target = rec.normal + np.random.randn(3) * 0.5
|
|
if np.linalg.norm(target) < 1e-8:
|
|
target = rec.normal.copy()
|
|
return True, self.albedo, Ray(rec.point, normalize(target))
|
|
|
|
|
|
class Mirror(Material):
|
|
__slots__ = ('albedo',)
|
|
|
|
def __init__(self, albedo):
|
|
self.albedo = np.asarray(albedo, dtype=np.float64)
|
|
|
|
def scatter(self, ray_in, rec, rng):
|
|
reflected = reflect(ray_in.direction, rec.normal)
|
|
if np.dot(reflected, rec.normal) <= 0:
|
|
return False, np.zeros(3), None
|
|
return True, self.albedo, Ray(rec.point, normalize(reflected))
|
|
|
|
|
|
class Dielectric(Material):
|
|
__slots__ = ('ref_idx', 'albedo')
|
|
|
|
def __init__(self, ref_idx, albedo=None):
|
|
self.ref_idx = ref_idx
|
|
self.albedo = np.asarray(albedo, dtype=np.float64) if albedo is not None else np.ones(3)
|
|
|
|
def scatter(self, ray_in, rec, rng):
|
|
refraction_ratio = 1.0 / self.ref_idx if rec.front_face else self.ref_idx
|
|
unit_direction = ray_in.direction.copy()
|
|
cos_theta = min(-np.dot(unit_direction, rec.normal), 1.0)
|
|
sin_theta = math.sqrt(max(0.0, 1.0 - cos_theta * cos_theta))
|
|
cannot_refract = refraction_ratio * sin_theta > 1.0
|
|
if cannot_refract or schlick(cos_theta, refraction_ratio) > np.random.random():
|
|
direction = reflect(unit_direction, rec.normal)
|
|
else:
|
|
direction = refract(unit_direction, rec.normal, refraction_ratio)
|
|
if direction is None:
|
|
direction = reflect(unit_direction, rec.normal)
|
|
return True, self.albedo, Ray(rec.point, normalize(direction))
|
|
|
|
|
|
class LightSource(Material):
|
|
__slots__ = ('emission',)
|
|
|
|
def __init__(self, emission):
|
|
self.emission = np.asarray(emission, dtype=np.float64)
|
|
|
|
def get_emission(self):
|
|
return self.emission
|
|
|
|
def scatter(self, ray_in, rec, rng):
|
|
return False, np.zeros(3), None
|
|
|
|
|
|
class Camera:
|
|
def __init__(self, lookfrom, lookat, vup, vfov, aspect, aperture=0.0, focus_dist=10.0):
|
|
self.origin = np.asarray(lookfrom, dtype=np.float64)
|
|
half_height = math.tan(math.radians(vfov) / 2.0)
|
|
half_width = aspect * half_height
|
|
w = normalize(self.origin - np.asarray(lookat, dtype=np.float64))
|
|
u = normalize(np.cross(vup, w))
|
|
v = np.cross(w, u)
|
|
self.horizontal = focus_dist * half_width * u
|
|
self.vertical = focus_dist * half_height * v
|
|
self.lower_left_corner = self.origin - self.horizontal / 2.0 - self.vertical / 2.0 - focus_dist * w
|
|
self.aperture = aperture
|
|
|
|
def get_ray(self, s, t, rng):
|
|
rd = np.zeros(3)
|
|
if self.aperture > 0:
|
|
rd = (np.random.rand(3) - 0.5) * self.aperture
|
|
rd[2] = 0.0
|
|
direction = self.lower_left_corner + s * self.horizontal + t * self.vertical - self.origin + rd
|
|
return Ray(self.origin, direction)
|
|
|
|
|
|
def cornell_box():
|
|
scene = Scene()
|
|
|
|
white = Diffuse(np.array([0.73, 0.73, 0.73]))
|
|
gray = Diffuse(np.array([0.5, 0.5, 0.5]))
|
|
red = Diffuse(np.array([0.65, 0.05, 0.05]))
|
|
green = Diffuse(np.array([0.12, 0.45, 0.15]))
|
|
mirror = Mirror(np.array([0.95, 0.95, 0.95]))
|
|
glass = Dielectric(ref_idx=1.5)
|
|
light = LightSource(np.array([50.0, 50.0, 50.0]))
|
|
|
|
scene.add(Plane(np.array([0, 0, 0]), np.array([0, 1, 0]), gray))
|
|
scene.add(Plane(np.array([0, 500, 0]), np.array([0, -1, 0]), white))
|
|
scene.add(Plane(np.array([0, 0, 500]), np.array([0, 0, -1]), white))
|
|
scene.add(Plane(np.array([0, 0, 0]), np.array([1, 0, 0]), red))
|
|
scene.add(Plane(np.array([500, 0, 0]), np.array([-1, 0, 0]), green))
|
|
|
|
scene.add(Box(np.array([0, 180, 180]), np.array([25, 220, 320]), light))
|
|
|
|
scene.add(Sphere(np.array([250, 180, 250]), 80, glass))
|
|
scene.add(Sphere(np.array([150, 180, 380]), 80, mirror))
|
|
scene.add(Sphere(np.array([350, 160, 300]), 60, red))
|
|
|
|
return scene
|
|
|
|
|
|
def trace(scene, ray, depth, max_depth, rng):
|
|
if depth >= max_depth:
|
|
return np.zeros(3)
|
|
|
|
rec = scene.hit(ray)
|
|
if rec is None:
|
|
t_val = 0.5 * (ray.direction[1] + 1.0)
|
|
return np.array([1.0, 1.0, 1.0]) * (1.0 - t_val) + np.array([0.5, 0.7, 1.0]) * t_val
|
|
|
|
emitted = rec.material.get_emission()
|
|
if isinstance(rec.material, LightSource):
|
|
return emitted
|
|
|
|
mat = rec.material
|
|
scattered, attenuation, scattered_ray = mat.scatter(ray, rec, rng)
|
|
if not scattered or scattered_ray is None:
|
|
return emitted
|
|
|
|
indirect = trace(scene, scattered_ray, depth + 1, max_depth, rng)
|
|
color = emitted + attenuation * indirect
|
|
|
|
if isinstance(mat, Diffuse):
|
|
direct = np.zeros(3)
|
|
for light_obj in scene.lights:
|
|
emission = light_obj.material.get_emission()
|
|
for _ in range(2):
|
|
face_idx = int(rng.random() * 6)
|
|
x = rng.uniform(light_obj.min_pt[0], light_obj.max_pt[0])
|
|
y = rng.uniform(light_obj.min_pt[1], light_obj.max_pt[1])
|
|
z = rng.uniform(light_obj.min_pt[2], light_obj.max_pt[2])
|
|
if face_idx == 0:
|
|
lp = np.array([light_obj.min_pt[0], y, z])
|
|
elif face_idx == 1:
|
|
lp = np.array([light_obj.max_pt[0], y, z])
|
|
elif face_idx == 2:
|
|
lp = np.array([x, light_obj.min_pt[1], z])
|
|
elif face_idx == 3:
|
|
lp = np.array([x, light_obj.max_pt[1], z])
|
|
elif face_idx == 4:
|
|
lp = np.array([x, y, light_obj.min_pt[2]])
|
|
else:
|
|
lp = np.array([x, y, light_obj.max_pt[2]])
|
|
|
|
to_light = lp - rec.point
|
|
dist_sq = np.dot(to_light, to_light)
|
|
if dist_sq < 1e-8:
|
|
continue
|
|
dist = math.sqrt(dist_sq)
|
|
d_to_light = to_light / dist
|
|
cos_at_rec = np.dot(rec.normal, d_to_light)
|
|
if cos_at_rec <= 0:
|
|
continue
|
|
|
|
if scene.blocked(rec.point + EPSILON * rec.normal, d_to_light, dist - EPSILON):
|
|
continue
|
|
|
|
light_area = 6 * (light_obj.max_pt[0] - light_obj.min_pt[0]) * \
|
|
(light_obj.max_pt[1] - light_obj.min_pt[1]) * \
|
|
(light_obj.max_pt[2] - light_obj.min_pt[2])
|
|
pdf = light_area / (dist_sq * cos_at_rec)
|
|
if pdf < 1e-12:
|
|
continue
|
|
brdf = mat.albedo / PI
|
|
direct += emission * cos_at_rec * brdf / pdf
|
|
break
|
|
|
|
if np.dot(direct, direct) > 0:
|
|
color = emitted + attenuation * (indirect + direct)
|
|
|
|
return color
|
|
|
|
|
|
def render_chunk(args):
|
|
scene, width, height, spp, max_d, \
|
|
lookfrom, lookat, vfov, chunk_start, chunk_end, seed_off = args
|
|
|
|
cam = Camera(lookfrom, lookat, np.array([0, 1, 0]), vfov, width / height, 0.0, 1078.0)
|
|
|
|
image = np.zeros((chunk_end - chunk_start, width, 3), dtype=np.float64)
|
|
|
|
for j_local, j in enumerate(range(chunk_start, chunk_end)):
|
|
for i in range(width):
|
|
color = np.zeros(3)
|
|
rng = np.random.default_rng(seed_off + i * 1000 + j * 7)
|
|
for _ in range(spp):
|
|
u = (i + rng.random()) / (width - 1)
|
|
v = (height - 1 - j + rng.random()) / (height - 1)
|
|
ray = cam.get_ray(u, v, rng)
|
|
color += trace(scene, ray, 0, max_d, rng)
|
|
color /= spp
|
|
color = np.sqrt(np.clip(color, 0, None))
|
|
image[j_local, i] = np.clip(color, 0.0, 1.0)
|
|
|
|
return (chunk_start, image)
|
|
|
|
|
|
def render(scene, width=800, height=600, samples_per_pixel=50, max_depth=15,
|
|
lookfrom=None, lookat=None, vfov=40, aperture=0.0, focus_dist=1078.0,
|
|
output="cornell_box.png", workers=16):
|
|
if lookfrom is None:
|
|
lookfrom = np.array([278.0, 278.0, -800.0])
|
|
if lookat is None:
|
|
lookat = np.array([278.0, 278.0, 278.0])
|
|
|
|
image = np.zeros((height, width, 3), dtype=np.float64)
|
|
start = time.time()
|
|
|
|
chunk = max(1, height // workers)
|
|
ranges = []
|
|
for w in range(workers):
|
|
s = w * chunk
|
|
e = height if w == workers - 1 else (w + 1) * chunk
|
|
ranges.append((scene, width, height, samples_per_pixel,
|
|
max_depth, lookfrom, lookat, vfov, s, e, w * 100000 + 42))
|
|
|
|
from concurrent.futures import ProcessPoolExecutor, as_completed
|
|
with ProcessPoolExecutor(max_workers=workers) as executor:
|
|
futures = {executor.submit(render_chunk, r): r[9] for r in ranges}
|
|
completed = 0
|
|
for f in as_completed(futures):
|
|
chunk_start, chunk_img = f.result()
|
|
end = chunk_start + chunk_img.shape[0]
|
|
image[chunk_start:end] = chunk_img
|
|
completed += 1
|
|
print(f"Chunk {completed}/{workers} done ({time.time() - start:.1f}s)")
|
|
|
|
img = (image * 255).astype(np.uint8)
|
|
Image.fromarray(img, 'RGB').save(output)
|
|
elapsed = time.time() - start
|
|
print(f"Saved {output} in {elapsed:.1f}s")
|
|
return image
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
print("=== Cornell Box Raytracer ===")
|
|
print("Features: reflections, refraction/caustics, multiple bounces, NEE")
|
|
print()
|
|
|
|
scene = cornell_box()
|
|
|
|
width = int(sys.argv[1]) if len(sys.argv) > 1 else 400
|
|
height = int(sys.argv[2]) if len(sys.argv) > 2 else 300
|
|
spp = int(sys.argv[3]) if len(sys.argv) > 3 else 50
|
|
max_depth = int(sys.argv[4]) if len(sys.argv) > 4 else 15
|
|
output = sys.argv[5] if len(sys.argv) > 5 else "cornell_box.png"
|
|
|
|
print(f"Resolution: {width}x{height}")
|
|
print(f"Samples per pixel: {spp}")
|
|
print(f"Max bounces: {max_depth}")
|
|
print(f"Output: {output}")
|
|
print()
|
|
|
|
render(scene, width=width, height=height,
|
|
samples_per_pixel=spp, max_depth=max_depth, output=output)
|