# Copyright (c) Meta Platforms, Inc. and affiliates. All rights reserved.
在本教程中,我们将学习如何使用可微渲染,根据参考图像获取摄像机的[x, y, z]位置。
我们将首先初始化一个渲染器,并为相机设置初始位置。然后,我们将使用它生成图像,计算与参考图像的损失,最后反向传播整个管道以更新相机的位置。
本教程展示了如何
.obj
文件加载网格Camera
、Shader
和Renderer
,确保已安装torch
和torchvision
。如果未安装pytorch3d
,请使用以下单元格安装它
import os
import sys
import torch
need_pytorch3d=False
try:
import pytorch3d
except ModuleNotFoundError:
need_pytorch3d=True
if need_pytorch3d:
if torch.__version__.startswith("2.2.") and sys.platform.startswith("linux"):
# We try to install PyTorch3D via a released wheel.
pyt_version_str=torch.__version__.split("+")[0].replace(".", "")
version_str="".join([
f"py3{sys.version_info.minor}_cu",
torch.version.cuda.replace(".",""),
f"_pyt{pyt_version_str}"
])
!pip install fvcore iopath
!pip install --no-index --no-cache-dir pytorch3d -f https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/{version_str}/download.html
else:
# We try to install PyTorch3D from source.
!pip install 'git+https://github.com/facebookresearch/pytorch3d.git@stable'
import os
import torch
import numpy as np
from tqdm.notebook import tqdm
import imageio
import torch.nn as nn
import torch.nn.functional as F
import matplotlib.pyplot as plt
from skimage import img_as_ubyte
# io utils
from pytorch3d.io import load_obj
# datastructures
from pytorch3d.structures import Meshes
# 3D transformations functions
from pytorch3d.transforms import Rotate, Translate
# rendering components
from pytorch3d.renderer import (
FoVPerspectiveCameras, look_at_view_transform, look_at_rotation,
RasterizationSettings, MeshRenderer, MeshRasterizer, BlendParams,
SoftSilhouetteShader, HardPhongShader, PointLights, TexturesVertex,
)
我们将加载一个obj文件并创建一个Meshes对象。Meshes是PyTorch3D中提供的一种独特的数据结构,用于处理不同大小的网格批次。它具有几个有用的类方法,这些方法用于渲染管道。
如果您在克隆PyTorch3D存储库后在本地运行此笔记本,则网格将已可用。如果使用Google Colab,请获取网格并将其保存在data/
路径下
!mkdir -p data
!wget -P data https://dl.fbaipublicfiles.com/pytorch3d/data/teapot/teapot.obj
# Set the cuda device
if torch.cuda.is_available():
device = torch.device("cuda:0")
torch.cuda.set_device(device)
else:
device = torch.device("cpu")
# Load the obj and ignore the textures and materials.
verts, faces_idx, _ = load_obj("./data/teapot.obj")
faces = faces_idx.verts_idx
# Initialize each vertex to be white in color.
verts_rgb = torch.ones_like(verts)[None] # (1, V, 3)
textures = TexturesVertex(verts_features=verts_rgb.to(device))
# Create a Meshes object for the teapot. Here we have only one mesh in the batch.
teapot_mesh = Meshes(
verts=[verts.to(device)],
faces=[faces.to(device)],
textures=textures
)
PyTorch3D中的渲染器由光栅化器和着色器组成,每个着色器都包含许多子组件,例如相机(正交/透视)。在这里,我们初始化其中一些组件,并为其余组件使用默认值。
为了优化相机位置,我们将使用一个渲染器,该渲染器仅生成对象的轮廓,并且不应用任何照明或阴影。我们还将初始化另一个应用完整Phong着色的渲染器,并将其用于可视化输出。
# Initialize a perspective camera.
cameras = FoVPerspectiveCameras(device=device)
# To blend the 100 faces we set a few parameters which control the opacity and the sharpness of
# edges. Refer to blending.py for more details.
blend_params = BlendParams(sigma=1e-4, gamma=1e-4)
# Define the settings for rasterization and shading. Here we set the output image to be of size
# 256x256. To form the blended image we use 100 faces for each pixel. We also set bin_size and max_faces_per_bin to None which ensure that
# the faster coarse-to-fine rasterization method is used. Refer to rasterize_meshes.py for
# explanations of these parameters. Refer to docs/notes/renderer.md for an explanation of
# the difference between naive and coarse-to-fine rasterization.
raster_settings = RasterizationSettings(
image_size=256,
blur_radius=np.log(1. / 1e-4 - 1.) * blend_params.sigma,
faces_per_pixel=100,
)
# Create a silhouette mesh renderer by composing a rasterizer and a shader.
silhouette_renderer = MeshRenderer(
rasterizer=MeshRasterizer(
cameras=cameras,
raster_settings=raster_settings
),
shader=SoftSilhouetteShader(blend_params=blend_params)
)
# We will also create a Phong renderer. This is simpler and only needs to render one face per pixel.
raster_settings = RasterizationSettings(
image_size=256,
blur_radius=0.0,
faces_per_pixel=1,
)
# We can add a point light in front of the object.
lights = PointLights(device=device, location=((2.0, 2.0, -2.0),))
phong_renderer = MeshRenderer(
rasterizer=MeshRasterizer(
cameras=cameras,
raster_settings=raster_settings
),
shader=HardPhongShader(device=device, cameras=cameras, lights=lights)
)
我们将首先定位茶壶并生成图像。我们使用辅助函数将茶壶旋转到所需的视点。然后,我们可以使用渲染器生成图像。在这里,我们将使用这两个渲染器并可视化轮廓和完全着色的图像。
世界坐标系定义为+Y向上,+X向左,+Z向内。世界坐标系中的茶壶壶嘴指向左侧。
我们定义了一个相机,它位于正z轴上,因此可以看到壶嘴在右侧。
# Select the viewpoint using spherical angles
distance = 3 # distance from camera to the object
elevation = 50.0 # angle of elevation in degrees
azimuth = 0.0 # No rotation so the camera is positioned on the +Z axis.
# Get the position of the camera based on the spherical angles
R, T = look_at_view_transform(distance, elevation, azimuth, device=device)
# Render the teapot providing the values of R and T.
silhouette = silhouette_renderer(meshes_world=teapot_mesh, R=R, T=T)
image_ref = phong_renderer(meshes_world=teapot_mesh, R=R, T=T)
silhouette = silhouette.cpu().numpy()
image_ref = image_ref.cpu().numpy()
plt.figure(figsize=(10, 10))
plt.subplot(1, 2, 1)
plt.imshow(silhouette.squeeze()[..., 3]) # only plot the alpha channel of the RGBA image
plt.grid(False)
plt.subplot(1, 2, 2)
plt.imshow(image_ref.squeeze())
plt.grid(False)
在这里,我们创建一个简单的模型类并为相机位置初始化一个参数。
class Model(nn.Module):
def __init__(self, meshes, renderer, image_ref):
super().__init__()
self.meshes = meshes
self.device = meshes.device
self.renderer = renderer
# Get the silhouette of the reference RGB image by finding all non-white pixel values.
image_ref = torch.from_numpy((image_ref[..., :3].max(-1) != 1).astype(np.float32))
self.register_buffer('image_ref', image_ref)
# Create an optimizable parameter for the x, y, z position of the camera.
self.camera_position = nn.Parameter(
torch.from_numpy(np.array([3.0, 6.9, +2.5], dtype=np.float32)).to(meshes.device))
def forward(self):
# Render the image using the updated camera position. Based on the new position of the
# camera we calculate the rotation and translation matrices
R = look_at_rotation(self.camera_position[None, :], device=self.device) # (1, 3, 3)
T = -torch.bmm(R.transpose(1, 2), self.camera_position[None, :, None])[:, :, 0] # (1, 3)
image = self.renderer(meshes_world=self.meshes.clone(), R=R, T=T)
# Calculate the silhouette loss
loss = torch.sum((image[..., 3] - self.image_ref) ** 2)
return loss, image
现在,我们可以创建上面模型的一个实例并为相机位置参数设置一个优化器。
# We will save images periodically and compose them into a GIF.
filename_output = "./teapot_optimization_demo.gif"
writer = imageio.get_writer(filename_output, mode='I', duration=0.3)
# Initialize a model using the renderer, mesh and reference image
model = Model(meshes=teapot_mesh, renderer=silhouette_renderer, image_ref=image_ref).to(device)
# Create an optimizer. Here we are using Adam and we pass in the parameters of the model
optimizer = torch.optim.Adam(model.parameters(), lr=0.05)
plt.figure(figsize=(10, 10))
_, image_init = model()
plt.subplot(1, 2, 1)
plt.imshow(image_init.detach().squeeze().cpu().numpy()[..., 3])
plt.grid(False)
plt.title("Starting position")
plt.subplot(1, 2, 2)
plt.imshow(model.image_ref.cpu().numpy().squeeze())
plt.grid(False)
plt.title("Reference silhouette");
我们运行前向和后向传递的几个迭代,并在每10次迭代保存输出。完成后,查看./teapot_optimization_demo.gif
以获取优化过程的酷炫gif!
loop = tqdm(range(200))
for i in loop:
optimizer.zero_grad()
loss, _ = model()
loss.backward()
optimizer.step()
loop.set_description('Optimizing (loss %.4f)' % loss.data)
if loss.item() < 200:
break
# Save outputs to create a GIF.
if i % 10 == 0:
R = look_at_rotation(model.camera_position[None, :], device=model.device)
T = -torch.bmm(R.transpose(1, 2), model.camera_position[None, :, None])[:, :, 0] # (1, 3)
image = phong_renderer(meshes_world=model.meshes.clone(), R=R, T=T)
image = image[0, ..., :3].detach().squeeze().cpu().numpy()
image = img_as_ubyte(image)
writer.append_data(image)
plt.figure()
plt.imshow(image[..., :3])
plt.title("iter: %d, loss: %0.2f" % (i, loss.data))
plt.axis("off")
writer.close()
在本教程中,我们学习了如何从obj文件加载网格,初始化称为Meshes的PyTorch3D数据结构,设置由光栅化器和着色器组成的渲染器,设置包括模型和损失函数的优化循环,以及运行优化。