# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.

レイマーチによるボリュームのフィッティング

このチュートリアルでは、微分可能なボリュームレンダリングを使用して、複数視点からの画像に合わせてボリュームをフィッティングさせる方法を説明します。

具体的には、以下の方法を説明します：

1. 微分可能なボリューメトリックレンダラーを作成します。
2. ボリュームモデルを作成します（Volumes クラスの使用方法を含む）。
3. 微分可能なボリュメトリックレンダラーを使用して、画像に基づいてボリュームをフィットします。
4. 予測されたボリュームを可視化します。

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

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

#!/usr/bin/pip3 install torch
#import sys
#import torch
#if torch.__version__=='1.6.0+cu101' and sys.platform.startswith('linux'):
#    !/usr/bin/pip3 install pytorch3d
#else:
#    !/usr/bin/pip3 install 'git+https://github.com/facebookresearch/pytorch3d.git@stable'

# 自分の場合はpython3.8, cuda11.2の環境なのと、condaは入れていないのでソースコードからインストールしました。
# git clone https://github.com/facebookresearch/pytorch3d.git
# cd pytorch3d && pip3 install -e .

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")

1. シーンの画像とマスクを生成する

次のセルは学習データを生成しています。fit_textured_mesh.ipynb チュートリアルの牛メッシュを複数の視点からレンダリングして返します：

1. 牛メッシュレンダラーによって生成された画像とシルエットのテンソルのバッチ
2. 各レンダリングに対応するカメラのセット

注意：このチュートリアルではvolumetricレンダリングの詳細を説明することを目的としているため、generate_cow_renders関数で実装されているメッシュレンダリングがどのように動作するかは説明していません。メッシュレンダリングの詳しい説明は fit_textured_mesh.ipynb を参照してください。

target_cameras, target_images, target_silhouettes = generate_cow_renders(num_views=40)
print(f'Generated {len(target_images)} images/silhouettes/cameras.')

Generated 40 images/silhouettes/cameras.

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

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

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

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

# render_size は、レンダリングされた画像の両側のサイズをピクセル単位で指定します。
# これをターゲット画像と同じサイズに設定します。
# つまり、正解画像と同じサイズでレンダリングします。
render_size = target_images.shape[1]

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

# 1) raysamplerのインスタンスを作成します。
# ここでは、NDCGridRaysampler を用いて、pytorch3d の座標規則に従った長方形の
# 画像グリッドを生成しています。128^3のボリュームを使用するので、n_pts_per_ray=150とします。
# さらに、カメラ平面から0.1ユニット以内には面がないので、min_depth=0.1とします。
raysampler = NDCGridRaysampler(
    image_width=render_size,
    image_height=render_size,
    n_pts_per_ray=150,
    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]の間の範囲、高いほど不透明）で記述されます。

濃度と色の範囲が[0-1]の間であることを確実にするために、私たちはボリュームの色と濃度の両方を対数空間で表現します。モデルの順伝搬関数の間、ログ空間の値はシグモイド関数に渡され、ログ空間の値が正しい範囲になるようになります。

さらに、VolumeModel にはレンダラオブジェクトが含まれています。このオブジェクトは最適化中も変更されません。

このセルでは、レンダリングされたカラーとマスクの間の不一致を計算するhuber損失関数も定義しています。

class VolumeModel(torch.nn.Module):
    def __init__(self, renderer, volume_size=[64] * 3, voxel_size=0.1):
        super().__init__()
        # torch.sigmoid(self.log_densities) を評価すると、密度がゼロに近いものが得られます。
        self.log_densities = torch.nn.Parameter(-4.0 * torch.ones(1, *volume_size))
        # torch.sigmoid(self.log_colors)を評価すると、どこもかしこも中性的な灰色になります。
        self.log_colors = torch.nn.Parameter(torch.zeros(3, *volume_size))
        self._voxel_size = voxel_size
        # レンダラモジュールも格納しておきます。
        self._renderer = renderer
        # 学習に使ったボリュームデータ参照用（追加）
        self.volumes = None
        
    def forward(self, cameras):
        batch_size = cameras.R.shape[0]

        # ログ空間の値を密度/色に変換します。
        densities = torch.sigmoid(self.log_densities)
        colors = torch.sigmoid(self.log_colors)
        
        # Volumesオブジェクトのインスタンスを作成し、密度と色が正しく展開された
        # batch_size-time であることを確認します。
        self.volumes = Volumes(
            densities = densities[None].expand(
                batch_size, *self.log_densities.shape),
            features = colors[None].expand(
                batch_size, *self.log_colors.shape),
            voxel_size=self._voxel_size,
        )
        
        # カメラとボリュームが与えられたら、レンダリングを実行し、最初の出力値のみを
        # 返します（2番目の出力は、サンプリングされた光線の表現で、この目的のために省略することができます）。
        return self._renderer(cameras=cameras, volumes=self.volumes)[0]
    
# レンダリングされたシルエットと色の間の滑らかなL1（huber）損失を評価するためのヘルパー関数です。
def huber(x, y, scaling=0.1):
    diff_sq = (x - y) ** 2
    loss = ((1 + diff_sq / (scaling**2)).clamp(1e-4).sqrt() - 1) * float(scaling)
    return loss

4. ボリュームのフィッティング

ここでは、微分可能なレンダリングを用いてボリュームフィッティングを行います。

ボリュームをフィットさせるために、ターゲットカメラの視点からレンダリングを行い、結果として得られたレンダリングと観測されたターゲット画像とターゲットシルエットを比較します。

比較は、target_images/rendered_imagesとtarget_silhouettes/rendered_silhouettesの対応するペア間の平均huber（smooth-l1）誤差を評価することによって行われます。

# まず、関連するすべての変数を正しいデバイスに移動させます。
target_cameras = target_cameras.to(device)
target_images = target_images.to(device)
target_silhouettes = target_silhouettes.to(device)

# ボリュームモデルをインスタンス化します。1辺のサイズが128の立方体のボリュームを使用します。
# ボリュームの各ボクセルのサイズは、volume_extent_world / volume_size s.t.に設定されています。
# ボリュームは、各辺のサイズが3に等しい(0, 0, 0)を中心とした3Dバウンディングボックスで囲まれた空間を表します。
volume_size = 128
volume_model = VolumeModel(
    renderer,
    volume_size=[volume_size] * 3, 
    voxel_size = volume_extent_world / volume_size,
).to(device)

# Adam optimizerのインスタンスを作成します。マスター学習率を0.1に設定します。
lr = 0.1
optimizer = torch.optim.Adam(volume_model.parameters(), lr=lr)

# Adamの反復処理を300回行い、各ミニバッチで10個のランダム画像をサンプリングします。
batch_size = 10
n_iter = 300
for iteration in range(n_iter):

    # 反復の最後の75%に達した場合は、Optimizerの学習率を0.1倍します。
    if iteration == round(n_iter * 0.75):
        print('Decreasing LR 10-fold ...')
        optimizer = torch.optim.Adam(
            volume_model.parameters(), lr=lr * 0.1
        )
    
    # オプティマイザの勾配をゼロ初期化します。
    optimizer.zero_grad()
    
    # ランダムなバッチ番号をサンプリングします。
    batch_idx = torch.randperm(len(target_cameras))[:batch_size]
    
    # カメラ座標のミニバッチもサンプリングします。
    batch_cameras = FoVPerspectiveCameras(
        R = target_cameras.R[batch_idx], 
        T = target_cameras.T[batch_idx], 
        znear = target_cameras.znear[batch_idx],
        zfar = target_cameras.zfar[batch_idx],
        aspect_ratio = target_cameras.aspect_ratio[batch_idx],
        fov = target_cameras.fov[batch_idx],
        device = device,
    )
    
    # ボリュームモデルを評価します。
    rendered_images, rendered_silhouettes = volume_model(
        batch_cameras
    ).split([3, 1], dim=-1)
    
    # 予測されたマスクとターゲットのシルエットの間の平均hubor損失としてシルエット誤差を計算します。
    sil_err = huber(
        rendered_silhouettes[..., 0], target_silhouettes[batch_idx],
    ).abs().mean()

    # レンダリングされた色とターゲットの正解画像の間の平均フーバー損失として色誤差を計算します。
    color_err = huber(
        rendered_images, target_images[batch_idx],
    ).abs().mean()
    
    # 最適化損失は単純に色の誤差とシルエット誤差の合計とします。
    loss = color_err + sil_err 
    
    # 現在のロスを表示します。
    if iteration % 10 == 0:
        print(
            f'Iteration {iteration:05d}:'
            + f' color_err = {float(color_err):1.2e}'
            + f' mask_err = {float(sil_err):1.2e}'
        )
    
    # 最適化ステップを実行します。
    loss.backward()
    optimizer.step()
    
    # ４０試行ごとに結果を可視化します
    if iteration % 40 == 0:
        # バッチからランダムに１つの要素だけ選んで可視化します
        im_show_idx = int(torch.randint(low=0, high=batch_size, size=(1,)))
        fig, ax = plt.subplots(2, 2, figsize=(10, 10))
        ax = ax.ravel()
        clamp_and_detach = lambda x: x.clamp(0.0, 1.0).cpu().detach().numpy()
        ax[0].imshow(clamp_and_detach(rendered_images[im_show_idx]))
        ax[1].imshow(clamp_and_detach(target_images[batch_idx[im_show_idx], ..., :3]))
        ax[2].imshow(clamp_and_detach(rendered_silhouettes[im_show_idx, ..., 0]))
        ax[3].imshow(clamp_and_detach(target_silhouettes[batch_idx[im_show_idx]]))
        for ax_, title_ in zip(
            ax, 
            ("rendered image", "target image", "rendered silhouette", "target silhouette")
        ):
            ax_.grid("off")
            ax_.axis("off")
            ax_.set_title(title_)
        fig.canvas.draw(); fig.show()
        display.clear_output(wait=True)
        display.display(fig)

Iteration 00290: color_err = 3.81e-03 mask_err = 9.96e-03

5. 最適化されたボリュームの可視化

最後に、ボリュームのy軸を中心に回転する複数の視点からレンダリングすることで、最適化されたボリュームを可視化します。

# AttributeError: 'FloatProgress' object has no attribute 'style'
# 上記のエラーが出たので下記を追加しました。
from tqdm import tqdm as tqdm
    
def generate_rotating_volume(volume_model, n_frames = 50):
    logRs = torch.zeros(n_frames, 3, device=device)
    logRs[:, 1] = torch.linspace(0.0, 2.0 * 3.14, n_frames, device=device)
    Rs = so3_exponential_map(logRs)
    Ts = torch.zeros(n_frames, 3, device=device)
    Ts[:, 2] = 2.7
    frames = []
    print('Generating rotating volume ...')
    for R, T in zip(tqdm(Rs), Ts):
        camera = FoVPerspectiveCameras(
            R=R[None], 
            T=T[None], 
            znear = target_cameras.znear[0],
            zfar = target_cameras.zfar[0],
            aspect_ratio = target_cameras.aspect_ratio[0],
            fov = target_cameras.fov[0],
            device=device,
        )
        frames.append(volume_model(camera)[..., :3].clamp(0.0, 1.0))
    return torch.cat(frames)
    
with torch.no_grad():
    rotating_volume_frames = generate_rotating_volume(volume_model, n_frames=7*4)

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

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

Generating rotating volume ...

# バッチサイズ分のボリュームデータ？
# 密度のサイズ（不透明度？）
print(volume_model.volumes.densities().size())
# 特徴（カラー？）
print(volume_model.volumes.features().size())
# グリッドサイズの表示
print(volume_model.volumes.get_grid_sizes())

# 最適化された結果のデータ？
print(torch.sigmoid(volume_model.log_densities).size())
print(torch.sigmoid(volume_model.log_colors).size())
print(volume_model._voxel_size)

# 0-1にスケーリングするためにシグモイド関数を通してchannelが最後に来るように並べ替える
d = torch.sigmoid(volume_model.log_densities).permute(1,2,3,0).to('cpu').detach().numpy().copy()
color = torch.sigmoid(volume_model.log_colors).permute(1,2,3,0).to('cpu').detach().numpy().copy()
print(d.shape)
print(color.shape)
# rgbとdensityを結合する
rgba = np.concatenate([color,d], 3)
print(rgba.shape)
# npzで保存
np.savez('cow.npz', rgba)

torch.Size([1, 1, 128, 128, 128])
torch.Size([1, 3, 128, 128, 128])
tensor([[128, 128, 128]], device='cuda:0')
torch.Size([1, 128, 128, 128])
torch.Size([3, 128, 128, 128])
0.0234375
(128, 128, 128, 1)
(128, 128, 128, 3)
(128, 128, 128, 4)

6. 結論

このチュートリアルでは、既知の視点からのボリュームのレンダリングが各視点の観測画像と一致するように、シーンの3Dボリューム表現を最適化する方法を示しました。レンダリングは、NDCGridRaysamplerとEmissionAbsorptionRaymarcherで構成されたPyTorch3Dのボリュームレンダラを用いて行いました。