花了2天时间，写了个基础版本的ray tracer
# coding=utf-8
'''
@author: guopei
'''
import taichi as ti
ti.init(arch=ti.gpu)
n = 640
ScreenWidth = n
ScreenHeight = n
ScreenBuffer = ti.Vector(3, dt=ti.f32, shape=(ScreenWidth, ScreenHeight))
# 屏幕射线
ScreenRayOrigins = ti.Vector(3, dt=ti.f32, shape=(ScreenWidth, ScreenHeight))
ScreenRayDirs = ti.Vector(3, dt=ti.f32, shape=(ScreenWidth, ScreenHeight))
ScreenRaySpecularRate = ti.var(dt = ti.i32, shape=(ScreenWidth, ScreenHeight))
# 屏幕射线是否有效
ScreenRayValid = ti.var(dt=ti.i32, shape=(ScreenWidth, ScreenHeight))
# 球的数量
SphereNum = 5
#材质标记 x:diffuse, y:specular, z:ambient
MaterialFlags = ti.Vector(3, dt=ti.i32, shape=(SphereNum))
DiffuseColors = ti.Vector(3, dt=ti.f32, shape=(SphereNum))
# 反射强度
SpecularRate = ti.var(dt = ti.f32, shape=(SphereNum))
# 环境光颜色
AmbientColors = ti.Vector(3, dt=ti.f32, shape=(SphereNum))
# 球体位置
SphereCenters = ti.Vector(3, dt=ti.f32, shape=(SphereNum))
# 球体半径
SphereRadius = ti.var(dt = ti.f32, shape=(SphereNum))
#---------------------------------------------------------
# camera参数
gCameraPos = ti.Vector(3, dt=ti.f32, shape=())
gCameraLookAt = ti.Vector(3, dt=ti.f32, shape=())
# camera fov
Fov = 3.1415926 * 0.25
# camera 宽高比
WidthHeightRatio = 1.0/1.0
BottomLeftCorner = ti.Vector(3, dt=ti.f32, shape=())
ProjPlaneRight = ti.Vector(3, dt=ti.f32, shape=())
ProjPlaneUp = ti.Vector(3, dt=ti.f32, shape=())
ProjPlaneWidth = ti.var(dt=ti.f32, shape=())
ProjPlaneHeight = ti.var(dt=ti.f32, shape=())
#---------------------------------------------------------
# 光源定义
# 光源数量
LightNumber = 2
PointLightPositions = ti.Vector(3, dt=ti.f32, shape=(LightNumber))
PointLightColors = ti.Vector(3, dt=ti.f32, shape=(LightNumber))
#---------------------------------------------------------
gCameraPos[None] = (0,0,0)
gCameraLookAt[None] = (0,0,1)
#---------------------------------------------------------
# 球列表
SphereCenters[0] = (1,0,5)
SphereRadius[0] = 1
MaterialFlags[0] = (1, 0, 1)
DiffuseColors[0] = (1,1, 0.5)
SpecularRate[0] = 1
AmbientColors[0] = (0.3,0.3,0.3)
SphereCenters[1] = (-2.5,0,5)
SphereRadius[1] = 2
MaterialFlags[1] = (0, 1, 1 )
DiffuseColors[1] = (0.5,1, 1)
SpecularRate[1] = 1
AmbientColors[1] = (0.1,0.1,0.1)
SphereCenters[2] = (0,2,5)
SphereRadius[2] = 1
MaterialFlags[2] = (1, 0, 1 )
DiffuseColors[2] = (1,0.2, 0.2)
SpecularRate[2] = 1
AmbientColors[2] = (0.1,0.1,0.1)
SphereCenters[3] = (-2.5,4,5)
SphereRadius[3] = 1
MaterialFlags[3] = (1, 0, 1 )
DiffuseColors[3] = (0.2,1, 0.2)
SpecularRate[3] = 1
AmbientColors[3] = (0.1,0.1,0.1)
SphereCenters[4] = (-2.5,-4,5)
SphereRadius[4] = 1
MaterialFlags[4] = (1, 0, 1 )
DiffuseColors[4] = (0.2,0.2, 1)
SpecularRate[4] = 1
AmbientColors[4] = (0.1,0.1,0.1)
#SphereCenters[5] = (10, 0,5)
#SphereRadius[5] = 7
#MaterialFlags[5] = (0, 1, 1 )
#DiffuseColors[5] = (0.2,0.2, 1)
#SpecularRate[5] = 1
#AmbientColors[5] = (0.1,0.1,0.1)
#---------------------------------------------------------
# 光源列表
PointLightPositions[0] = (0,0,3)
PointLightColors[0] = (1, 1, 1)
PointLightPositions[1] = (1,0,5)
PointLightColors[1] = (4, 1, 1)
# 构造camera
@ti.kernel
def BuildCameraSetting():
#Forward = func1()
Forward = (gCameraLookAt[None] - gCameraPos[None]).normalized()
#Up = ti.Vector((0, 1, 0))
#Right = Up.cross( Forward)
HalfScreenWidth = ti.tan(Fov)
HalfScreenHeight = HalfScreenWidth / WidthHeightRatio
ProjPlaneRight[None] = ti.Vector((0,1,0)).cross(Forward)
ProjPlaneUp[None] = Forward.cross(ProjPlaneRight[None])
ProjPlaneWidth[None] = HalfScreenWidth * 2
ProjPlaneHeight[None] = HalfScreenHeight * 2
BottomLeftCorner[None] = gCameraPos[None] + Forward - ProjPlaneRight[None] * HalfScreenWidth - ProjPlaneUp[None] * HalfScreenHeight
#print("Fov", Fov, "Forward:",Forward, "HalfScreenWidth:", HalfScreenWidth, "BottomLeftCorner,", BottomLeftCorner[None])
pass
@ti.func
def BuildRay(u, v):
target_pt = BottomLeftCorner[None] + ProjPlaneRight[
None] * ProjPlaneWidth[None] * u + ProjPlaneUp[
None] * ProjPlaneWidth[None] * v
return (gCameraPos[None], (target_pt - gCameraPos[None]).normalized())
@ti.func
def Reflect(v, n):
return 2 * v.dot(n) * n - v
@ti.func
def IntersectWithSphere( ray_origin, ray_dir, sphere_center, sphere_radius):
'''
计算射线与球的交点
@param ray_origin :ti.Vector
'''
hit_point = ti.Vector((0.0,0.0,0.0))
hit_normal = ti.Vector((0.0,0.0,1.0))
IsHit = False
t = 0
dist = (ray_origin - sphere_center).norm()
if dist > sphere_radius:
origin2center = sphere_center - ray_origin
#计算投影向量
proj_vec = origin2center.dot(ray_dir) * ray_dir;
vertical_vec = proj_vec - origin2center
d = vertical_vec.norm()
if d < sphere_radius:
# 两个交点，只取最近的那个
dist = ti.sqrt(sphere_radius ** 2 - d ** 2)
nor_negproj = (-proj_vec).normalized()
hit_point = ray_origin + origin2center + vertical_vec + nor_negproj * dist
hit_normal = (hit_point - sphere_center).normalized()
is_samedir = ray_dir.dot(hit_normal) >= 0
if is_samedir:
IsHit = False
else:
IsHit = True
t = (hit_point - ray_origin).norm()
#print("IntersectWithSphere", "t", t, "hit_point", hit_point, "hit_normal", hit_normal, "ray_origin", ray_origin,
# "ray_dir", ray_dir, "origin2center", origin2center, "vertical_vec",vertical_vec,
# "nor_negproj", nor_negproj, "dist", dist, "hit_point", hit_point, "sphere_center", hit_point)
pass
elif d == sphere_radius:
# 1个交点
hit_point = ray_origin + origin2center + vertical_vec
hit_normal = (hit_point - sphere_center).normalized()
is_samedir = ray_dir.dot(hit_normal) >= 0
if is_samedir:
IsHit = False
else:
IsHit = True
t = (hit_point - ray_origin).norm()
pass
else:
# 没有交点
pass
else:
IsHit = False
return (IsHit, t,hit_point, hit_normal)
@ti.func
def IntersectForShadow(ray_origin, ray_dir):
tmin = 9999999
is_hit = False
for k in range(SphereNum):
(hit, t, _, _) = IntersectWithSphere(ray_origin, ray_dir, SphereCenters[
k], SphereRadius[k])
if hit and t < tmin:
is_hit = True
tmin = t
pass
return (is_hit, tmin)
@ti.func
def IntersectWithSpheres(ray_origin, ray_dir):
tmin = 9999999
hitpoint_min = ti.Vector((0.0,0.0,0.0))
hitnormal_min = ti.Vector((0.0,0.0,1.0))
is_hit = False
hit_idx = -1
for k in range(SphereNum):
(hit, t, hit_point, hit_normal) = IntersectWithSphere(ray_origin, ray_dir, SphereCenters[
k], SphereRadius[k])
if hit and t < tmin:
is_hit = True
hitpoint_min = hit_point
hitnormal_min = hit_normal
tmin = t
hit_idx = k
pass
return (is_hit, hit_idx, hitpoint_min, hitnormal_min, tmin)
#ti.func
def CalcUV(i, j, screen_width, screen_height):
'''
计算uv
'''
u = (i + 0.5) / screen_width
v = (j + 0.5) / screen_height
return (u,v)
@ti.kernel
def BuildRays():
for i, j in ScreenBuffer:
(u,v) = CalcUV(i,j, ScreenWidth, ScreenHeight)
(ray_origin, ray_dir) = BuildRay(u,v)
ScreenRayOrigins[i,j] = ray_origin
ScreenRayDirs[i,j] = ray_dir
ScreenRaySpecularRate[i,j] = 1.0
ScreenRayValid[i,j]=1
@ti.func
def ProcRay(i, j):
ray_origin = ScreenRayOrigins[i,j]
ray_dir = ScreenRayDirs[i,j]
(is_hit, hit_idx, hitpoint_min, hitnormal_min, t) = IntersectWithSpheres(
ray_origin, ray_dir)
if is_hit:
#print("hit! ray_origin", ray_origin, "ray_dir", ray_dir)
Shade(i, j, hit_idx, hitpoint_min, hitnormal_min)
pass
@ti.func
def Shade(i, j, sphere_idx, hit_point, hit_normal):
if MaterialFlags[sphere_idx][0] == 1:
# diffuse
# 遍历所有光源，看看能不能击中光源，能的话，应用这个光源
for k in range(LightNumber):
light_pos = PointLightPositions[k]
distance = (light_pos - hit_point).norm()
ray_dir = (light_pos - hit_point).normalized()
(is_hit, t) = IntersectForShadow(hit_point, ray_dir)
if not is_hit or (is_hit and t >= distance):
# 没有击中球，或者击中的球比光源还要远，那么这个光源能照到当前位置
color = ShadeForDiffuse(hit_point, hit_normal, sphere_idx, k) * ScreenRaySpecularRate[i, j]
ScreenBuffer[i,j] += color
ScreenRayValid[i,j] = 0
if MaterialFlags[sphere_idx][1] == 1:
# 镜面反射
# 首先计算新的反射率
new_specrate = SpecularRate[sphere_idx] * ScreenRaySpecularRate[i, j]
if new_specrate < 0.01:
# 如果这个反射率太低了，那么就不再构造反射射线了
ScreenRaySpecularRate[i, j] = 0
ScreenRayValid[i,j] = 0
else:
# 当前这个位置反射率还够高，构造反射射线
ScreenRaySpecularRate[i, j] = new_specrate
ScreenRayOrigins[i,j] = hit_point
ray_dir = ScreenRayDirs[i,j]
ScreenRayDirs[i,j] = Reflect(-ray_dir, hit_normal)
ScreenRayValid[i,j] = 1
pass
if MaterialFlags[sphere_idx][2] == 1:
# 环境光
ScreenBuffer[i,j] += AmbientColors[sphere_idx]
pass
@ti.func
def ShadeForDiffuse(hit_point, hit_normal, sphere_idx, light_idx):
light_dir = ( PointLightPositions[light_idx] - hit_point ).normalized()
ndotl = light_dir.dot(hit_normal)
if ndotl < 0:
ndotl = 0
return DiffuseColors[sphere_idx] * PointLightColors[light_idx] * ndotl
@ti.kernel
def Render():
for i,j in ScreenBuffer:
# 最多反射5次
for _ in range(5):
if ScreenRayValid[i,j] != 0:
ProcRay(i, j)
pass
#func1()
gui = ti.GUI("screen", (ScreenWidth, ScreenHeight))
# 构造 camera
BuildCameraSetting()
#构造所有射线
BuildRays()
# 清空back buffer
ScreenBuffer.fill(0)
Render()
while True:
if gui.get_event(ti.GUI.PRESS):
if gui.event.key == ti.GUI.LMB:
x, y = gui.get_cursor_pos()
gCameraLookAt[None][0] = x * 4 - 2
gCameraLookAt[None][1] = y * 2 - 1
# 构造 camera
BuildCameraSetting()
#构造所有射线
BuildRays()
# 清空back buffer
ScreenBuffer.fill(0)
Render()
gui.set_image(ScreenBuffer.to_numpy())
gui.show()
print("done")
