generated from taichiCourse01/taichi_course_homework
-
Notifications
You must be signed in to change notification settings - Fork 0
/
movable_ray_tracer.py
170 lines (159 loc) · 6.84 KB
/
movable_ray_tracer.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
import taichi as ti
import numpy as np
import argparse
from ray_tracing_models import Ray, Camera, Hittable_list, Sphere, PI, random_in_unit_sphere, refract, reflect, reflectance, random_unit_vector
ti.init(arch=ti.gpu)
diff = ti.Vector.field(3, ti.f32, ())
# Canvas
aspect_ratio = 1.0
image_width = 800
image_height = int(image_width / aspect_ratio)
canvas = ti.Vector.field(3, dtype=ti.f32, shape=(image_width, image_height))
# Rendering parameters
samples_per_pixel = 4
max_depth = 10
sample_on_unit_sphere_surface = True
@ti.kernel
def black_out():
for i, j in canvas:
canvas[i, j] = ti.Vector([0.0, 0.0, 0.0])
@ti.kernel
def render():
for i, j in canvas:
u = (i + ti.random()) / image_width
v = (j + ti.random()) / image_height
color = ti.Vector([0.0, 0.0, 0.0])
for n in range(samples_per_pixel):
ray = camera.get_ray(u, v)
color += ray_color(ray)
color /= samples_per_pixel
canvas[i, j] += color
# Path tracing
@ti.func
def ray_color(ray):
color_buffer = ti.Vector([0.0, 0.0, 0.0])
brightness = ti.Vector([1.0, 1.0, 1.0])
scattered_origin = ray.origin
scattered_direction = ray.direction
p_RR = 0.8
for n in range(max_depth):
if ti.random() > p_RR:
break
is_hit, hit_point, hit_point_normal, front_face, material, color = scene.hit(Ray(scattered_origin, scattered_direction))
if is_hit:
if material == 0:
color_buffer = color * brightness
break
else:
# Diffuse
if material == 1:
target = hit_point + hit_point_normal
if sample_on_unit_sphere_surface:
target += random_unit_vector()
else:
target += random_in_unit_sphere()
scattered_direction = target - hit_point
scattered_origin = hit_point
brightness *= color
# Metal and Fuzz Metal
elif material == 2 or material == 4:
fuzz = 0.0
if material == 4:
fuzz = 0.4
scattered_direction = reflect(scattered_direction.normalized(),
hit_point_normal)
if sample_on_unit_sphere_surface:
scattered_direction += fuzz * random_unit_vector()
else:
scattered_direction += fuzz * random_in_unit_sphere()
scattered_origin = hit_point
if scattered_direction.dot(hit_point_normal) < 0:
break
else:
brightness *= color
# Dielectric
elif material == 3:
refraction_ratio = 1.5
if front_face:
refraction_ratio = 1 / refraction_ratio
cos_theta = min(-scattered_direction.normalized().dot(hit_point_normal), 1.0)
sin_theta = ti.sqrt(1 - cos_theta * cos_theta)
# total internal reflection
if refraction_ratio * sin_theta > 1.0 or reflectance(cos_theta, refraction_ratio) > ti.random():
scattered_direction = reflect(scattered_direction.normalized(), hit_point_normal)
else:
scattered_direction = refract(scattered_direction.normalized(), hit_point_normal, refraction_ratio)
scattered_origin = hit_point
brightness *= color
brightness /= p_RR
return color_buffer
if __name__ == "__main__":
parser = argparse.ArgumentParser(description='Naive Ray Tracing')
parser.add_argument(
'--max_depth', type=int, default=10, help='max depth (default: 10)')
parser.add_argument(
'--samples_per_pixel', type=int, default=4, help='samples_per_pixel (default: 4)')
parser.add_argument(
'--samples_in_unit_sphere', action='store_true', help='whether sample in a unit sphere')
args = parser.parse_args()
max_depth = args.max_depth
samples_per_pixel = args.samples_per_pixel
sample_on_unit_sphere_surface = not args.samples_in_unit_sphere
scene = Hittable_list()
# Light source
scene.add(Sphere(center=ti.Vector([0, 5.4, -1]), radius=3.0, material=0, color=ti.Vector([10.0, 10.0, 10.0])))
# Ground
scene.add(Sphere(center=ti.Vector([0, -100.5, -1]), radius=100.0, material=1, color=ti.Vector([0.8, 0.8, 0.8])))
# ceiling
scene.add(Sphere(center=ti.Vector([0, 102.5, -1]), radius=100.0, material=1, color=ti.Vector([0.8, 0.8, 0.8])))
# back wall
scene.add(Sphere(center=ti.Vector([0, 1, 101]), radius=100.0, material=1, color=ti.Vector([0.8, 0.8, 0.8])))
# right wall
scene.add(Sphere(center=ti.Vector([-101.5, 0, -1]), radius=100.0, material=1, color=ti.Vector([0.6, 0.0, 0.0])))
# left wall
scene.add(Sphere(center=ti.Vector([101.5, 0, -1]), radius=100.0, material=1, color=ti.Vector([0.0, 0.6, 0.0])))
# Diffuse ball
scene.add(Sphere(center=ti.Vector([0, -0.2, -1.5]), radius=0.3, material=1, color=ti.Vector([0.8, 0.3, 0.3])))
# Metal ball
scene.add(Sphere(center=ti.Vector([-0.8, 0.2, -1]), radius=0.7, material=2, color=ti.Vector([0.6, 0.8, 0.8])))
# Glass ball
scene.add(Sphere(center=ti.Vector([0.7, 0, -0.5]), radius=0.5, material=3, color=ti.Vector([1.0, 1.0, 1.0])))
# Metal ball-2
scene.add(Sphere(center=ti.Vector([0.6, -0.3, -2.0]), radius=0.2, material=4, color=ti.Vector([0.8, 0.6, 0.2])))
camera = Camera()
gui = ti.GUI("Ray Tracing", res=(image_width, image_height))
canvas.fill(0)
cnt = 0
while gui.running:
flag = False
for e in gui.get_events(ti.GUI.PRESS, ti.GUI.MOTION):
if e.key == ti.GUI.LEFT:
camera.set_camera(0.2, 0, 0)
flag = True
elif e.key == ti.GUI.RIGHT:
camera.set_camera(-0.2, 0, 0)
flag = True
elif e.key == ti.GUI.UP:
camera.set_camera(0, 0.2, 0)
flag = True
elif e.key == ti.GUI.DOWN:
camera.set_camera(0, -0.2, 0)
flag = True
elif e.key == ti.GUI.WHEEL:
# scroll: zoom
if e.delta[1] > 0:
camera.set_camera(0, 0, 0.2)
flag = True
elif e.delta[1] < 0:
camera.set_camera(0, 0, -0.2)
flag = True
elif e.key == 'r' or e.key == 'R':
flag = True
if flag:
black_out()
cnt = 0
camera.reset()
render()
cnt += 1
gui.set_image(np.sqrt(canvas.to_numpy() / cnt))
gui.show()