在 [ ]

# Copyright (c) Meta Platforms, Inc. and affiliates. All rights reserved.

ShapeNetCore 和 R2N2 的数据加载器¶

本教程展示如何

使用 PyTorch3D 的数据加载器加载 ShapeNetCore 和 R2N2 中的模型。
将加载的数据集传递给 torch.utils.data.DataLoader。
使用 PyTorch3D 的渲染器渲染 ShapeNetCore 模型。
以与数据集中原始渲染相同的方位渲染 R2N2 模型。
可视化 R2N2 模型体素。

0. 安装并导入模块¶

确保已安装 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 numpy as np
import torch

from pytorch3d.datasets import (
    R2N2,
    ShapeNetCore,
    collate_batched_meshes,
    render_cubified_voxels,
)
from pytorch3d.renderer import (
    OpenGLPerspectiveCameras,
    PointLights,
    RasterizationSettings,
    TexturesVertex,
    look_at_view_transform,
)

from pytorch3d.structures import Meshes
from torch.utils.data import DataLoader

# add path for demo utils functions 
import sys
import os
sys.path.append(os.path.abspath(''))

如果使用 Google Colab，请获取用于绘制图像网格的 utils 文件

在 [ ]

!wget https://raw.githubusercontent.com/facebookresearch/pytorch3d/main/docs/tutorials/utils/plot_image_grid.py
from plot_image_grid import image_grid

或者，如果在本地运行，请取消注释并运行以下单元格

在 [ ]

# from utils import image_grid

1. 加载数据集¶

如果您尚未下载 ShapeNetCore 数据集，请先按照此处的说明进行下载：https://www.shapenet.org/。ShapeNetCore 是 ShapeNet 数据集的一个子集。在 PyTorch3D 中，我们支持版本 1（57 个类别）和版本 2（55 个类别）。

然后将下面的 SHAPENET_PATH 修改为您本地 ShapeNetCore 数据集文件夹的路径。

在 [ ]

# Setup
if torch.cuda.is_available():
    device = torch.device("cuda:0")
    torch.cuda.set_device(device)
else:
    device = torch.device("cpu")
    
SHAPENET_PATH = ""
shapenet_dataset = ShapeNetCore(SHAPENET_PATH)

R2N2 数据集可以通过此处的说明下载：http://3d-r2n2.stanford.edu/。查看 ShapeNetRendering 和 ShapeNetVox32 的链接。R2N2 数据集包含 13 个类别，这些类别是 ShapeNetCore v.1 数据集的子集。R2N2 数据集还包含其自己的每个对象的 24 个渲染和体素化模型。

然后将下面的 R2N2_PATH 和 SPLITS_PATH 分别修改为您的本地 R2N2 数据集文件夹路径和分割文件路径。在这里，我们将加载 R2N2 的 train 分割，并要求返回每个模型的体素。

在 [ ]

R2N2_PATH = ""
SPLITS_PATH = "None"
r2n2_dataset = R2N2("train", SHAPENET_PATH, R2N2_PATH, SPLITS_PATH, return_voxels=True)

我们可以通过索引加载的数据集来检索模型。对于 ShapeNetCore 和 R2N2，我们可以检查此模型所属的类别（以 synset id 的形式，等同于 ImageNet API 中描述的 wnid：http://image-net.org/download-API）、其模型 ID 以及其顶点和面。

在 [ ]

shapenet_model = shapenet_dataset[6]
print("This model belongs to the category " + shapenet_model["synset_id"] + ".")
print("This model has model id " + shapenet_model["model_id"] + ".")
model_verts, model_faces = shapenet_model["verts"], shapenet_model["faces"]

我们可以使用其顶点和面来形成一个 Meshes 对象，这是一个用于处理批处理网格的 PyTorch3D 数据结构。

在 [ ]

model_textures = TexturesVertex(verts_features=torch.ones_like(model_verts, device=device)[None])
shapenet_model_mesh = Meshes(
    verts=[model_verts.to(device)],   
    faces=[model_faces.to(device)],
    textures=model_textures
)

对于 R2N2，我们可以进一步检查 R2N2 的原始渲染。例如，如果我们想查看 R2N2 数据集中第 11 个对象的第二个和第三个视图，我们可以执行以下操作

在 [ ]

r2n2_renderings = r2n2_dataset[10,[1,2]]
image_grid(r2n2_renderings.numpy(), rows=1, cols=2, rgb=True)

2. 使用 `torch.utils.data.DataLoader` 使用数据集¶

训练深度学习模型通常需要传入批量的输入。PyTorch 中的 torch.utils.data.DataLoader 可以帮助我们做到这一点。PyTorch3D 提供了一个函数 collate_batched_meshes 用于将输入网格组合成一个表示批次的 Meshes 对象。然后，Meshes 数据结构可以直接由深度学习模型中可能包含的其他 PyTorch3D 操作（例如 graph_conv）使用。

对于 R2N2，如果批次中的所有模型都具有相同数量的视图，则视图、旋转矩阵、平移矩阵、内参矩阵和体素也将被堆叠成批处理张量。

注意：R2N2 的 val 分割中的所有模型都有 24 个视图，但有 8 个模型将其 24 个视图分割到 train 和 test 分割中，在这种情况下，collate_batched_meshes 只能将矩阵、视图和体素作为列表连接。但是，可以通过设置 return_all_views = False 只加载每个模型的一个视图来避免这种情况。

在 [ ]

batch_size = 12
r2n2_single_view = R2N2("train", SHAPENET_PATH, R2N2_PATH, SPLITS_PATH, return_all_views=False, return_voxels=True)
r2n2_loader = DataLoader(r2n2_single_view, batch_size=batch_size, collate_fn=collate_batched_meshes)

让我们可视化批次中的所有视图（每个模型一个）

在 [ ]

it = iter(r2n2_loader)
r2n2_batch = next(it)
batch_renderings = r2n2_batch["images"] # (N, V, H, W, 3), and in this case V is 1.
image_grid(batch_renderings.squeeze().numpy(), rows=3, cols=4, rgb=True)

3. 使用 PyTorch3D 的可微渲染器渲染 ShapeNetCore 模型¶

ShapeNetCore 和 R2N2 数据加载器都具有自定义的 render 函数，这些函数支持通过指定其模型 ID、类别或索引来使用 PyTorch3D 的可微渲染器实现渲染模型。

在 [ ]

# Rendering settings.
R, T = look_at_view_transform(1.0, 1.0, 90)
cameras = OpenGLPerspectiveCameras(R=R, T=T, device=device)
raster_settings = RasterizationSettings(image_size=512)
lights = PointLights(location=torch.tensor([0.0, 1.0, -2.0], device=device)[None],device=device)

首先，我们将尝试根据其模型 ID 渲染三个模型

在 [ ]

images_by_model_ids = shapenet_dataset.render(
    model_ids=[
        "13394ca47c89f91525a3aaf903a41c90",
        "14755c2ee8e693aba508f621166382b0",
        "156c4207af6d2c8f1fdc97905708b8ea",
    ],
    device=device,
    cameras=cameras,
    raster_settings=raster_settings,
    lights=lights,
)
image_grid(images_by_model_ids.cpu().numpy(), rows=1, cols=3, rgb=True)

假设我们想渲染数据集中前三个模型，我们可以根据其索引渲染模型

在 [ ]

images_by_idxs = shapenet_dataset.render(
    idxs=list(range(3)),
    device=device,
    cameras=cameras,
    raster_settings=raster_settings,
    lights=lights,
)
image_grid(images_by_idxs.cpu().numpy(), rows=1, cols=3, rgb=True)

或者，如果我们对任何特定模型不感兴趣，但想查看来自某些特定类别的随机模型，我们可以通过指定 categories 和 sample_nums 来做到这一点。例如，如果我们想从“水龙头”类别渲染 2 个模型，从“椅子”类别渲染 3 个模型，我们可以执行以下操作

在 [ ]

images_by_categories = shapenet_dataset.render(
    categories=["faucet", "chair"],
    sample_nums=[2, 3],
    device=device,
    cameras=cameras,
    raster_settings=raster_settings,
    lights=lights,
)
image_grid(images_by_categories.cpu().numpy(), rows=1, cols=5, rgb=True)

如果我们对任何特定类别不感兴趣，并且只想从整个数据集中渲染一些随机模型，我们可以设置要在 sample_nums 中渲染的模型数量，并且不指定任何 categories

在 [ ]

random_model_images = shapenet_dataset.render(
    sample_nums=[3],
    device=device,
    cameras=cameras,
    raster_settings=raster_settings,
    lights=lights,
)
image_grid(random_model_images.cpu().numpy(), rows=1, cols=5, rgb=True)

4. 以与数据集中原始渲染相同的方位渲染 R2N2 模型¶

我们可以像上面渲染 ShapeNetCore 模型一样渲染 R2N2 模型。此外，我们还可以以与数据集中原始渲染相同的方位渲染 R2N2 模型。为此，我们将使用 R2N2 的自定义 render 函数和一种不同类型的 PyTorch3D 相机，称为 BlenderCamera。

在此示例中，我们将以与其第二个和第三个视图相同的方位渲染第七个模型。首先，我们将检索 R2N2 的原始渲染以与结果进行比较。

在 [ ]

original_rendering = r2n2_dataset[6,[1,2]]["images"]
image_grid(original_rendering.numpy(), rows=1, cols=2, rgb=True)

接下来，我们将可视化 PyTorch3D 的渲染结果

在 [ ]

r2n2_oriented_images = r2n2_dataset.render(
    idxs=[6],
    view_idxs=[1,2],
    device=device,
    raster_settings=raster_settings,
    lights=lights,
)
image_grid(r2n2_oriented_images.cpu().numpy(), rows=1, cols=2, rgb=True)

5. 可视化 R2N2 模型的体素¶

R2N2 数据加载器还返回模型的体素。我们可以通过利用 R2N2 的 render_vox_to_mesh 函数来可视化它们。这会将体素立方化成一个 Meshes 对象，然后对其进行渲染。

在此示例中，我们将以与其第二个和第三个视图相同的方位可视化数据集中第十个模型。首先，我们将检索 R2N2 的原始渲染以与结果进行比较。

在 [ ]

r2n2_model = r2n2_dataset[9,[1,2]]
original_rendering = r2n2_model["images"]
image_grid(original_rendering.numpy(), rows=1, cols=2, rgb=True)

接下来，我们将体素传递给 render_vox_to_mesh

在 [ ]

vox_render = render_cubified_voxels(r2n2_model["voxels"], device=device)
image_grid(vox_render.cpu().numpy(), rows=1, cols=2, rgb=True)

ShapeNetCore 和 R2N2 的数据加载器¶

0. 安装并导入模块¶

1. 加载数据集¶

2. 使用 torch.utils.data.DataLoader 使用数据集¶

3. 使用 PyTorch3D 的可微渲染器渲染 ShapeNetCore 模型¶

4. 以与数据集中原始渲染相同的方位渲染 R2N2 模型¶

5. 可视化 R2N2 模型的体素¶

2. 使用 `torch.utils.data.DataLoader` 使用数据集¶