pytorch3d-数据集

ShapeNet 是一个大规模的 3D 形状数据集，包含超过 55 个类别、51,300 个 3D 模型。在 PyTorch3D 中使用 ShapeNet 数据集需要几个关键步骤：下载数据、创建数据加载器、可视化样本和集成到训练流程中。通过以上步骤，你可以有效地使用 ShapeNet 数据集训练和评估 2D 到 3D 的模型。：ShapeNet 包含大量 3D 模型，加载时可能占用较多内存，建议使

dongf2019

931人浏览 · 2025-07-02 10:58:21

dongf2019 · 2025-07-02 10:58:21 发布

ShapeNet 数据集使用指南

ShapeNet 是一个大规模的 3D 形状数据集，包含超过 55 个类别、51,300 个 3D 模型。在 PyTorch3D 中使用 ShapeNet 数据集需要几个关键步骤：下载数据、创建数据加载器、可视化样本和集成到训练流程中。

一、下载 ShapeNet 数据集

注册并下载
- 访问ShapeNet 官网注册账号
- 下载所需类别的数据（例如02691156对应飞机）

数据结构
下载后的数据结构通常为：

ShapeNetCore/
├── 02691156/  # 飞机类别
│   ├── 1a04e3eab45ca15dd86060f189eb133/
│   │   ├── models/
│   │   │   ├── model_normalized.obj
│   │   │   └── ...
│   │   └── ...
│   └── ...
├── 02958343/  # 汽车类别
└── ...

二、创建 ShapeNet 数据加载器

PyTorch3D 提供了ShapeNetCore数据集类，可直接加载数据：

from pytorch3d.datasets import ShapeNetCore

# 数据集路径
shapenet_dir = "/path/to/ShapeNetCore"

# 加载特定类别的数据（例如飞机和汽车）
categories = ["airplane", "car"]  # 也可以使用类别ID如 ["02691156", "02958343"]
dataset = ShapeNetCore(
    shapenet_dir,
    categories=categories,
    version=2,  # 版本2使用.obj文件
    load_textures=True,  # 是否加载纹理
    texture_resolution=4,  # 纹理分辨率
)

# 创建数据加载器
from torch.utils.data import DataLoader

dataloader = DataLoader(
    dataset,
    batch_size=4,
    shuffle=True,
    num_workers=4
)

三、可视化 ShapeNet 样本

使用 PyTorch3D 的渲染工具可视化数据集中的 3D 模型：

import matplotlib.pyplot as plt
from pytorch3d.renderer import (
    look_at_view_transform,
    FoVPerspectiveCameras,
    PointLights,
    RasterizationSettings,
    MeshRenderer,
    MeshRasterizer,
    SoftPhongShader,
)

# 初始化渲染器
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 设置相机和渲染参数
R, T = look_at_view_transform(2.7, 0, 180)
cameras = FoVPerspectiveCameras(device=device, R=R, T=T)

raster_settings = RasterizationSettings(
    image_size=512,
    blur_radius=0.0,
    faces_per_pixel=1,
)

lights = PointLights(device=device, location=[[0.0, 0.0, 3.0]])

renderer = MeshRenderer(
    rasterizer=MeshRasterizer(cameras=cameras, raster_settings=raster_settings),
    shader=SoftPhongShader(cameras=cameras, lights=lights, device=device)
)

# 可视化样本
def visualize_shapenet_sample(mesh, title="ShapeNet Sample"):
    # 渲染网格
    images = renderer(mesh)
    
    # 显示图像
    plt.figure(figsize=(10, 10))
    plt.imshow(images[0, ..., :3].cpu().numpy())
    plt.title(title)
    plt.axis("off")
    plt.show()

# 从数据集中获取一个样本
sample = dataset[0]
mesh = sample["mesh"].to(device)
category = sample["synset_id"]
model_id = sample["model_id"]

print(f"类别: {category}, 模型ID: {model_id}")
visualize_shapenet_sample(mesh)

四、集成 ShapeNet 到训练流程

将 ShapeNet 数据集与之前的 2D 到 3D 模型训练流程结合：

# 准备ShapeNet数据集
shapenet_dir = "/path/to/ShapeNetCore"
categories = ["airplane", "car", "chair"]  # 选择感兴趣的类别
dataset = ShapeNetCore(shapenet_dir, categories=categories, version=2)

# 创建自定义数据集，将3D模型与随机渲染的2D图像配对
class ShapeNet2D3DDataset(Dataset):
    def __init__(self, shapenet_dataset, image_size=256, num_views=1):
        self.shapenet_dataset = shapenet_dataset
        self.image_size = image_size
        self.num_views = num_views
        self.renderer = self._create_renderer()
        
    def _create_renderer(self):
        # 创建用于生成2D图像的渲染器
        R, T = look_at_view_transform(2.7, 0, 0)  # 初始视角
        cameras = FoVPerspectiveCameras(device=device, R=R, T=T)
        
        raster_settings = RasterizationSettings(
            image_size=self.image_size,
            blur_radius=0.0,
            faces_per_pixel=1,
        )
        
        lights = PointLights(device=device, location=[[0.0, 0.0, 3.0]])
        
        return MeshRenderer(
            rasterizer=MeshRasterizer(cameras=cameras, raster_settings=raster_settings),
            shader=SoftPhongShader(cameras=cameras, lights=lights, device=device)
        )
    
    def __len__(self):
        return len(self.shapenet_dataset) * self.num_views
    
    def __getitem__(self, idx):
        # 获取ShapeNet样本
        shapenet_idx = idx // self.num_views
        view_idx = idx % self.num_views
        
        sample = self.shapenet_dataset[shapenet_idx]
        mesh = sample["mesh"].to(device)
        
        # 随机生成视角
        elev = torch.rand(1) * 45  # 仰角: 0-45度
        azim = torch.rand(1) * 360  # 方位角: 0-360度
        
        R, T = look_at_view_transform(dist=2.7, elev=elev, azim=azim)
        self.renderer.rasterizer.cameras = FoVPerspectiveCameras(device=device, R=R, T=T)
        
        # 渲染2D图像
        image = self.renderer(mesh)[0, ..., :3].permute(2, 0, 1)  # [C, H, W]
        
        # 数据预处理
        transform = transforms.Compose([
            transforms.Resize(256),
            transforms.CenterCrop(256),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        ])
        
        image = transform(image)
        
        return image, mesh

# 创建数据加载器
combined_dataset = ShapeNet2D3DDataset(dataset, num_views=5)  # 每个3D模型生成5个不同视角的图像
train_size = int(0.8 * len(combined_dataset))
val_size = len(combined_dataset) - train_size
train_dataset, val_dataset = torch.utils.data.random_split(combined_dataset, [train_size, val_size])

train_loader = DataLoader(train_dataset, batch_size=4, shuffle=True, num_workers=4)
val_loader = DataLoader(val_dataset, batch_size=4, shuffle=False, num_workers=4)

# 训练模型
model = ImageTo3DModel().to(device)
renderer = create_renderer()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

# 训练循环
train_model(model, renderer, train_loader, val_loader, optimizer, num_epochs=50)