def main(size, k, steps, output_path):
# utils.random_seed(2**7)
utils.random_seed(2**10)
size = size, size
camera = Camera(vector(0, 0, -1))
objects = ObjectList(randomize_objects())
with Pool(os.cpu_count()) as pool:
image = pool.imap(
partial(
render_row,
size=size,
camera=camera,
objects=objects,
k=k,
max_steps=steps,
),
range(size[0]),
)
image = list(tqdm(image, total=size[0]))
image = torch.stack(image, 1)
image = image.clamp(0, 1)
image = image.flip(1)
image = to_pil_image(image)
os.makedirs(output_path, exist_ok=True)
image.save(
os.path.join(output_path, "{}_{}_{}.png".format(size[0], k, steps)))
plt.imshow(image)
plt.show()
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 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 randomize_objects():
floor = Sphere(vector(0, -102, 1), 100,
Diffuse(vector(1 / 3, 1 / 3, 1 / 3)))
objects = [floor]
for _ in range(8):
floor_to_center = vector(0, -2, 1) - floor.center
floor_to_center[0] += random.uniform(-4, 4)
floor_to_center[2] += random.uniform(0, 8)
radius = random.uniform(0.25, 1.5)
center = floor.center + normalize(
floor_to_center) * floor.radius + radius
mat = random.choice([Diffuse, Metal])
object = Sphere(center, radius, mat(vector(0, 0, 0).uniform_(0, 1)))
objects.append(object)
objects[8].material = Light(vector(1, 1, 1))
return objects
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 emit(self):
return vector(0, 0, 0)
def ray_to_position(self, x, y):
position = vector(2 * x - 1, 2 * y - 1, 0)
return Ray(self.origin, position - self.origin)