def color_at(object: Object, position, normal, scene: Scene):
color = vector()
color = object.material.ambient * color
to_cam = scene.camera - position
specular_k = 50
for light in scene.lights:
to_light = Ray(position, light.position - position)
color += (
object.material.color
* object.material.diffuse
* max(torch.dot(normal, to_light.direction), 0)
)
half_vector = normalize(to_light.direction + to_cam)
color += (
light.color
* object.material.specular
* max(torch.dot(normal, half_vector), 0) ** specular_k
)
color = color.clamp(0, 1)
return color
def reflect(self, ray, t, normal):
reflected = reflect(ray.direction, normal)
reflected = Ray(ray.position_at(t), reflected)
if reflected.direction.dot(normal) <= 0:
return None
return Reflection(reflected, self.color)
def ray_trace(ray: Ray, objects: ObjectList, max_steps):
if max_steps == 0:
return vector(0, 0, 0)
intersection = objects.intersects(ray)
if intersection is None:
return vector(0.2, 0.2, 0.2)
position = ray.position_at(intersection.t)
normal = intersection.object.normal_at(position)
emitted = intersection.object.material.emit()
reflection = intersection.object.material.reflect(ray, intersection.t,
normal)
if reflection is None:
return emitted
return emitted + reflection.attenuation * ray_trace(
reflection.ray, objects, max_steps - 1)
def ray_trace(ray: Ray, scene: Scene):
color = vector()
ot = None
for object in scene.objects:
t = object.intersects(ray)
if t is None:
continue
if ot is None:
ot = object, t
if t < ot[1]:
ot = object, t
if ot is None:
return color
object, t = ot
del ot
position = ray.position_at(t)
normal = object.normal_at(position)
color += color_at(object, position, normal, scene)
return color
def main():
size = 256, 256
camera = vector(0, 0, -1)
objects = [
Sphere(vector(0, 0, 0), 0.5, Metal(vector(1, 0, 0))),
Sphere(vector(0, -1, 1), 0.5, Metal(vector(0, 1, 0))),
]
lights = [
Light(vector(1.5, -0.5, -10), vector(1, 1, 1)),
]
scene = Scene(camera=camera, objects=objects, lights=lights)
view = build_view(size)
image = torch.zeros(3, *size, dtype=torch.float)
for i, j in tqdm(product(range(size[0]), range(size[1])), total=size[0] * size[1]):
ray = Ray(camera, view[:, i, j] - camera)
image[:, i, j] = ray_trace(ray, scene)
image = to_pil_image(image)
image.save("./ray_tracing/output.png")
plt.imshow(image)
plt.show()
def reflect_hemi(self, ray, t, normal):
ray = Ray(ray.position_at(t), random_in_hemisphere(normal))
return Reflection(ray, self.color)
def reflect_sphere(self, ray, t, normal):
ray = Ray(ray.position_at(t), normal + random_unit())
return Reflection(ray, self.color)
def ray_to_position(self, x, y):
position = vector(2 * x - 1, 2 * y - 1, 0)
return Ray(self.origin, position - self.origin)