In [1]:
# .objファイルの読み込み
# レンダラの設定
# メッシュのレンダリング
# 光源やカメラ位置などのレンダリング設定の変更
# 異なる視点でメッシュをバッチ処理

import os
import torch
import matplotlib.pyplot as plt
from skimage.io import imread

# Util function for loading meshes
from pytorch3d.io import load_obj,load_objs_as_meshes

# Data structures and functions for rendering
from pytorch3d.structures import Meshes, Textures
from pytorch3d.renderer import (
    look_at_view_transform,
    OpenGLPerspectiveCameras, 
    PointLights, 
    DirectionalLights, 
    Materials, 
    RasterizationSettings, 
    MeshRenderer, 
    MeshRasterizer,  
    TexturedPhongShader
)

# add path for demo utils functions 
import sys
import os
sys.path.append(os.path.abspath(''))
In [2]:
# 予めplot_image_grid.pyをダウンロードして同じディレクトリに置いています
#!wget https://raw.githubusercontent.com/facebookresearch/pytorch3d/master/docs/tutorials/utils/plot_image_grid.py
from plot_image_grid import image_grid
In [3]:
# データは事前にサーバに置いている

# Setup
device = torch.device("cuda:0")
torch.cuda.set_device(device)

# Set paths
DATA_DIR = "./data"
#obj_filename = os.path.join(DATA_DIR, "cow_mesh/cow.obj")
obj_filename = os.path.join(DATA_DIR, "tda_append/tda_append.obj")

# Load the obj and ignore the textures and materials.
# load_obj関数を使って.objのメッシュを読み込むが、この関数は1枚のテクスチャを持つメッシュにしか使えない？
#mesh = load_objs_as_meshes([obj_filename], device=device)

verts, faces, aux = load_obj(obj_filename)
verts_uvs = aux.verts_uvs[None,...]
faces_uvs = faces.textures_idx[None,...]
tex_maps = aux.texture_images
texture_image = list(tex_maps.values())[0]
texture_image = texture_image[None, ...]
tex = Textures(verts_uvs=verts_uvs, faces_uvs=faces_uvs, maps=texture_image)

center = verts.mean(0)
verts = verts - center
scale = max(verts.abs().max(0)[0])
verts = 2.0*verts / scale

meshes = Meshes(verts=[verts], faces=[faces.verts_idx], textures=tex)

# Meshesには、頂点と面のデータを表現する3つの異なる方法がある
# list: 他の表現に変換するための開始点として使用される
# Padded: バッチの次元を持つ
# Packed: バッチの次元を持たない。Padded表現へのインデックスに使用する補助変数あり
# 今回は、メッシュのテクスチャデータをバッチの次元を持つPaddedで取得
#texture_image=mesh.textures.maps_padded()
# バッチの次元を持つので、[1, 1024, 1024, 3]
print(texture_image.shape)

torch.Size([1, 2048, 7216, 3])
In [4]:
# テクスチャマップの表示
plt.figure(figsize=(7,7))
# squeeze(): size 1の入力を取り除く[1, 1024, 1024, 3] -> [1024, 1024, 3]
# cpu(): オブジェクトのコピーをCPUメモリに返す
# numpy() -> numpy.ndarray: tensorをNumPyのndarrayに変換する
plt.imshow(texture_image.squeeze().cpu().numpy())
plt.grid("off")
plt.axis('off')
Out[4]:
(-0.5, 7215.5, 2047.5, -0.5)
In [5]:
# 2. レンダラの作成
# OpenGL perspective cameraの初期化
R, T = look_at_view_transform(2.7, 10, 180)
cameras = OpenGLPerspectiveCameras(device=device, R=R, T=T)

# ラスタライザとシェーダの設定。512x512で出力。
# faces_per_pixel=1: 各ピクセルに保存する面の数。1だと一番手前の面の情報だけ保持して、2以上だと近い順に保持？
# blur_radius=0.0: [0, 2]の間で設定するFloat値の距離で、ブラーを掛ける範囲を設定？
# bin_size=0: coarse-to-fine rasterizationに使うビンのサイズらしいです。coarse-to-fine:粗密。最初は粗く、重要な部分だけ細かくラスタライズ？
raster_settings = RasterizationSettings(
    image_size=512, 
    blur_radius=0.0, 
    faces_per_pixel=1, 
    bin_size=0
)

# ポイントライトの設置
lights = PointLights(device=device, location=[[1.0, 1.0, 2.0]])

# フォンのレンダラの作成。フォンシェーダはテクスチャのuv座標を補間します。
renderer = MeshRenderer(
    rasterizer=MeshRasterizer(
        cameras=cameras, 
        raster_settings=raster_settings
    ),
    shader=TexturedPhongShader(
        device=device, 
        cameras=cameras,
        lights=lights
    )
)

# 3. render the mesh
mesh = meshes.to(device)
images = renderer(mesh)
plt.figure(figsize=(10, 10))
plt.imshow(images[0, ..., :3].cpu().numpy())
plt.grid("off")
plt.axis("off")
Out[5]:
(-0.5, 511.5, 511.5, -0.5)
In [6]:
# 4. 光源をオブジェクトの後ろに移動させて再描画
print("before:", lights.location.shape, lights.location)
lights.location = torch.tensor([0.0, 0.0, -2.0], device=device)[None]
print("after:", lights.location)

images = renderer(mesh, lights=lights)

plt.figure(figsize=(10, 10))
plt.imshow(images[0, ..., :3].cpu().numpy())
plt.grid("off")
plt.axis("off")

before: torch.Size([1, 3]) tensor([[1., 1., 2.]], device='cuda:0')
after: tensor([[ 0.,  0., -2.]], device='cuda:0')
Out[6]:
(-0.5, 511.5, 511.5, -0.5)
In [7]:
# 5. オブジェクトの回転、材質もしくは光源属性の変更
# カメラのアングルを変更して物体を回転
R, T = look_at_view_transform(dist=2.7, elev=30, azim=150)
cameras = OpenGLPerspectiveCameras(device=device, R=R, T=T)

# 光源の位置を元に戻す
lights.location = torch.tensor([[5.0, 5.0, 2.0]], device=device)

# 鏡面反射色を緑色にして、shininessを変更
materials = Materials(
  device=device,
  specular_color=[[0.0,1.0,0.0]],
  shininess=10.0
)

# 再描画
images = renderer(mesh, lights=lights, materials=materials, cameras=cameras)

# 表示
plt.figure(figsize=(10, 10))
plt.imshow(images[0, ..., :3].cpu().numpy())
plt.grid("off")
plt.axis("off")
Out[7]:
(-0.5, 511.5, 511.5, -0.5)
In [8]:
# 6. Batched Rendering
# PyTorch3d APIのコアデザインの一つとして、バッチ処理をサポートしていることが挙げられます。
# レンダラは、１つのforward passで、バッチ分の出力画像を得ることができます。

# バッチサイズの設定 - レンダリングする異なる視点の数
batch_size = 20

# バッチサイズ分のメッシュを作る必要があります。
# `extend`関数を使うと、簡単に作成できます。テクスチャも拡張してくれます。
meshes = mesh.extend(batch_size)
print(meshes)

# Get a batch of viewing angles.
elev = torch.linspace(0, 360, batch_size)
azim = torch.linspace(0, 360, batch_size)
print("elev.shape=", elev.shape)
print("azim.shape=", azim.shape)

# camerasのヘルパー関数は、入力とブロードキャストの混合型をサポートしています。
# そのため、距離は2.7で固定して、角度だけ変更するようなことも可能です。
R, T = look_at_view_transform(dist=2.7, elev=elev, azim=azim)
cameras = OpenGLPerspectiveCameras(device=device, R=R, T=T)
print("R.shape=", R.shape)
print("T.shape=", T.shape)

# 光源の位置を戻す
lights.location = torch.tensor([[1.0, 1.0, -5.0]], device=device)

# レンダリング
images = renderer(meshes, cameras=cameras, lights=lights)

image_grid(images.cpu().numpy(), rows=4, cols=5, rgb=True)

<pytorch3d.structures.meshes.Meshes object at 0x7f58d0710e80>
elev.shape= torch.Size([20])
azim.shape= torch.Size([20])
R.shape= torch.Size([20, 3, 3])
T.shape= torch.Size([20, 3])
In [ ]: