PyTorch3D：3Dコンピュータービジョンライブラリ

ついにFacebook Research公式から3Dディープラーニング用のPyTorchベースのライブラリが登場した。
ライセンスはBSD-3-Clauseライセンス。

PyTorch3D

PyTorch3Dは、PyTorchを使用した3Dコンピュータービジョン研究のための効率的で再利用可能なコンポーネントを提供します。

主要な機能は以下の通りです：

三角形Meshを保存・操作するためのデータ構造

三角形Mesh上での効率的な操作(射影変換、グラフ畳み込み、サンプリング、損失関数)

微分可能なMeshレンダラー

PyTorch3Dは3Dデータを予測・操作するためにディープラーニング手法とスムーズに統合できるよう設計されています。そのため、PyTorch3Dでの操作は全て以下に対応しています：

PyTorchのtensorを使用した実装

異種データのミニバッチ処理

微分可能

GPUによる処理の高速化

Facebook Research内では、PyTorch3DをMesh R-CNNなどの研究プロジェクトで使用しています。

ソースコードはGitHubにあり、ドキュメントはこちら。
INSTALL.mdを読むと、まだWindowsには対応していないみたいですね。

PyTorch3DまだWindows対応していないので、MacやLinux環境ないけどサクッと試したいだけの人はGoogleColabがおススメ。Notebookの先頭に以下でインスコすれば動きます。
!pip install 'git+https://t.co/8Cs3tKNij8'

— Teppei Kurita (@kuritateppei) February 10, 2020

追記：Facebookの公式ブログの記事も公開された↓
https://ai.facebook.com/blog/-introducing-pytorch3d-an-open-source-library-for-3d-deep-learning/

We just released PyTorch3D, a new toolkit for researchers and engineers that’s fast and modular for 3D deep learning research: https://t.co/RmMSXJoWCm pic.twitter.com/H75Ai9gH2I

— Facebook AI (@facebookai) February 6, 2020

言及されているMesh R-CNNもソースコードが公開されている↓
https://github.com/facebookresearch/meshrcnn
https://ai.facebook.com/blog/pushing-state-of-the-art-in-3d-content-understanding/

個人的に興味があるのはDifferentiable Rendering周りのAPI↓

アーキテクチャの概要

レンダラーは、モジュール式で拡張可能に設計されており、全ての入力に対してバッチ処理、勾配(微分)をサポートしています。以下の図はレンダリングパイプラインの全てのコンポーネントを示しています。

Fragments
rasterizerは4つの出力tensorを以下のような名前付きtupleで返します。

pix_to_face：(N, image_size, image_size, faces_per_pixel)形状のlong型tensor。画像の各ピクセルに重なるfaceのインデックスを(packed faceで)指定します。

zbuf：(N, image_size, image_size, faces_per_pixel)形状のfloat型tensor。各ピクセルで最も近いfaceのworld座標系でのz座標をz値の昇順で提供します。

bary_coords：(N, image_size, image_size, faces_per_pixel, 3)形状のfloat型tensor。各ピクセルで最も近いfaceの重心座標をNDC単位でz値の昇順で提供します。

pix_dists：(N, image_size, image_size, faces_per_pixel)形状のflaot型tensor。ピクセルに最も近い各点のx/y平面での符号付きユークリッド距離(NDC単位)を提供します。

パイプラインの各コンポーネントの詳細については、レンダラーAPIリファレンスを参照してください。

注意：differentiable renderer APIは試験的なものであり、仕様は今後変更される可能性があります。

座標変換の規則

レンダリングでは、座標空間を異なる座標空間へ変換する処理を何度か行う必要があります：world空間、view/camera空間、NDC空間、screen空間があります。これら各座標空間でカメラがどこに位置しているか、x, y, z軸がどのような配置関係でどのような範囲の値を取るのかを知ることが重要です。以下の図は、PyTorch3dが使用する規則を示しています。

注意：PyTorch3d vs OpenGL
PyTorch3dはOpenGLに似せるようにしていますが、OpenGLの左手系のNDC座標系と違い、PyTorch3dのNDC座標系は右手系となっています。(射影行列によって利き手が切り替わります)
OpenGLでは、camera空間の原点にあるカメラが奥行きをz軸マイナス方向に捉えるのに対し、PyTorch3dではNDC空間のz軸プラス方向に奥行きを捉えます。

シンプルなレンダラー

PyTorch3dのレンダラーはrasterizerとshaderで構成されています。レンダラーを作成するには以下のように記述します。
# Imports
from pytorch3d.renderer import (
    OpenGLPerspectiveCameras, look_at_view_transform,
    RasterizationSettings, BlendParams,
    MeshRenderer, MeshRasterizer, PhongShader
)

# Initialize an OpenGL perspective camera.
R, T = look_at_view_transform(2.7, 10, 20)
cameras = OpenGLPerspectiveCameras(device=device, R=R, T=T)

# Define the settings for rasterization and shading. Here we set the output image to be of size
# 512x512. As we are rendering images for visualization purposes only we will set faces_per_pixel=1
# and blur_radius=0.0. Refer to rasterize_meshes.py for explanations of these parameters.
raster_settings = RasterizationSettings(
    image_size=512,
    blur_radius=0.0,
    faces_per_pixel=1,
    bin_size=0
)

# Create a phong renderer by composing a rasterizer and a shader. Here we can use a predefined
# PhongShader, passing in the device on which to initialize the default parameters
renderer = MeshRenderer(
    rasterizer=MeshRasterizer(cameras=cameras, raster_settings=raster_settings),
    shader=PhongShader(device=device, cameras=cameras)
)