【小白深度教程 1.16】手把手教你使用 Pytorch3D(1)使用 3D 损失函数来拟合 Mesh
本教程展示了如何:
- 从
.obj文件中加载网格和纹理。 - 设置渲染器
- 渲染网格
- 调整渲染设置,例如光照和相机位置
- 使用 PyTorch3D API 的批处理功能从不同视角渲染网格
渲染带纹理的网格(Mesh)_files/3ee3bfdbbf8f445383ff034b8d04e017.gif)
1. 安装和导入模块
确保安装了“torch”和“torchvision”。如果’ pytorch3d '未安装,请使用以下命令安装:
pip install --no-index --no-cache-dir pytorch3d -f https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/{version_str}/download.html
或者从 github 源码安装:
pip install 'git+https://github.com/facebookresearch/pytorch3d.git@stable'
然后导入模块:
import os
import torch
import matplotlib.pyplot as plt
# Util function for loading meshes
from pytorch3d.io import load_objs_as_meshes, load_obj
# Data structures and functions for rendering
from pytorch3d.structures import Meshes
from pytorch3d.vis.plotly_vis import AxisArgs, plot_batch_individually, plot_scene
from pytorch3d.vis.texture_vis import texturesuv_image_matplotlib
from pytorch3d.renderer import (
look_at_view_transform,
FoVPerspectiveCameras,
PointLights,
DirectionalLights,
Materials,
RasterizationSettings,
MeshRenderer,
MeshRasterizer,
SoftPhongShader,
TexturesUV,
TexturesVertex
)
# add path for demo utils functions
import sys
import os
sys.path.append(os.path.abspath(''))
下载用于可视化的代码(如果使用 Colab 的话):
!wget https://raw.githubusercontent.com/facebookresearch/pytorch3d/main/docs/tutorials/utils/plot_image_grid.py
from plot_image_grid import image_grid
2. 加载网格和纹理文件
加载一个 .obj 文件及其关联的 .mtl 文件,并创建一个 Textures 和 Meshes 对象。
Meshes 是 PyTorch3D 提供的一种独特数据结构,用于处理不同大小的网格批次。
TexturesUV 是一个辅助数据结构,用于存储网格的顶点 UV 和纹理映射。
Meshes 有多个类方法,可在整个渲染管道中使用。
如果使用 Google Colab 运行此笔记本,请运行以下单元格以获取网格 obj 和纹理文件,并将其保存在路径 data/cow_mesh 中:
如果本地运行,数据已经保存在正确的路径中。
!mkdir -p data/cow_mesh
!wget -P data/cow_mesh https://dl.fbaipublicfiles.com/pytorch3d/data/cow_mesh/cow.obj
!wget -P data/cow_mesh https://dl.fbaipublicfiles.com/pytorch3d/data/cow_mesh/cow.mtl
!wget -P data/cow_mesh https://dl.fbaipublicfiles.com/pytorch3d/data/cow_mesh/cow_texture.png
# Setup
if torch.cuda.is_available():
device = torch.device("cuda:0")
torch.cuda.set_device(device)
else:
device = torch.device("cpu")
# Set paths
DATA_DIR = "./data"
obj_filename = os.path.join(DATA_DIR, "cow_mesh/cow.obj")
# Load obj file
mesh = load_objs_as_meshes([obj_filename], device=device)
3. 可视化纹理图(Texture Map)
plt.figure(figsize=(7,7))
texture_image=mesh.textures.maps_padded()
plt.imshow(texture_image.squeeze().cpu().numpy())
plt.axis("off");
渲染带纹理的网格(Mesh)_files/b65889e0d93e4659a9982456577cd424.png)
PyTorch3D有一个内置的方法来查看纹理图,以及对应的顶点。
plt.figure(figsize=(7,7))
texturesuv_image_matplotlib(mesh.textures, subsample=None)
plt.axis("off")
渲染带纹理的网格(Mesh)_files/0d875608357e4e3793e473823e41248b.png)
4. 创建渲染器(Renderer)
在 PyTorch3D 中,渲染器由 光栅化器 和 着色器 组成,每个组件都有一些子组件,例如 相机(正交/透视)。在这里,我们将初始化其中一些组件,并对其余部分使用默认值。
在本例中,我们将首先创建一个使用透视相机、点光源并应用Phong 着色的渲染器。然后,我们学习如何使用模块化 API 来调整不同的组件。
# Initialize a camera.
# With world coordinates +Y up, +X left and +Z in, the front of the cow is facing the -Z direction.
# So we move the camera by 180 in the azimuth direction so it is facing the front of the cow.
R, T = look_at_view_transform(2.7, 0, 180)
cameras = FoVPerspectiveCameras(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. 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=512,
blur_radius=0.0,
faces_per_pixel=1,
)
# Place a point light in front of the object. As mentioned above, the front of the cow is facing the
# -z direction.
lights = PointLights(device=device, location=[[0.0, 0.0, -3.0]])
# Create a Phong renderer by composing a rasterizer and a shader. The textured Phong shader will
# interpolate the texture uv coordinates for each vertex, sample from a texture image and
# apply the Phong lighting model
renderer = MeshRenderer(
rasterizer=MeshRasterizer(
cameras=cameras,
raster_settings=raster_settings
),
shader=SoftPhongShader(
device=device,
cameras=cameras,
lights=lights
)
)
5. 渲染 Mesh
光在物体的前面,所以它是明亮的,图像有镜面高光。
images = renderer(mesh)
plt.figure(figsize=(10, 10))
plt.imshow(images[0, ..., :3].cpu().numpy())
plt.axis("off")
4. 移动物体后的光源并重新渲染
我们可以通过调用“渲染器”将任意关键字参数传递给“光栅化器”/“着色器”,这样如果任何设置发生变化,渲染器就不需要重新初始化
在这种情况下,我们可以简单地更新灯光的位置,并将它们传递给渲染器的调用。
图像现在是黑暗的,因为只有环境照明,没有镜面高光。
# Now move the light so it is on the +Z axis which will be behind the cow.
lights.location = torch.tensor([0.0, 0.0, +1.0], device=device)[None]
images = renderer(mesh, lights=lights)
plt.figure(figsize=(10, 10))
plt.imshow(images[0, ..., :3].cpu().numpy())
plt.axis("off")
5. 旋转目标,修改材料属性或光属性
我们还可以更改渲染管道中的许多其他设置。我们在这里:
- 改变视角的相机
- 改变点光源的位置
- 改变网格的材料反射率属性
# Rotate the object by increasing the elevation and azimuth angles
R, T = look_at_view_transform(dist=2.7, elev=10, azim=-150)
cameras = FoVPerspectiveCameras(device=device, R=R, T=T)
# Move the light location so the light is shining on the cow's face.
lights.location = torch.tensor([[2.0, 2.0, -2.0]], device=device)
# Change specular color to green and change material shininess
materials = Materials(
device=device,
specular_color=[[0.0, 1.0, 0.0]],
shininess=10.0
)
# Re render the mesh, passing in keyword arguments for the modified components.
images = renderer(mesh, lights=lights, materials=materials, cameras=cameras)
plt.figure(figsize=(10, 10))
plt.imshow(images[0, ..., :3].cpu().numpy())
plt.axis("off")
6. 批量渲染
PyTorch3D API的核心设计选择之一是支持所有组件的批处理输入。
渲染器和相关组件可以接受批量输入,并在一次转发中渲染一批输出图像。现在我们将使用这个特性从许多不同的角度渲染网格。
# Set batch size - this is the number of different viewpoints from which we want to render the mesh.
batch_size = 20
# Create a batch of meshes by repeating the cow mesh and associated textures.
# Meshes has a useful `extend` method which allows us do this very easily.
# This also extends the textures.
meshes = mesh.extend(batch_size)
# Get a batch of viewing angles.
elev = torch.linspace(0, 180, batch_size)
azim = torch.linspace(-180, 180, batch_size)
# All the cameras helper methods support mixed type inputs and broadcasting. So we can
# view the camera from the same distance and specify dist=2.7 as a float,
# and then specify elevation and azimuth angles for each viewpoint as tensors.
R, T = look_at_view_transform(dist=2.7, elev=elev, azim=azim)
cameras = FoVPerspectiveCameras(device=device, R=R, T=T)
# Move the light back in front of the cow which is facing the -z direction.
lights.location = torch.tensor([[0.0, 0.0, -3.0]], device=device)
# We can pass arbitrary keyword arguments to the rasterizer/shader via the renderer
# so the renderer does not need to be reinitialized if any of the settings change.
images = renderer(meshes, cameras=cameras, lights=lights)
image_grid(images.cpu().numpy(), rows=4, cols=5, rgb=True)