├── result.png ├── skybox.jpg ├── README.md └── pt.py /result.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Danielmelody/taichi_pt/HEAD/result.png -------------------------------------------------------------------------------- /skybox.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Danielmelody/taichi_pt/HEAD/skybox.jpg -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Taichi Path tracer 2 | 3 | ![](result.png) 4 | 5 | A progressive path tracer written in [taichi](https://github.com/taichi-dev/taichi) 6 | 7 | #### Features: 8 | * Global illumination via unbiased Monte Carlo path tracing 9 | * Physically based Specular shading(GGX) 10 | * Lambert diffuse shading 11 | * Ray-Sphere intersection 12 | * Unbiasd russain roule 13 | * Antialiasing via super-sampling 14 | * Depth of field effect of lens camera 15 | * Multiple Importance Sampling 16 | * Balance heuristic 17 | * Cosine-weighted pdf 18 | * ggx normal weighted pdf 19 | * ACES Tone mapping 20 | * Bloom effect 21 | * Gamma correction of final result 22 | 23 | The 5 balls in the scene are: 24 | 1. rough golden ball 25 | 2. smooth ceramics ball 26 | 3. the light source ball 27 | 4. huge ground rough iron ball 28 | 5. smooth metal ball 29 | #### Usage 30 | 31 | ```bash 32 | pip3 install taichi 33 | python3 pt.py 34 | ``` 35 | 36 | and you a ready to go. 37 | 38 | Scroll mouse to zoom in/out, press and drag the mouse to see a interactive result. 39 | You can also press shift while scrolling mouse to change the focal length of the camera. 40 | -------------------------------------------------------------------------------- /pt.py: -------------------------------------------------------------------------------- 1 | import taichi as ti 2 | import numpy as np 3 | import math 4 | 5 | ti.init(arch=ti.gpu) 6 | 7 | inf = 1e10 8 | eps = 1e-4 9 | 10 | max_luminance = 100 11 | 12 | class Image: 13 | def __init__(self, path): 14 | self.img = ti.tools.imread(path) 15 | tex_w, tex_h = self.img.shape[0:2] 16 | self.field = ti.Vector.field(3, dtype=ti.u8, shape=(tex_w, tex_h)) 17 | 18 | def load(self): 19 | self.field.from_numpy(self.img) 20 | 21 | 22 | @ti.func 23 | def tex2d(tex_field, uv): 24 | tex_w, tex_h = tex_field.shape[0:2] 25 | p = uv * ti.Vector([tex_w, tex_h]) 26 | l = int(ti.floor(p)) 27 | t = p - l 28 | 29 | # bilinear interp 30 | return ( 31 | ((tex_field[l] * (1 - t[0])) + tex_field[l + ti.Vector([1, 0])] * t[0]) * (1 - t[1]) 32 | + (tex_field[l + ti.Vector([0, 1])] * (1 - t[0]) + tex_field[l + ti.Vector([1, 1])] * t[0]) * t[1]) / 255.0 33 | 34 | def gaussian_kernel(l=5, sig=1.): 35 | """\ 36 | creates gaussian kernel with side length `l` and a sigma of `sig` 37 | """ 38 | ax = np.linspace(-(l - 1) / 2., (l - 1) / 2., l, dtype=np.single) 39 | gauss = np.exp(-0.5 * np.square(ax) / np.square(sig), dtype=np.single) 40 | kernel = np.outer(gauss, gauss) 41 | return kernel / np.sum(kernel) 42 | 43 | 44 | @ti.data_oriented 45 | class SphereSOA: 46 | def __init__(self, array_length): 47 | self.center_radius = ti.Vector.field(4, dtype=ti.f32, shape=(array_length)) 48 | self.albedos = ti.Vector.field(3, dtype=ti.f32, shape=(array_length)) 49 | self.emissions = ti.Vector.field(3, dtype=ti.f32, shape=(array_length)) 50 | self.roughness = ti.field(dtype=ti.f32, shape=(array_length)) 51 | self.metallics = ti.field(dtype=ti.f32, shape=(array_length)) 52 | self.iors = ti.field(dtype=ti.f32, shape=(array_length)) 53 | self.array_length = array_length 54 | 55 | @ ti.func 56 | def intersect(self, o, d): 57 | min_t = inf 58 | min_index = -1 59 | sp_index = 0 60 | for i in ti.ndrange(self.array_length): 61 | c_r = self.center_radius[i] 62 | r = c_r[3] 63 | p = ti.Vector([c_r[0], c_r[1], c_r[2]]) 64 | op = p - o 65 | b = op.dot(d) 66 | det = b * b - op.dot(op) + r * r 67 | t = inf 68 | if det > 0: 69 | det = ti.sqrt(det) 70 | t = (b - det) if ((b - det) > eps) else ((b + det) 71 | if (b + det > eps) else inf) 72 | 73 | if t < min_t: 74 | min_index = sp_index 75 | min_t = t 76 | 77 | sp_index += 1 78 | return ti.Vector([min_t, min_index]) 79 | 80 | 81 | width, height = 1080, 720 82 | 83 | 84 | linear_pixel = ti.Vector.field(3, dtype=ti.f32, shape=(width, height)) 85 | errors = ti.field(dtype=ti.f32, shape=(width, height)) 86 | ray_count = ti.field(dtype=ti.i32, shape=(width, height)) 87 | bloom_pixel = ti.Vector.field(3, dtype=ti.f32, shape=(width, height)) 88 | pixel = ti.Vector.field(3, dtype=ti.f32, shape=(width, height)) 89 | 90 | gaussian_kernel_size = 9 91 | bloom_strength = 0.02 92 | ambient_weight = 0.2 93 | start_cursor = [0.0, 0.0] 94 | 95 | 96 | 97 | blur_kernel = ti.field(ti.f32, shape=(gaussian_kernel_size, gaussian_kernel_size)) 98 | blur_kernel.from_numpy(gaussian_kernel(gaussian_kernel_size, 1.)) 99 | 100 | skybox = Image('skybox.jpg') # z > 0 101 | 102 | 103 | last_camera_pos = ti.field(ti.f32, 3) 104 | camera_pos = ti.field(ti.f32, 3) 105 | focal_length = ti.field(ti.f32, shape=()) 106 | 107 | class SphereAOS: 108 | def __init__(self, center_radius, albedo, emission, roughness, metallic, ior): 109 | self.center_radius = center_radius 110 | self.albedo = albedo 111 | self.emission = emission 112 | self.roughness = roughness 113 | self.metallic = metallic 114 | self.ior = ior 115 | 116 | 117 | spheres_aos = [ 118 | SphereAOS( 119 | center_radius=ti.Vector([-2.,0.0, -3.0, 2.5]), albedo=ti.Vector([1.0, 1.0, 0.0]), emission=ti.Vector([0, 0, 0]), roughness=0.2, metallic=1.0, ior=2.495), 120 | SphereAOS( 121 | center_radius=ti.Vector([1.0, -1.5, 1.0, 1.0]), albedo=ti.Vector([1.0, 1.0, 1.0]), emission=ti.Vector([0, 0, 0]), roughness=0.0, metallic=0.0, ior=1.4), 122 | SphereAOS( 123 | center_radius=ti.Vector([0.0, -500., 0., 497.5]), albedo=ti.Vector([1.0, 0.8, 0.6]), emission=ti.Vector([0, 0, 0]), roughness=0.5, metallic=0.95, ior=2.90), 124 | SphereAOS( 125 | center_radius=ti.Vector([-1, -2.0, 0.5, 0.5]), albedo=ti.Vector([1.0, 1.0, 1.0]), emission=ti.Vector([0, 0, 0]), roughness=0.0, metallic=1.0, ior=2.5), 126 | SphereAOS( 127 | center_radius=ti.Vector([2.0, 0, -2.0, 0.7]), albedo=ti.Vector([0, 0, 0]), emission=ti.Vector([15, 1, 1]), roughness=0.0, metallic=0.8, ior=2.0), 128 | SphereAOS( 129 | center_radius=ti.Vector([-2.0, 1.0, 1.0, 0.7]), albedo=ti.Vector([0, 0, 0]), emission=ti.Vector([0.2, 15, 15]), roughness=0.0, metallic=0.8, ior=2.0), 130 | SphereAOS( 131 | center_radius=ti.Vector([-8.0, -1.5, -5, 1.01]), albedo=ti.Vector([1.0, 1.0, 1.0]), emission=ti.Vector([4, 4, 0.5]), roughness=0.5, metallic=0.2, ior=2.0), 132 | 133 | ] 134 | 135 | spheres_soa = SphereSOA(len(spheres_aos)) 136 | 137 | for i, sphere in enumerate(spheres_aos): 138 | spheres_soa.center_radius[i] = sphere.center_radius 139 | spheres_soa.albedos[i] = sphere.albedo 140 | spheres_soa.emissions[i] = sphere.emission 141 | spheres_soa.roughness[i] = sphere.roughness 142 | spheres_soa.metallics[i] = sphere.metallic 143 | spheres_soa.iors[i] = sphere.ior 144 | 145 | skybox.load() 146 | 147 | 148 | @ ti.func 149 | def reflect(I, N): 150 | return I - 2 * N.dot(I) * N 151 | 152 | 153 | @ ti.func 154 | def lerp(vl, vr, frac): 155 | return vl + frac * (vr - vl) 156 | 157 | 158 | @ ti.func 159 | def cubemap_coord(dir): 160 | eps = 1e-7 161 | coor = ti.Vector([0., 0.]) 162 | if dir.z >= 0 and dir.z >= abs(dir.y) - eps and dir.z >= abs(dir.x) - eps: 163 | coor = ti.Vector([3 / 8, 1 / 2]) + \ 164 | ti.Vector([dir.x / 8, dir.y / 6]) / dir.z 165 | if dir.z <= 0 and -dir.z >= abs(dir.y) - eps and -dir.z >= abs(dir.x) - eps: 166 | coor = ti.Vector([7 / 8, 1 / 2]) + \ 167 | ti.Vector([-dir.x / 8, dir.y / 6]) / -dir.z 168 | if dir.x <= 0 and -dir.x >= abs(dir.y) - eps and -dir.x >= abs(dir.z) - eps: 169 | coor = ti.Vector([1 / 8, 1 / 2]) + \ 170 | ti.Vector([dir.z / 8, dir.y / 6]) / -dir.x 171 | if dir.x >= 0 and dir.x >= abs(dir.y) - eps and dir.x >= abs(dir.z) - eps: 172 | coor = ti.Vector([5 / 8, 1 / 2]) + \ 173 | ti.Vector([-dir.z / 8, dir.y / 6]) / dir.x 174 | if dir.y >= 0 and dir.y >= abs(dir.x) - eps and dir.y >= abs(dir.z) - eps: 175 | coor = ti.Vector([3 / 8, 5 / 6]) + \ 176 | ti.Vector([dir.x / 8, -dir.z / 6]) / dir.y 177 | if dir.y <= 0 and -dir.y >= abs(dir.x) - eps and -dir.y >= abs(dir.z) - eps: 178 | coor = ti.Vector([3 / 8, 1 / 6]) + \ 179 | ti.Vector([dir.x / 8, dir.z / 6]) / -dir.y 180 | return coor 181 | 182 | 183 | @ ti.func 184 | def cosine_sample(n): 185 | u = ti.Vector([1.0, 0.0, 0.0]) 186 | if abs(n[1]) < 1 - eps: 187 | u = n.cross(ti.Vector([0.0, 1.0, 0.0])).normalized() 188 | v = n.cross(u) 189 | phi = 2.0 * math.pi * ti.random() 190 | ay = ti.sqrt(ti.random()) 191 | ax = ti.sqrt(1.0 - ay**2.0) 192 | return ax * (ti.cos(phi) * u + ti.sin(phi) * v) + ay * n 193 | 194 | 195 | @ ti.func 196 | def ggx_sample(n, wo, roughness): 197 | u = ti.Vector([1.0, 0.0, 0.0]) 198 | if abs(n[1]) < 1.0 - eps: 199 | u = n.cross(ti.Vector([0.0, 1.0, 0.0])).normalized() 200 | v = n.cross(u) 201 | r0 = ti.random() 202 | r1 = ti.random() 203 | a = roughness ** 2.0 204 | a2 = a * a 205 | theta = ti.acos(ti.sqrt((1.0 - r0) / ((a2-1.0)*r0 + 1.0))) 206 | phi = 2.0 * math.pi * r1 207 | wm = ti.cos(theta) * n + ti.sin(theta) * ti.cos(phi) * \ 208 | u + ti.sin(theta) * ti.sin(phi) * v 209 | wi = reflect(-wo, wm) 210 | return wi 211 | 212 | 213 | @ ti.func 214 | def ggx_ndf(wi, wo, n, roughness): 215 | m = (wi + wo).normalized() 216 | a = roughness * roughness 217 | nm2 = n.dot(m) ** 2.0 218 | return (a * a) / (math.pi * nm2*(a*a-1.0)+1.0**2.0) 219 | 220 | 221 | @ ti.func 222 | def ggx_ndf_wi(wi, wo, n, roughness): 223 | m = (wi + wo).normalized() 224 | wm_pdf = ggx_ndf(wi, wo, n, roughness) 225 | wi_pf = n.dot(m) / (4 * m.dot(wi)) 226 | return wi_pf 227 | 228 | 229 | @ ti.func 230 | def ggx_pdf(roughness, hdotn, vdoth): 231 | t = hdotn*hdotn*roughness*roughness - (hdotn*hdotn - 1.0) 232 | D = (roughness*roughness) * math.pi / (t*t) 233 | return D*hdotn / (4.0 * abs(vdoth)) 234 | 235 | 236 | @ ti.func 237 | def schlick2(wi, n, f0): 238 | return f0 + (1.0-f0) * (1-n.dot(wi))**5.0 239 | 240 | 241 | @ ti.func 242 | def smith_g(NoV, NoL, roughness): 243 | a2 = roughness * roughness 244 | ggx_v = NoL * ti.sqrt(NoV * NoV * (1.0 - a2) + a2) 245 | ggx_l = NoV * ti.sqrt(NoL * NoL * (1.0 - a2) + a2) 246 | return 0.5 / (ggx_v + ggx_l) 247 | 248 | 249 | @ti.func 250 | def ggx_smith_uncorrelated(roughness, hdotn, vdotn, ldotn, fresnel): 251 | t = hdotn*hdotn*roughness*roughness - (hdotn*hdotn - 1.0) 252 | D = (roughness*roughness) * math.pi / (t*t) 253 | F = fresnel 254 | Gv = vdotn * ti.sqrt(roughness*roughness + 255 | (1.0 - roughness * roughness)*ldotn*ldotn) 256 | Gl = ldotn * ti.sqrt(roughness*roughness + 257 | (1.0 - roughness * roughness)*vdotn*vdotn) 258 | G = 1.0 / (Gv + Gl) 259 | return F*G*D / 2.0 260 | 261 | 262 | @ ti.func 263 | def luma(albedo): 264 | return albedo.dot(ti.Vector([0.2126, 0.7152, 0.0722])) 265 | 266 | @ti.func 267 | def tone(x): 268 | A = 2.51 269 | B = 0.03 270 | C = 2.43 271 | D = 0.59 272 | E = 0.14 273 | return (x * (A * x + B)) / (x * (C * x + D) + E) 274 | 275 | 276 | 277 | gui = ti.GUI("Path Tracer", res=(width, height)) 278 | 279 | 280 | @ ti.kernel 281 | def trace(): 282 | for i, j in linear_pixel: 283 | adaptive_sampler_count = int(16 * tone(errors[i, j] / ray_count[i, j])) 284 | # pixel[i, j].fill(adaptive_sampler_count / 16) 285 | for addition_sampler in ti.ndrange(16): 286 | if addition_sampler > adaptive_sampler_count: 287 | continue 288 | ray_count[i, j] += 1 289 | sample = ray_count[i, j] 290 | o = ti.Vector([camera_pos[0], camera_pos[1], camera_pos[2]]) 291 | aperture_size = 0.1 292 | forward = -o.normalized() 293 | 294 | u = ti.Vector([0.0, 1.0, 0.0]).cross(forward).normalized() 295 | v = forward.cross(u).normalized() 296 | u = -u 297 | 298 | d = ((i + ti.random() - width / 2.0) / width * u 299 | + (j + ti.random() - height / 2.0) / width * v 300 | + width / width * forward).normalized() 301 | 302 | focal_point = d * focal_length[None] / d.dot(forward) + o 303 | 304 | # assuming a circle-like aperture 305 | phi = 2.0 * math.pi * ti.random() 306 | aperture_radius = ti.random() * aperture_size 307 | o += u * aperture_radius * ti.cos(phi) + \ 308 | v * aperture_radius * ti.sin(phi) 309 | 310 | d = (focal_point - o).normalized() 311 | uv = cubemap_coord(d) 312 | albedo_factor = ti.Vector([1.0, 1.0, 1.0]) 313 | radiance = ti.Vector([0.0, 0.0, 0.0]) 314 | for step in ti.ndrange(32): 315 | sp = spheres_soa.intersect(o, d) 316 | uv = cubemap_coord(d) 317 | if sp[1] > -1: 318 | sp_index = int(sp[1]) 319 | p = sp[0] * d + o 320 | c_ = spheres_soa.center_radius[sp_index] 321 | c = ti.Vector([c_[0], c_[1], c_[2]]) 322 | radius = c_[3] 323 | n = (p - c).normalized() 324 | wo = -d 325 | albedo = spheres_soa.albedos[sp_index] 326 | metallic = spheres_soa.metallics[sp_index] 327 | ior = spheres_soa.iors[sp_index] 328 | roughness = ti.max(0.04, spheres_soa.roughness[sp_index]) 329 | f0 = (1.0 - ior) / (1.0 + ior) 330 | f0 = f0 * f0 331 | f0 = lerp(f0, luma(albedo), metallic) 332 | wi = reflect(-wo, n) 333 | radiance += spheres_soa.emissions[sp_index] / (radius * radius) * albedo_factor 334 | 335 | view_fresnel = schlick2(wo, n, f0) 336 | sample_weights = ti.Vector([1.0 - view_fresnel, view_fresnel]) 337 | 338 | weight = 0.0 339 | 340 | h = (wi + wo).normalized() 341 | 342 | shaded = ti.Vector([0.0, 0.0, 0.0]) 343 | 344 | if ti.random() < sample_weights[0]: 345 | wi = cosine_sample(n).normalized() 346 | h = (wi + wo).normalized() 347 | shaded = ti.max(0.00, wi.dot(n) * albedo / math.pi) 348 | 349 | else: 350 | wi = ggx_sample(n, wo, roughness).normalized() 351 | h = (wi + wo).normalized() 352 | F = schlick2(wi, n, f0) 353 | shaded = ti.max(0.0, wi.dot(n) * albedo * ggx_smith_uncorrelated( 354 | roughness, h.dot(n), wo.dot(n), wi.dot(n), F)) 355 | 356 | pdf_lambert = wi.dot(n) / math.pi 357 | pdf_ggx = ggx_pdf(roughness, h.dot(n), wo.dot(h)) 358 | 359 | weight = ti.max(0.0, 1.0 / (sample_weights[0] * 360 | pdf_lambert + sample_weights[1] * pdf_ggx)) 361 | 362 | # russian roule 363 | albedo_factor *= shaded * weight 364 | if step > 5: 365 | if luma(albedo) < ti.random(): 366 | break 367 | else: 368 | albedo_factor /= luma(albedo) 369 | d = wi 370 | o = p + eps * d 371 | 372 | else: 373 | radiance += albedo_factor * (tex2d(skybox.field, uv)) * ambient_weight 374 | break 375 | linear_color = radiance 376 | 377 | linear_pixel[i, j] = (linear_pixel[i, j] * (sample - 1) + linear_color) / sample 378 | 379 | errorVec = linear_color - linear_pixel[i, j] 380 | error = errorVec.dot(errorVec) 381 | errors[i, j] += error 382 | 383 | luminance = radiance.dot(ti.Vector([0.2126, 0.7152, 0.0722])) 384 | 385 | if(luminance < max_luminance): 386 | if (luminance > 1.0): 387 | bloom_pixel[i, j] = (bloom_pixel[i, j] * (sample - 1) + linear_color) / sample 388 | else: 389 | bloom_pixel[i, j] = (bloom_pixel[i, j] * (sample - 1)) / sample 390 | 391 | for i, j in bloom_pixel: 392 | hdr_blur = ti.Vector([0.0, 0.0, 0.0]) 393 | for kx in ti.ndrange(gaussian_kernel_size): 394 | for ky in ti.ndrange(gaussian_kernel_size): 395 | hdr_blur += bloom_pixel[i + kx - gaussian_kernel_size // 2, j + ky - gaussian_kernel_size // 2] * blur_kernel[kx, ky] 396 | 397 | bloom_pixel[i, j] = hdr_blur 398 | 399 | for i, j in pixel: 400 | pixel[i, j] = pow(tone(linear_pixel[i, j] + bloom_pixel[i, j] * bloom_strength), 1.0 / 2.2) 401 | 402 | 403 | def try_reset(t): 404 | for e in gui.get_events(gui.PRESS, gui.MOTION): 405 | if e.key == gui.LMB: 406 | last_camera_pos[0] = camera_pos[0] 407 | last_camera_pos[1] = camera_pos[1] 408 | last_camera_pos[2] = camera_pos[2] 409 | start_cursor_x, start_cursor_y = gui.get_cursor_pos() 410 | start_cursor[0] = start_cursor_x 411 | start_cursor[1] = start_cursor_y 412 | 413 | if e.key == gui.MOVE and gui.is_pressed(gui.LMB, gui.RMB): 414 | (current_cursor_x, current_cursor_y) = gui.get_cursor_pos() 415 | 416 | rotateX = (current_cursor_x - start_cursor[0]) 417 | rotateY = (current_cursor_y - start_cursor[1]) 418 | 419 | x = math.cos( 420 | rotateX) * last_camera_pos[0] - math.sin(rotateX) * last_camera_pos[2] 421 | z = math.sin( 422 | rotateX) * last_camera_pos[0] + math.cos(rotateX) * last_camera_pos[2] 423 | y = last_camera_pos[1] 424 | camera_pos[1] = math.cos( 425 | rotateY) * y - math.sin(rotateY) * z 426 | camera_pos[2] = math.sin( 427 | rotateY) * y + math.cos(rotateY) * z 428 | camera_pos[0] = x 429 | return True 430 | 431 | if e.key == gui.WHEEL: 432 | dt = e.delta[1] / 1000.0 433 | if gui.is_pressed(gui.SHIFT): 434 | focal_length[None] *= (1.0 + dt) 435 | else: 436 | for i in range(3): 437 | camera_pos[i] *= (1.0 + dt) 438 | return True 439 | 440 | return False 441 | 442 | # Initialize the camera 443 | sample = 0 444 | t = 1 445 | focal_length[None] = 15.0 446 | last_camera_pos[2] = focal_length[None] 447 | camera_pos[2] = focal_length[None] 448 | 449 | print("Start tracing...") 450 | print("0 spp") 451 | 452 | while True: 453 | 454 | if try_reset(t): 455 | if sample > 1: 456 | print ("\033[A \033[A") 457 | print("Camera moved, restart tracing...") 458 | print("0 spp") 459 | sample = 0 460 | linear_pixel.fill(0) 461 | bloom_pixel.fill(0) 462 | ray_count.fill(0) 463 | errors.fill(0) 464 | 465 | print ("\033[A \033[A") 466 | print(f"[{sample} spp]" ) 467 | 468 | sample += 1 469 | trace() 470 | t += 1 471 | gui.set_image(pixel) 472 | gui.show() 473 | --------------------------------------------------------------------------------