ai_pathtracer/raytracer.py

576 lines
20 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
from rich.console import Console
from rich.progress import Progress, SpinnerColumn, BarColumn, TextColumn, TimeElapsedColumn, TimeRemainingColumn
from rich.table import Table
from rich.panel import Panel
from matplotlib import pyplot as plt
console = Console()
console.print()
console.print(Panel.fit(
"[bold magenta]Cornell Box Raytracer[/bold magenta]\n"
"[dim]Reflections | Caustics | Multiple Bounces | NEE[/dim]",
border_style="bright_blue", padding=(0, 2)
))
settings = Table(title="Scene Configuration", show_header=False)
settings.add_column("Parameter", style="cyan", no_wrap=True)
settings.add_column("Value", style="white")
settings.add_row("Resolution", f"[bold]{width}x{height}[/bold]")
settings.add_row("Samples/pixel", f"[bold yellow]{samples_per_pixel}[/bold yellow]")
settings.add_row("Max bounces", f"[bold green]{max_depth}[/bold green]")
settings.add_row("Workers", f"[bold magenta]{workers}[/bold magenta]")
settings.add_row("Output", f"[bold]{output}[/bold]")
console.print(settings)
console.print()
progress = Progress(
SpinnerColumn(finished_text="[DONE]", style="green"),
TextColumn("[progress.description]{task.description}"),
BarColumn(bar_width=50, style="bright_blue", complete_style="bright_green"),
TextColumn("[progress.percentage]{task.percentage:>3.0f}%"),
TextColumn("[bold]{task.completed}/{task.total}[/bold] rows"),
TimeElapsedColumn(),
TimeRemainingColumn(),
)
brightness_history = []
timestamps = []
completed_rows = 0
total_rows = height
with progress:
render_task = progress.add_task("Rendering", total=total_rows)
preview_path = "preview.png"
with ProcessPoolExecutor(max_workers=workers) as executor:
futures = {executor.submit(render_chunk, r): r[9] for r in ranges}
for f in as_completed(futures):
chunk_start, chunk_img = f.result()
chunk_end = chunk_start + chunk_img.shape[0]
image[chunk_start:chunk_end] = chunk_img
completed_rows += (chunk_end - chunk_start)
progress.update(render_task, completed=completed_rows)
img = (np.clip(image, 0.0, 1.0) * 255).astype(np.uint8)
Image.fromarray(img, 'RGB').save(preview_path)
elapsed = time.time() - start
brightness = float(image.mean())
brightness_history.append(brightness)
timestamps.append(elapsed)
elapsed = time.time() - start
img = (np.clip(image, 0.0, 1.0) * 255).astype(np.uint8)
Image.fromarray(img, 'RGB').save(output)
console.print()
console.print(Panel.fit(
f"[bold green]Rendered {output}[/bold green]\n"
f"[dim]Total time: {elapsed:.1f}s | "
f"Speed: {width * height * samples_per_pixel / elapsed / 1e6:.2f} MPix*samp/s[/dim]",
border_style="bright_green", padding=(0, 2)
))
region_table = Table(title="Region Brightness Analysis")
region_table.add_column("Region", style="cyan", no_wrap=True)
region_table.add_column("Brightness", justify="right", style="white")
region_table.add_column("Visual", justify="center")
h_start, h_end = height // 3, 2 * height // 3
regions = [
("Left Wall", image[:, :width // 4]),
("Right Wall", image[:, 3 * width // 4:]),
("Floor", image[2 * height // 3:, h_start:h_end]),
("Ceiling", image[:height // 3, h_start:h_end]),
("Center", image[h_start:h_end, h_start:h_end]),
]
for name, region in regions:
brightness = float(region.mean()) * 255
bar_len = 20
filled = int(brightness / 255 * bar_len)
bar = "#" * filled + "-" * (bar_len - filled)
color = "green" if brightness > 80 else ("yellow" if brightness > 40 else "red")
region_table.add_row(name, f"{brightness:.1f}", f"[{color}]{bar}[/{color}]")
console.print(region_table)
console.print()
fig, axes = plt.subplots(1, 2, figsize=(14, 5))
ax1, ax2 = axes
t_arr = np.array(timestamps)
b_arr = np.array(brightness_history)
ax1.plot(t_arr, b_arr, color='#22cc66', linewidth=2)
ax1.set_xlabel("Time (s)", fontsize=11)
ax1.set_ylabel("Mean Brightness", fontsize=11)
ax1.set_title("Brightness Convergence", fontsize=13, fontweight='bold')
ax1.grid(True, alpha=0.3)
ax1.fill_between(t_arr, b_arr, alpha=0.15, color='#22cc66')
preview = (np.clip(image, 0.0, 1.0) * 255).astype(np.uint8)
ax2.imshow(preview)
ax2.set_title("Render Preview", fontsize=13, fontweight='bold')
ax2.axis('off')
plt.tight_layout()
metrics_path = "render_metrics.png"
plt.savefig(metrics_path, dpi=150, bbox_inches='tight')
plt.close()
console.print(f"[dim]Metrics graph saved to [bold]{metrics_path}[/bold][/dim]")
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)