雨上がりの町でライトが水たまりに反射した夜景、綺麗ですよね。こうした景色はコンピュータ上でも再現することができて、それを実現する技術の一つにレイトレーシングというものがあります。本記事では簡易的なレイトレーシングをPyTorchを用いて実装してみました。

　基本的にmain.pyの1ファイルで完結しています。また、importは省略しています。

環境

　主な環境を書いておきます

• OS: Windows11
• Python: Python3.12.3
• torch: 2.3.1+cu121

レイトレーシングについて

　レイトレーシングの基本的なアイデアは「すべての光線の軌跡をたどる」ということです。しかし本当にすべての光線をたどると非効率的です。最終的に得たいものは視点となるカメラからの景色です。そのため、カメラからたくさん光線(レイ)を飛ばし、それがどういった経路をたどるかを求めてカメラからの見え方を逆算するという手法を採用しています。

　レイトレーシングといってもいろいろ手法があるようで、今回扱う実装以外にも、レイマーチングという物もあるそうです。今回は、愚直に個々のレイに対して、描画対象となる空間(ワールド)内のすべてのオブジェクトと当たり判定処理を行い、反射などの処理を行います。

基本的なワールド描画

PyTorchを用いたベクトルの計算

　PyTorchを用いるので、計算で使う型は基本的に torch.Tensor になります。これを3次元ベクトルの配列(Nx3行列)としてみて色々計算していきます。まずはベクトルの基本的な計算を行うための関数を用意します。

def lerp(a: torch.Tensor, b: torch.Tensor, t: torch.Tensor):
    t = torch.clamp(t, 0, 1)
    return a + (b - a) * t

def tensor_dot(a: torch.Tensor, b: torch.Tensor):
    return torch.sum(a * b, dim=1)

def tensor_cross(a: torch.Tensor, b: torch.Tensor):
    return a.cross(b, dim=1)

def tensor_norm(a: torch.Tensor):
    return torch.sqrt(torch.sum(a ** 2, dim=1))

def tensor_normalize(a: torch.Tensor):
    return a / tensor_norm(a).view(-1, 1)

　あまり難しいところはないです。torch.Tensor.cross で指定された次元でのベクトルで外積を取れるのは意外でした。

　この記事の核心はほとんどここになります。実際、他の部分は通常のレイを一本ずつ処理する実装を借りて、テンソルの形を合わせるだけです。

レイ

　レイトレーシングの核となるレイの実装です。レイは始点から飛んでいく(光)線なので、0以上のスカラ値tと二つのベクトルで表現します。

class Rays:
    def __init__(self, origin: torch.Tensor, direction: torch.Tensor):
        self.origin = origin
        self.direction = direction

    def at(self, t: torch.Tensor):
        t = t.view(-1, 1)
        return self.origin + self.direction * t
    
    def __len__(self):
        return len(self.origin)

　またコードからもわかるように、この時点で複数のレイを一度に扱う実装にしています。最終的にはカメラから飛んでいるすべてのレイを一気に処理していきます。

カメラ

　カメラは描画地点や向き等の情報を保持し、描画時の最初のレイの作成を担当します。今回はカメラ位置と向いている場所からカメラの向きを決定し、さらにスクリーンの上方向を指定してスクリーンの回転を決定、そして(縦方向の)視野角を指定することでスクリーンの視界を決定します。これらの情報から、レイ作成に必要となる

• スクリーン左上の座標
• スクリーン横方向の向きと幅
• スクリーン縦方向の向きと高さ

の三つを計算します。詳しい数学的な解説は こちら を参照してください。とても分かりやすく書かれています。

class Camera:
    def __init__(self, width: float, height: float, origin: torch.Tensor, lookat: torch.Tensor, vup: torch.Tensor, vfov, device: torch.device='cpu'):
        self.device = device
        self.width = width
        self.height = height
        self.lookfrom = origin.to(device)
        self.lookat = lookat.to(device)
        self.vup = vup.to(device)
        self.vfov = vfov
        self.aspect_ratio = width / height
        self.build_uvw()

    def build_uvw(self):
        theta = self.vfov * pi / 180
        h = tan(theta / 2)
        viewport_height = 2.0 * h
        viewport_width = self.aspect_ratio * viewport_height
        
        w = tensor_normalize((self.lookfrom - self.lookat).view(1, -1))
        u = tensor_normalize(tensor_cross(self.vup.view(1, -1), w))
        v = tensor_cross(w, u)
        
        self.origin = self.lookfrom
        self.horizontal = viewport_width * u
        self.vertical = viewport_height * v
        self.top_left_corner = self.origin - self.horizontal / 2 + self.vertical / 2 - w

        self.horizontal = self.horizontal[0]
        self.vertical = -self.vertical[0]
        self.top_left_corner = self.top_left_corner[0]
        
    def get_ray(self, u: torch.Tensor, v: torch.Tensor):
        u = u.view(-1, 1)
        v = v.view(-1, 1)
        return Rays(self.origin.repeat(len(u)).view(-1, 3), self.top_left_corner + u * self.horizontal + self.vertical * v - self.origin)

　参考記事では左下を原点としていましたが、PILを用いた画像書き込みの都合上、左上を原点としています。するとワールドで使用する座標系と縦方向の向き(y)が合わなくなるのでvは-1倍されています。

ワールド

　まずはワールド内にオブジェクトが何もなく、レンダリングを行っても空しかない状態を作ります。

　必要になる機能はレンダリング機能だけになります。レンダリングは、(1)レイを作成、(2)各レイについてその方向で見える色を取得、(3)描画という流れで行います。

class World:
    def __init__(self, device: torch.device='cpu'):
        self.device = device

    def color(self, rays: Rays):
        return self.background_skys(rays)

    def background_skys(self, rays: Rays):
        dir = tensor_normalize(rays.direction)
        RANGE = 0.1
        t = (dir[:,1] + RANGE) / (RANGE * 2)
        return lerp(torch.tensor([0.4, 0.4, 0.4], device=self.device), torch.tensor([0.5, 0.7, 1.0], device=self.device), t.unsqueeze(1))

    def render(self, camera: Camera) -> torch.Tensor:
        u = torch.linspace(0, 1, camera.width).repeat(camera.height, 1).flatten().to(self.device)
        v = torch.linspace(0, 1, camera.height).repeat(camera.width, 1).t().flatten().to(self.device)

        rays = camera.get_ray(u, v)
        colors = self.color(rays)
        
        colors = torch.clamp(colors, 0, 1)
        colors = colors ** (1/2.2)
        colors = colors.view(camera.height, camera.width, 3)

        return colors

　最後に色を1/2.2乗しているのは、ディスプレイでのガンマ補正を打ち消して画像が暗くなることを防ぐためです。

　render関数でのレイの作成がポイントです。まずスクリーン上のuv座標を列挙したtensorを作成し、それをcamera.get_rayに渡すことでこれから飛ばす全てのレイを一気に取得します。

レンダリングしてみる

　とりあえずレンダリングができるところまで実装できたので、レンダリングをしてみます。

DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'

world = World(device=DEVICE)
camera = Camera(1600, 900, 
                torch.tensor([3, 2.4, 4], dtype=torch.float64) / 2, 
                torch.tensor([-3, -1.6, -4], dtype=torch.float64) / 2, 
                torch.tensor([0, 1, 0], dtype=torch.float64), 
                60, device=DEVICE)
colors = world.render(camera).cpu().numpy()
colors = (colors * 255).astype(np.uint8)
Image.fromarray(colors).save('output.png')

するとこんな画像が出力されます。

　いい感じですね。次はこれにオブジェクトを追加して、光線の反射も実装していきます。

球の描画

レイの衝突の情報

　ここからはレイをオブジェクトに衝突させることになるので、この衝突の情報を保持するためのクラスが必要です。このクラスも、Raysクラス同様に複数のレイの衝突の情報を一度に保持させます。

class HitRecords:
    def __init__(self, t: torch.Tensor, p: torch.Tensor, normal: torch.Tensor):
        self.t = t
        self.p = p
        self.normal = normal

　運用として、衝突しないレイはtを大きな数にしておきます(距離無限大の先で空と衝突)。

球

　次は球を作っていきます。レンダリングを行う上で球に必要な機能は、飛んできたレイが球と当たるのかの判定と、当たる場合はその当たった場所と表面の法線の取得です。これらの判定はレイのパラメータtに関する二次方程式を解くことで行うことができます。

レイの始点と向きを$a$、$b$として、球の中心位置を$s$、半径を$r$とします。
レイが急に衝突する時、とある$t$について

$$|a+tb-s|=r$$

が成り立ちます。両辺を二乗し、tについて整理すると次のようになります。

$$|b{|}^{2}t+2(a-s)\cdot bt+|a-s{|}^2-r^2=0$$

これを解けばいいわけです。

　後は判別式を持ちいて衝突があるかを判定し、解が複数ある場合は(指定された範囲内で)小さい方を採用します。

class Sphere:
    def __init__(self, center: torch.Tensor, radius: float, device: torch.device='cpu'):
        self.center = center.to(device)
        self.radius = radius
        self.device = device

    def hits(self, ray: Rays, t_min: float, t_max: float):
        oc = (ray.origin - self.center.view(1, -1))
        a = tensor_dot(ray.direction, ray.direction)
        b = 2.0 * tensor_dot(oc, ray.direction)
        c = tensor_dot(oc, oc) - self.radius ** 2
        discriminants = b ** 2 - 4 * a * c

        roots = torch.sqrt(discriminants)
        t1 = (-b - roots) / (2.0 * a)
        t2 = (-b + roots) / (2.0 * a)
        # t1とt2のうち、t_minとt_maxの範囲にあるものを選択
        # ただし、t1とt2が両方とも範囲外の場合はinfを返す
        # どちらも範囲内の場合はt1を返す

        t1_in_range = (t1 > t_min) & (t1 < t_max)
        t2_in_range = (t2 > t_min) & (t2 < t_max)
        both_out_of_range = ~t1_in_range & ~t2_in_range
        t2_in_range = t2_in_range & ~t1_in_range
        t1 = torch.where(t1_in_range, t1, torch.zeros_like(t1))
        t2 = torch.where(t2_in_range, t2, torch.zeros_like(t2))
        t = t1 + t2
        t = torch.where(both_out_of_range, torch.tensor(float('inf')), t)

        p = ray.at(t)
        normal = tensor_normalize(p - self.center)
        return HitRecords(t, p, normal)

　PyTorchでは各要素ごとに簡単な判定を行ってマスクを作成できるので、その機能を利用して適切な値を設定していきます。

オブジェクトをワールドに追加する

　球ができたので、これをワールドに追加させるための処理を書きます。ここでは特別な処理はせず、単にSphereのインスタンスを保持することにします。

    def __init__(self, device: torch.device='cpu'):
        self.objects = []
        self.device = device

    def add(self, obj):
        self.objects.append(obj)

objectsというメンバ変数を追加しました。

レイの衝突判定

　次はレイをオブジェクトに衝突させて反射させてみます。流れとしては

1. 各レイについて最初に衝突したオブジェクトを取得し、衝突の情報を取得
2. 衝突がないなら空を描画、衝突があった場合新しいレイを作成
3. 新しく作成されたレイを飛ばす
4. 1-3を再帰的に繰り返して、色を減衰させながら元のレイの色を計算

となります。衝突の情報を取得するために新しくメンバ関数hitsを定義し、メンバ関数colorを修正します。

    def color(self, rays: Rays, depth=0):
        if depth > 50:
            return torch.zeros(len(rays), 3, device=self.device, dtype=torch.float64)
        t_min = 0.001
        t_max = torch.tensor(float('inf'))
        
        hr = self.hits(rays, t_min, t_max)
        
        # t=t_maxのものはレイが当たっていないので背景色に
        # それ以外のレイは次のレイを生成して、再帰的に色を計算
        mask = hr.t < t_max
        
        normal = hr.normal[mask]
        direction = rays.direction[mask] - 2 * tensor_dot(rays.direction[mask], normal).view(-1, 1) * normal
        next_rays = Rays(hr.p[mask], direction)
        next_color = self.color(next_rays, depth + 1) * 0.64
        
        color = self.background_skys(rays)
        color[mask] = next_color
        return color

    def hits(self, rays: Rays, t_min: float, t_max: float):
        hit_records = HitRecords(
            t=torch.tensor(float('inf'), device=self.device).repeat(len(rays)),
            p=rays.origin,
            normal=torch.zeros_like(rays.origin, device=self.device)
        )
        for obj in self.objects:
            hits = obj.hits(rays, t_min, t_max)
            t = hits.t
            mask = t < hit_records.t
            hit_records.t = torch.where(mask, t, hit_records.t)
            hit_records.p = torch.where(mask.view(-1, 1), hits.p, hit_records.p)
            hit_records.normal = torch.where(mask.view(-1, 1), hits.normal, hit_records.normal)
        return hit_records

　ここでもマスク機能を利用して、より近い場所での衝突のふるい分け(hits)や、衝突したレイの分別(color)を行います。

　また、実装が再帰関数である以上、限りなく再帰してしまう場合もあるかもしれません。この場合の対策として、50回再帰した場合はその向きからの光はカメラに届かないとして、強制的に0を返すようにします。実際、反射するごとに色が0.64倍されるので、50回も反射すると色が消えて黒色になっています。

レンダリング

　これでオブジェクトの追加ができたので、実際に球を追加してレンダリングをしてみます。

DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'

world = World(device=DEVICE)
camera = Camera(1600, 900, 
                torch.tensor([3, 2.4, 4], dtype=torch.float64) / 2, 
                torch.tensor([-3, -1.6, -4], dtype=torch.float64) / 2, 
                torch.tensor([0, 1, 0], dtype=torch.float64), 
                60, device=DEVICE)

world.add(Sphere(torch.tensor([0, 0.5, 0]), 0.5, device=DEVICE))

colors = world.render(camera).cpu().numpy()

colors = (colors * 255).astype(np.uint8)
Image.fromarray(colors).save('output.png')

実際に球が出てきました！

平面の追加

　球だけだと物足りないので、球の下に平面を追加して、球がその上に置かれているようにします。実装はPlaneクラスを定義するだけです。球よりも処理は簡単なので解説は省略します。

class Plane:
    def __init__(self, point: torch.Tensor, normal: torch.Tensor, device: torch.device='cpu'):
        self.device = device
        self.point = point.to(device).view(1, -1)
        self.normal = normal.to(device).view(1, -1)

    def hits(self, ray: Rays, t_min: float, t_max: float):
        # ray.originからplane.pointへのベクトル
        p0 = self.point - ray.origin
        # ray.directionとplane.normalの内積
        n_dot_d = tensor_dot(ray.direction, self.normal)
        # p0とplane.normalの内積
        n_dot_p0 = tensor_dot(p0, self.normal)
        # rayがplaneに当たるかどうか
        mask = n_dot_d != 0
        t = torch.where(mask, n_dot_p0 / n_dot_d, torch.tensor(float('inf'), device=self.device))
        mask = (t > t_min) & (t < t_max)
        t = torch.where(mask, t, torch.tensor(float('inf'), device=self.device))
        p = ray.at(t)
        return HitRecords(t, p, self.normal)

　これで面を追加してレンダリングをします。

　綺麗に描画できました！

アンチエイリアス

　今まで画像を1600x900で描画していたのでわかりづらいですが、今までのレンダリングの方法だとオブジェクトの縁などがカクカクしてあまり良くありません。先ほどのワールドを解像度400x225で描画してみます。

　小さいですが、カクカクしているのがわかります。これへの対処として、アンチエイリアスを行います。

　今まで、1ピクセルの色を一つのuv座標からのレイで決めていました。これが問題で、画像がカクカクした原因になっています。ここでは、一つのピクセル内の複数のuv座標についてレイを飛ばし、それらで得られた色の平均値をとることでアンチエイリアスを行うことにします。

　render関数を修正して、新しく引数n_samplesをとります。そしてその数だけレイを1ピクセルに対して作成、計算します。また、すべてのレイを一気に計算するとメモリ容量が足りなくなるので、いくつかバッチに分けて計算をしてメモリの節約を図っています。BATCH_SIZEをお好みの数値に設定してください。

    def render(self, camera: Camera, n_samples: int) -> torch.Tensor:
        u_origin = torch.linspace(0, 1, camera.width).repeat(camera.height, 1).flatten().to(self.device)
        v_origin = torch.linspace(0, 1, camera.height).repeat(camera.width, 1).t().flatten().to(self.device)

        colors = torch.zeros(camera.width * camera.height, 3, dtype=torch.float64, device=self.device)

        u = []
        v = []

        for i in range(n_samples):
            noize_u = torch.rand_like(u_origin) / camera.width
            noize_v = torch.rand_like(v_origin) / camera.height
            u.append(u_origin + noize_u)
            v.append(v_origin + noize_v)

        u = torch.stack(u, dim=0).flatten()
        v = torch.stack(v, dim=0).flatten()
        colors = torch.tensor([])

        BATCH_SIZE = (1 << 23)

        for i in track(range(0, len(u), BATCH_SIZE), description='Rendering...'):
            batch_u = u[i:i+BATCH_SIZE]
            batch_v = v[i:i+BATCH_SIZE]
            rays = camera.get_ray(batch_u, batch_v)
            color = self.color(rays).to('cpu')
            colors = torch.cat([colors, color], dim=0)

        colors = colors.view(n_samples, camera.height * camera.width, 3)
        colors = colors.sum(dim=0) / n_samples
        colors = colors ** (1/2.2)
        colors = torch.clamp(colors, 0, 1)
        colors = colors.view(camera.height, camera.width, 3)

        return colors

　これでアンチエイリアスが実装できたので、実際にサンプリング数を指定して実行してみます。ここでは1ピクセル当たり100本のレイを飛ばします。

　先ほどと比較して色の境界が滑らかになったと思います。

ぼかしてみる

　ここまで、光の反射は全て鏡面反射のように実装していました。しかし、実際の物体は表面に細かい凹凸があって光の反射する向きが一定とは限りません。ここでは、反射時の法線の向きを少しランダムに変化させてこれを再現してみます。

color関数を修正します

    def color(self, rays: Rays, depth=0):
        if depth > 50:
            return torch.zeros(len(rays), 3, device=self.device, dtype=torch.float64)
        t_min = 0.001
        t_max = torch.tensor(float('inf'))
        
        hr = self.hits(rays, t_min, t_max)
        
        # t=t_maxのものはレイが当たっていないので背景色に
        # それ以外のレイは次のレイを生成して、再帰的に色を計算
        mask = hr.t < t_max
        
        normal = hr.normal[mask]

        rnd_r_max = tensor_norm(normal)
        rnd_r = torch.rand(len(normal), device=self.device) * rnd_r_max * 0.5
        rnd_phi = torch.rand(len(normal), device=self.device) * 2 * pi
        rnd_theta = torch.rand(len(normal), device=self.device) * pi
        rnd = torch.stack([rnd_r * torch.sin(rnd_theta) * torch.cos(rnd_phi), rnd_r * torch.sin(rnd_theta) * torch.sin(rnd_phi), rnd_r * torch.cos(rnd_theta)], dim=1)
        normal += rnd

        direction = rays.direction[mask] - 2 * tensor_dot(rays.direction[mask], normal).view(-1, 1) * normal
        next_rays = Rays(hr.p[mask], direction)
        next_color = self.color(next_rays, depth + 1) * 0.64
        
        color = self.background_skys(rays)
        color[mask] = next_color
        return color

　各法線の長さの定数倍を半径とした球内におけるランダムな点を選び、それを法線ベクトルに加算しています。ランダムに点を選ぶ操作は球面座標系を用いて行います。

　これでもう一度レンダリングをしてみます。

　ぼかすことに成功しました。

　ぼかす方法は他にも、先ほどランダムにずらした法線をそのまま反射光とする場合もあります。これには名前があり「ランバート」と言うようです。この場合はこんな感じのレンダリングになります。

　こちらの方がもっとくすんだ感じがしますね。

球を増やしてみる

　では球をさらに増やしてみます。真ん中の大きな球の周囲に小さな球を4つ配置してレンダリングしてみます。また、ぼかし処理でずらす法線の大きさを法線の長さの0.1倍にしてより鏡っぽくします。

DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'

world = World(device=DEVICE)
camera = Camera(1600, 900, 
                torch.tensor([3, 2.4, 4], dtype=torch.float64) / 2, 
                torch.tensor([-3, -1.6, -4], dtype=torch.float64) / 2, 
                torch.tensor([0, 1, 0], dtype=torch.float64), 
                60, device=DEVICE)

world.add(Sphere(torch.tensor([0, 0.5, 0]), 0.5, device=DEVICE))
world.add(Sphere(torch.tensor([1.2, 0.2, 0]), 0.2, device=DEVICE))
world.add(Sphere(torch.tensor([-1.2, 0.2, 0]), 0.2, device=DEVICE))
world.add(Sphere(torch.tensor([0, 0.2, 1.2]), 0.2, device=DEVICE))
world.add(Sphere(torch.tensor([0, 0.2, -1.2]), 0.2, device=DEVICE))

world.add(Plane(torch.tensor([0, 0, 0]), torch.tensor([0, 1, 0]), device=DEVICE))

colors = world.render(camera, 100).cpu().numpy()

colors = (colors * 255).astype(np.uint8)
Image.fromarray(colors).save('output.png')

　これを実行するとこんな感じになります。

　しっかりと大きい球に小さい球の反射が写っていていい感じです。美しい。

　ちなみに、反射の処理をランバートで行うとこんな感じになります。

　これもこれで柔らかい感じがあっていいですね。

まとめ

　今回はPyTorchを使って高速にレイトレーシングを実装してみました。そこそこに手軽に現実に近いものを描画することができました。個人的には初めてHelloWorldプログラムを書いたときくらいの感動がありました。現実で見るもの、することをコンピュータ上で再現できると嬉しいですね。

　本来GPUは描画処理のために開発されたもので、今回はPyTorchというGPGPUを行うモジュールを用いて描画を行ったと...。何かとんでもない大回りをした気がします。いっそのことCUDA直書きしてもよかったのかもしれない。とはいえ、手軽にGPUで計算ができるのはPyTorchの強みだと思います。並列計算はお任せあれといった感じです。本当に速い。CPUで一本一本レイを処理していたら1時間はかかったと思います。

　そういえば「Pythonで動作」していて「3Dのモデリング、レンダリングが可能」な「GPUも使える」ソフトがあったような気がします。オレンジ色のまつげが三本ついた目玉が思い浮かびますが、気のせいでしょう。

　本来この日に新歓記事を書く予定はなかったですが、速度と成果物にあまりにも感動してしまったので記事にすることにしました。明日は vPhos さんによる記事です！お楽しみに！