In [1]:
# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.

ボリュームレンダリングのテスト

  1. 微分可能なボリューメトリックレンダラーを作成します。
  2. ボリュームモデルを作成します（Volumes クラスの使用方法を含む）。
  3. 微分可能なボリュメトリックレンダラーを使用して、ボリュームを可視化します。

0. インストールとモジュールのインポート

torchとpytorch3dがインストールされていない場合は、インストールしてください。

In [2]:
import os
import sys
import time
import json
import glob
import torch
import math
from tqdm.notebook import tqdm
import matplotlib.pyplot as plt
import numpy as np
from PIL import Image
from IPython import display

# レンダリングのためのデータ構造と関数
from pytorch3d.structures import Volumes
from pytorch3d.renderer import (
    FoVPerspectiveCameras, 
    VolumeRenderer,
    NDCGridRaysampler,
    EmissionAbsorptionRaymarcher
)
from pytorch3d.transforms import so3_exponential_map

# デモ／ユーティリティ関数のためのパスの追加
sys.path.append(os.path.abspath(''))
from utils.plot_image_grid import image_grid
from utils.generate_cow_renders import generate_cow_renders

# 利用可能なデバイスの取得
if torch.cuda.is_available():
    device = torch.device("cuda:0")
    torch.cuda.set_device(device)
else:
    device = torch.device("cpu")

2. ボリュームレンダラの初期化

以下は、ターゲット画像の各ピクセルから光線を放射し、光線に沿って等間隔の点のセットをサンプリングするボリュームレンダラを初期化します。各レイポイントで、対応する密度と色の値は、シーンのボリュームメトリックモデル（モデルは後のセルで説明され、インスタンス化されます）の対応する位置を照会することによって得られます。

レンダラはraymarcherとraysamplerで構成されています。

  • raysamplerは画像のピクセルから光線を放出し、それに沿って点をサンプリングする役割を担っています。ここでは、標準的な PyTorch3D の座標グリッド（+X は右から左へ、+Y は下から上へ、+Z はユーザーから離れた位置）に従う NDCGridRaysampler を使用しています。
  • raymarcher は、各光線に沿ってサンプリングされた密度と色を取り、各光線を色と光線のソースピクセルの不透明度にレンダリングします。ここでは、標準的な Emission-Absorption レイマーチングアルゴリズムを実装した EmissionAbsorptionRaymarcher を使用しています。
In [3]:
# render_size は、レンダリングされた画像の両側のサイズをピクセル単位で指定します。
render_size = 512

# 私たちのレンダリングされたシーンは、(0,0,0)を中心にして、
# 側面が6.0(ワールド単位)にほぼ等しいバウンディングボックスの中に囲まれています。
volume_extent_world = 6.0

# 1) raysamplerのインスタンスを作成します。
# ここでは、NDCGridRaysampler を用いて、pytorch3d の座標規則に従った長方形の
# 画像グリッドを生成しています。512x512x310のボリュームを使用するので、n_pts_per_ray=550とします。
# さらに、カメラ平面から0.1ユニット以内には面がないので、min_depth=0.1とします。
raysampler = NDCGridRaysampler(
    image_width=render_size,
    image_height=render_size,
    n_pts_per_ray=550,
    min_depth=0.1,
    max_depth=volume_extent_world,
)


# 2) raymarcherのインスタンスを作成します。
# ここでは、標準的なEmissionAbsorptionRaymarcherを使用して、各光線を
# 単一の3Dカラーベクトルと不透明度スカラーにレンダリングするために、各光線に沿って行進します。
raymarcher = EmissionAbsorptionRaymarcher()

# 最後に、raysampler と raymarcher オブジェクトを使って 
# volumetric レンダラのインスタンスを作成します。
renderer = VolumeRenderer(
    raysampler=raysampler, raymarcher=raymarcher,
)

3. ボリュームデータの読み込み

次に、シーンのボリュームモデルをインスタンス化します。各ボクセルは、ボクセルのRGBカラーを表す3Dベクトルと、ボクセルの不透明度を表す密度スカラー（[0-1]の間の範囲、高いほど不透明）で記述されます。

npz形式で保存されたボリュームデータを読み込み、tensorに変換し、Volumesモデルに渡すために並び替えを行います。

In [4]:
from pprint import pprint
bonsai = np.load('../../../dataset/volumedata/Bonsai1.npz')['arr_0']
# 最大値と最小値の確認
pprint("min={},max={}".format(bonsai.min(), bonsai.max()))
# ヒストグラムを表示
pprint(np.histogram(bonsai, bins=4))
# tensorに変換
bonsai = torch.from_numpy(bonsai.astype(np.float32)).clone()
# Volumesに使用するために並べ替え
bonsai = bonsai.permute(3, 1, 2, 0)
print(bonsai.shape)

'min=0.0,max=0.0624847412109375'
(array([80528822,   696895,    38543,      380]),
 array([0.        , 0.01562119, 0.03124237, 0.04686356, 0.06248474],
      dtype=float32))
torch.Size([1, 512, 512, 310])

4. カラー情報の設定

ボクセルのRGBカラーを生成します。伝達関数のようにボリュームデータの密度を使って色を決定した後、Tensorに変換しています。

In [5]:
# 色情報
colors = np.zeros((3, 512, 512, 310))
# 伝達関数もどき
colors[2,:]=np.where((0.0006 <= bonsai) & (bonsai < 0.0625), 1.0, bonsai)
colors[1,:]=np.where((0.0006 <= bonsai) & (bonsai < 0.025), 1.0, bonsai)
colors[0,:]=np.where((bonsai < 0.01), bonsai/0.00025, bonsai)
# tensorに変換
colors = torch.from_numpy(colors.astype(np.float32)).clone()
bonsai = bonsai * 2.0
print(colors.shape)

torch.Size([3, 512, 512, 310])

5. Volumesデータの生成

ボリュームデータの密度と色情報を渡して、Volumesのインスタンスを生成します。

In [6]:
volume_size=512
volumes = Volumes(
            densities = bonsai[None].expand(
                1, *bonsai.shape),
            features = colors[None].expand(
                1, *colors.shape),
            voxel_size=volume_extent_world / volume_size,
        ).to(device)

6. FovPerspectiveCamerasの生成

回転行列・平行移動成分を設定してカメラを作成します。

In [7]:
from tqdm import tqdm as tqdm
# (nframes, 3)のTensorを作って０初期化する
logRs = torch.zeros(1, 3, device=device)
# ステップの数を指定した数列
# torch.linspace(start, end, num_step)という形式で作ります。
# 0～2πをn_frames分割, y軸周りの回転角度を設定？
logRs[:, 1] = torch.linspace(2.0, 2.0 * 3.14, 1, device=device)
# logRsを3x3の回転行列に変換
Rs = so3_exponential_map(logRs)
# 平行移動成分（nframes, 3)
Ts = torch.zeros(1, 3, device=device)
# Z軸方向にvolume_extent_world移動した位置
Ts[:, 2] = volume_extent_world
frames = []
# 回転行列と平行移動成分を取得
for R, T in zip(Rs, Ts):
    print(R.shape)
    print(R[None].shape)
    print(T[None].shape)
    # カメラの作成
    camera = FoVPerspectiveCameras(
        # [1,3,3]
        R=R[None], 
        # [1,3]
        T=T[None], 
        znear = 0.1,
        zfar = 100.0,
        aspect_ratio = 1.0,
        fov = 60.0,
        device=device,
    )

torch.Size([3, 3])
torch.Size([1, 3, 3])
torch.Size([1, 3])

7. rendererでレンダリングして結果を表示

rendererにカメラとボリュームデータを渡して可視化します。

In [8]:
frame = renderer(cameras=camera, volumes=volumes)[0]
frames.append(frame[..., :3].clamp(0.0, 1.0))
ret = torch.cat(frames)
ret = ret.clamp(0., 1.).cpu().numpy()
#貼り付け
plt.imshow(ret[0])
#表示
plt.show()
#保存
pilImg = Image.fromarray(np.uint8(ret[0]*255.0))
pilImg.save('result.png')

8. 複数視点からのボリュームの可視化

最後に、ボリュームのy軸を中心に回転する複数の視点からレンダリングします。

In [9]:
# AttributeError: 'FloatProgress' object has no attribute 'style'
# 上記のエラーが出たので下記を追加しました。
from tqdm import tqdm as tqdm
    
def generate_rotating_volume(volume_model, n_frames = 50):
    # (nframes, 3)のTensorを作って０初期化する
    logRs = torch.zeros(n_frames, 3, device=device)
    # ステップの数を指定した数列
    # torch.linspace(start, end, num_step)という形式で作ります。
    # 0～2πをn_frames分割, y軸周りの回転角度を設定？
    logRs[:, 1] = torch.linspace(0.0, 2.0 * 3.14, n_frames, device=device)
    # logRsを3x3の回転行列に変換
    Rs = so3_exponential_map(logRs)
    # 平行移動成分（nframes, 3)
    Ts = torch.zeros(n_frames, 3, device=device)
    # Z軸方向に2.7移動した位置
    Ts[:, 2] = volume_extent_world
    frames = []
    print('Generating rotating volume ...')
    # 回転行列と平行移動成分を取得
    for R, T in zip(tqdm(Rs), Ts):
        # カメラの作成
        camera = FoVPerspectiveCameras(
            # [1,3,3]
            R=R[None], 
            # [1,3]
            T=T[None], 
            znear = 0.1,
            zfar = 100.0,
            aspect_ratio = 1.0,
            fov = 60.0,
            device=device,
        )
        # volume_modelにカメラを渡してレンダリングした結果を0-1にクリップして結果をframesに追加
        frames.append(renderer(cameras=camera, volumes=volumes)[0][..., :3].clamp(0.0, 1.0))
    return torch.cat(frames)
    
rotating_volume_frames = generate_rotating_volume(volumes, n_frames=7*4)

image_grid(rotating_volume_frames.clamp(0., 1.).cpu().numpy(), rows=4, cols=7, rgb=True, fill=True)
plt.show()

  7%|▋         | 2/28 [00:00<00:02, 11.91it/s]
Generating rotating volume ...

100%|██████████| 28/28 [00:02<00:00, 11.25it/s]

９. 結論

レンダリングは、NDCGridRaysamplerとEmissionAbsorptionRaymarcherで構成されたPyTorch3Dのボリュームレンダラを用いて行いました。