mirror of
https://github.com/facebookresearch/pytorch3d.git
synced 2025-08-02 03:42:50 +08:00
778383eef7
Summary: - Add support for loading textures from ShapeNet Obj files as a texture atlas. - Support textured rendering of shapenet models Reviewed By: gkioxari Differential Revision: D23141143 fbshipit-source-id: 26eb81758d4cdbd6d820b072b58f5c6c08cb90bc
286 lines
11 KiB
Python
286 lines
11 KiB
Python
# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
|
|
"""
|
|
Sanity checks for loading ShapeNetCore.
|
|
"""
|
|
import os
|
|
import unittest
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import torch
|
|
from common_testing import TestCaseMixin, load_rgb_image
|
|
from PIL import Image
|
|
from pytorch3d.datasets import ShapeNetCore, collate_batched_meshes
|
|
from pytorch3d.renderer import (
|
|
FoVPerspectiveCameras,
|
|
PointLights,
|
|
RasterizationSettings,
|
|
look_at_view_transform,
|
|
)
|
|
from torch.utils.data import DataLoader
|
|
|
|
|
|
# Set the SHAPENET_PATH to the local path to the dataset
|
|
SHAPENET_PATH = None
|
|
VERSION = 1
|
|
# If DEBUG=True, save out images generated in the tests for debugging.
|
|
# All saved images have prefix DEBUG_
|
|
DEBUG = False
|
|
DATA_DIR = Path(__file__).resolve().parent / "data"
|
|
|
|
|
|
class TestShapenetCore(TestCaseMixin, unittest.TestCase):
|
|
def setUp(self):
|
|
"""
|
|
Check if the ShapeNet dataset is provided in the repo.
|
|
If not, download this separately and update the shapenet_path`
|
|
with the location of the dataset in order to run the tests.
|
|
"""
|
|
if SHAPENET_PATH is None or not os.path.exists(SHAPENET_PATH):
|
|
url = "https://www.shapenet.org/"
|
|
msg = (
|
|
"ShapeNet data not found, download from %s, update "
|
|
"SHAPENET_PATH at the top of the file, and rerun."
|
|
)
|
|
|
|
self.skipTest(msg % url)
|
|
|
|
def test_load_shapenet_core(self):
|
|
"""
|
|
Test loading both the entire ShapeNetCore dataset and a subset of the ShapeNetCore
|
|
dataset. Check the loaded datasets return items of the correct shapes and types.
|
|
"""
|
|
# Try loading ShapeNetCore with an invalid version number and catch error.
|
|
with self.assertRaises(ValueError) as err:
|
|
ShapeNetCore(SHAPENET_PATH, version=3)
|
|
self.assertTrue("Version number must be either 1 or 2." in str(err.exception))
|
|
|
|
# Load ShapeNetCore without specifying any particular categories.
|
|
shapenet_dataset = ShapeNetCore(SHAPENET_PATH, version=VERSION)
|
|
|
|
# Count the number of grandchildren directories (which should be equal to
|
|
# the total number of objects in the dataset) by walking through the given
|
|
# directory.
|
|
wnsynset_list = [
|
|
wnsynset
|
|
for wnsynset in os.listdir(SHAPENET_PATH)
|
|
if os.path.isdir(os.path.join(SHAPENET_PATH, wnsynset))
|
|
]
|
|
model_num_list = [
|
|
(len(next(os.walk(os.path.join(SHAPENET_PATH, wnsynset)))[1]))
|
|
for wnsynset in wnsynset_list
|
|
]
|
|
# Check total number of objects in the dataset is correct.
|
|
self.assertEqual(len(shapenet_dataset), sum(model_num_list))
|
|
|
|
# Randomly retrieve an object from the dataset.
|
|
rand_obj = shapenet_dataset[torch.randint(len(shapenet_dataset), (1,))]
|
|
# Check that data types and shapes of items returned by __getitem__ are correct.
|
|
verts, faces = rand_obj["verts"], rand_obj["faces"]
|
|
self.assertTrue(verts.dtype == torch.float32)
|
|
self.assertTrue(faces.dtype == torch.int64)
|
|
self.assertEqual(verts.ndim, 2)
|
|
self.assertEqual(verts.shape[-1], 3)
|
|
self.assertEqual(faces.ndim, 2)
|
|
self.assertEqual(faces.shape[-1], 3)
|
|
|
|
# Load six categories from ShapeNetCore.
|
|
# Specify categories with a combination of offsets and labels.
|
|
shapenet_subset = ShapeNetCore(
|
|
SHAPENET_PATH,
|
|
synsets=[
|
|
"04330267",
|
|
"guitar",
|
|
"02801938",
|
|
"birdhouse",
|
|
"03991062",
|
|
"tower",
|
|
],
|
|
version=1,
|
|
)
|
|
subset_offsets = [
|
|
"04330267",
|
|
"03467517",
|
|
"02801938",
|
|
"02843684",
|
|
"03991062",
|
|
"04460130",
|
|
]
|
|
subset_model_nums = [
|
|
(len(next(os.walk(os.path.join(SHAPENET_PATH, offset)))[1]))
|
|
for offset in subset_offsets
|
|
]
|
|
self.assertEqual(len(shapenet_subset), sum(subset_model_nums))
|
|
|
|
def test_collate_models(self):
|
|
"""
|
|
Test collate_batched_meshes returns items of the correct shapes and types.
|
|
Check that when collate_batched_meshes is passed to Dataloader, batches of
|
|
the correct shapes and types are returned.
|
|
"""
|
|
# Load ShapeNetCore without specifying any particular categories.
|
|
shapenet_dataset = ShapeNetCore(SHAPENET_PATH)
|
|
# Randomly retrieve several objects from the dataset.
|
|
rand_idxs = torch.randint(len(shapenet_dataset), (6,))
|
|
rand_objs = [shapenet_dataset[idx] for idx in rand_idxs]
|
|
|
|
# Collate the randomly selected objects
|
|
collated_meshes = collate_batched_meshes(rand_objs)
|
|
verts, faces = (collated_meshes["verts"], collated_meshes["faces"])
|
|
self.assertEqual(len(verts), 6)
|
|
self.assertEqual(len(faces), 6)
|
|
|
|
# Pass the custom collate_fn function to DataLoader and check elements
|
|
# in batch have the correct shape.
|
|
batch_size = 12
|
|
shapenet_core_loader = DataLoader(
|
|
shapenet_dataset, batch_size=batch_size, collate_fn=collate_batched_meshes
|
|
)
|
|
it = iter(shapenet_core_loader)
|
|
object_batch = next(it)
|
|
self.assertEqual(len(object_batch["synset_id"]), batch_size)
|
|
self.assertEqual(len(object_batch["model_id"]), batch_size)
|
|
self.assertEqual(len(object_batch["label"]), batch_size)
|
|
self.assertEqual(object_batch["mesh"].verts_padded().shape[0], batch_size)
|
|
self.assertEqual(object_batch["mesh"].faces_padded().shape[0], batch_size)
|
|
|
|
def test_catch_render_arg_errors(self):
|
|
"""
|
|
Test rendering ShapeNetCore with invalid model_ids, categories or indices,
|
|
and catch corresponding errors.
|
|
"""
|
|
# Load ShapeNetCore.
|
|
shapenet_dataset = ShapeNetCore(SHAPENET_PATH)
|
|
|
|
# Try loading with an invalid model_id and catch error.
|
|
with self.assertRaises(ValueError) as err:
|
|
shapenet_dataset.render(model_ids=["piano0"])
|
|
self.assertTrue("not found in the loaded dataset" in str(err.exception))
|
|
|
|
# Try loading with an index out of bounds and catch error.
|
|
with self.assertRaises(IndexError) as err:
|
|
shapenet_dataset.render(idxs=[100000])
|
|
self.assertTrue("are out of bounds" in str(err.exception))
|
|
|
|
def test_render_shapenet_core(self):
|
|
"""
|
|
Test rendering objects from ShapeNetCore.
|
|
"""
|
|
# Setup device and seed for random selections.
|
|
device = torch.device("cuda:0")
|
|
torch.manual_seed(39)
|
|
|
|
# Load category piano from ShapeNetCore.
|
|
piano_dataset = ShapeNetCore(SHAPENET_PATH, synsets=["piano"])
|
|
|
|
# Rendering settings.
|
|
R, T = look_at_view_transform(1.0, 1.0, 90)
|
|
cameras = FoVPerspectiveCameras(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],
|
|
# TODO: debug the source of the discrepancy in two images when rendering on GPU.
|
|
diffuse_color=((0, 0, 0),),
|
|
specular_color=((0, 0, 0),),
|
|
device=device,
|
|
)
|
|
|
|
# Render first three models in the piano category.
|
|
pianos = piano_dataset.render(
|
|
idxs=list(range(3)),
|
|
device=device,
|
|
cameras=cameras,
|
|
raster_settings=raster_settings,
|
|
lights=lights,
|
|
)
|
|
# Check that there are three images in the batch.
|
|
self.assertEqual(pianos.shape[0], 3)
|
|
|
|
# Compare the rendered models to the reference images.
|
|
for idx in range(3):
|
|
piano_rgb = pianos[idx, ..., :3].squeeze().cpu()
|
|
if DEBUG:
|
|
Image.fromarray((piano_rgb.numpy() * 255).astype(np.uint8)).save(
|
|
DATA_DIR / ("DEBUG_shapenet_core_render_piano_by_idxs_%s.png" % idx)
|
|
)
|
|
image_ref = load_rgb_image(
|
|
"test_shapenet_core_render_piano_%s.png" % idx, DATA_DIR
|
|
)
|
|
self.assertClose(piano_rgb, image_ref, atol=0.05)
|
|
|
|
# Render the same piano models but by model_ids this time.
|
|
pianos_2 = piano_dataset.render(
|
|
model_ids=[
|
|
"13394ca47c89f91525a3aaf903a41c90",
|
|
"14755c2ee8e693aba508f621166382b0",
|
|
"156c4207af6d2c8f1fdc97905708b8ea",
|
|
],
|
|
device=device,
|
|
cameras=cameras,
|
|
raster_settings=raster_settings,
|
|
lights=lights,
|
|
)
|
|
|
|
# Compare the rendered models to the reference images.
|
|
for idx in range(3):
|
|
piano_rgb_2 = pianos_2[idx, ..., :3].squeeze().cpu()
|
|
if DEBUG:
|
|
Image.fromarray((piano_rgb_2.numpy() * 255).astype(np.uint8)).save(
|
|
DATA_DIR / ("DEBUG_shapenet_core_render_piano_by_ids_%s.png" % idx)
|
|
)
|
|
image_ref = load_rgb_image(
|
|
"test_shapenet_core_render_piano_%s.png" % idx, DATA_DIR
|
|
)
|
|
self.assertClose(piano_rgb_2, image_ref, atol=0.05)
|
|
|
|
#######################
|
|
# Render by categories
|
|
#######################
|
|
|
|
# Load ShapeNetCore.
|
|
shapenet_dataset = ShapeNetCore(SHAPENET_PATH)
|
|
|
|
# Render a mixture of categories and specify the number of models to be
|
|
# randomly sampled from each category.
|
|
mixed_objs = shapenet_dataset.render(
|
|
categories=["faucet", "chair"],
|
|
sample_nums=[2, 1],
|
|
device=device,
|
|
cameras=cameras,
|
|
raster_settings=raster_settings,
|
|
lights=lights,
|
|
)
|
|
# Compare the rendered models to the reference images.
|
|
for idx in range(3):
|
|
mixed_rgb = mixed_objs[idx, ..., :3].squeeze().cpu()
|
|
if DEBUG:
|
|
Image.fromarray((mixed_rgb.numpy() * 255).astype(np.uint8)).save(
|
|
DATA_DIR
|
|
/ ("DEBUG_shapenet_core_render_mixed_by_categories_%s.png" % idx)
|
|
)
|
|
image_ref = load_rgb_image(
|
|
"test_shapenet_core_render_mixed_by_categories_%s.png" % idx, DATA_DIR
|
|
)
|
|
self.assertClose(mixed_rgb, image_ref, atol=0.05)
|
|
|
|
# Render a mixture of categories without specifying sample_nums.
|
|
mixed_objs_2 = shapenet_dataset.render(
|
|
categories=["faucet", "chair"],
|
|
device=device,
|
|
cameras=cameras,
|
|
raster_settings=raster_settings,
|
|
lights=lights,
|
|
)
|
|
# Compare the rendered models to the reference images.
|
|
for idx in range(2):
|
|
mixed_rgb_2 = mixed_objs_2[idx, ..., :3].squeeze().cpu()
|
|
if DEBUG:
|
|
Image.fromarray((mixed_rgb_2.numpy() * 255).astype(np.uint8)).save(
|
|
DATA_DIR
|
|
/ ("DEBUG_shapenet_core_render_without_sample_nums_%s.png" % idx)
|
|
)
|
|
image_ref = load_rgb_image(
|
|
"test_shapenet_core_render_without_sample_nums_%s.png" % idx, DATA_DIR
|
|
)
|
|
self.assertClose(mixed_rgb_2, image_ref, atol=0.05)
|