def test_render_voxels(self):
"""
Test rendering meshes formed from voxels.
"""
# Set up device and seed for random selections.
device = torch.device("cuda:0")
# Load dataset in the train split with only a single view returned for each model.
r2n2_dataset = R2N2(
"train",
SHAPENET_PATH,
R2N2_PATH,
SPLITS_PATH,
return_voxels=True,
voxels_rel_path=VOXELS_REL_PATH,
)
r2n2_model = r2n2_dataset[6, [5]]
vox_render = render_cubified_voxels(r2n2_model["voxels"], device=device)
vox_render_rgb = vox_render[0, ..., :3].squeeze().cpu()
if DEBUG:
Image.fromarray((vox_render_rgb.numpy() * 255).astype(np.uint8)).save(
DATA_DIR / ("DEBUG_r2n2_voxel_to_mesh_render.png")
)
image_ref = load_rgb_image("test_r2n2_voxel_to_mesh_render.png", DATA_DIR)
self.assertClose(vox_render_rgb, image_ref, atol=0.05)
def test_catch_render_arg_errors(self):
"""
Test rendering R2N2 with an invalid model_id, category or index, and
catch corresponding errors.
"""
# Load dataset in the train split.
r2n2_dataset = R2N2("train", SHAPENET_PATH, R2N2_PATH, SPLITS_PATH)
# Try loading with an invalid model_id and catch error.
with self.assertRaises(ValueError) as err:
r2n2_dataset.render(model_ids=["lamp0"])
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:
r2n2_dataset.render(idxs=[1000000])
self.assertTrue("are out of bounds" in str(err.exception))
blend_cameras = BlenderCamera(R=torch.rand((3, 3, 3)),
T=torch.rand((3, 3)),
K=torch.rand((3, 4, 4)))
with self.assertRaises(ValueError) as err:
r2n2_dataset.render(idxs=[10, 11], cameras=blend_cameras)
self.assertTrue("Mismatch between batch dims" in str(err.exception))
def test_load_R2N2(self):
"""
Test the loaded train split of R2N2 return items of the correct shapes and types. Also
check the first image returned is correct.
"""
# Load dataset in the train split.
r2n2_dataset = R2N2("train", SHAPENET_PATH, R2N2_PATH, SPLITS_PATH)
# Check total number of objects in the dataset is correct.
with open(SPLITS_PATH) as splits:
split_dict = json.load(splits)["train"]
model_nums = [len(split_dict[synset].keys()) for synset in split_dict.keys()]
self.assertEqual(len(r2n2_dataset), sum(model_nums))
# Randomly retrieve an object from the dataset.
rand_idx = torch.randint(len(r2n2_dataset), (1,))
rand_obj = r2n2_dataset[rand_idx]
# Check that verts and faces returned by __getitem__ have the correct shapes and types.
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)
# Check that image batch returned by __getitem__ has the correct shape.
self.assertEqual(rand_obj["images"].shape[0], 24)
self.assertEqual(rand_obj["images"].shape[1], 137)
self.assertEqual(rand_obj["images"].shape[2], 137)
self.assertEqual(rand_obj["images"].shape[-1], 3)
self.assertEqual(r2n2_dataset[rand_idx, [21]]["images"].shape[0], 1)
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 dataset in the train split.
r2n2_dataset = R2N2(
"val",
SHAPENET_PATH,
R2N2_PATH,
SPLITS_PATH,
return_voxels=True,
voxels_rel_path=VOXELS_REL_PATH,
)
# Randomly retrieve several objects from the dataset and collate them.
collated_meshes = collate_batched_R2N2([
r2n2_dataset[idx]
for idx in torch.randint(len(r2n2_dataset), (6, ))
])
# Check the collated verts and faces have the correct shapes.
verts, faces = collated_meshes["verts"], collated_meshes["faces"]
self.assertEqual(len(verts), 6)
self.assertEqual(len(faces), 6)
self.assertEqual(verts[0].shape[-1], 3)
self.assertEqual(faces[0].shape[-1], 3)
# Check the collated mesh has the correct shape.
mesh = collated_meshes["mesh"]
self.assertEqual(mesh.verts_padded().shape[0], 6)
self.assertEqual(mesh.verts_padded().shape[-1], 3)
self.assertEqual(mesh.faces_padded().shape[0], 6)
self.assertEqual(mesh.faces_padded().shape[-1], 3)
# Pass the custom collate_fn function to DataLoader and check elements
# in batch have the correct shape.
batch_size = 12
r2n2_loader = DataLoader(r2n2_dataset,
batch_size=batch_size,
collate_fn=collate_batched_R2N2)
it = iter(r2n2_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)
self.assertEqual(object_batch["images"].shape[0], batch_size)
self.assertEqual(object_batch["R"].shape[0], batch_size)
self.assertEqual(object_batch["T"].shape[0], batch_size)
self.assertEqual(object_batch["K"].shape[0], batch_size)
self.assertEqual(len(object_batch["voxels"]), batch_size)
def test_load_R2N2(self):
"""
Test the loaded train split of R2N2 return items of the correct shapes and types. Also
check the first image returned is correct.
"""
# Load dataset in the test split.
r2n2_dataset = R2N2("test", SHAPENET_PATH, R2N2_PATH, SPLITS_PATH)
# Check total number of objects in the dataset is correct.
with open(SPLITS_PATH) as splits:
split_dict = json.load(splits)["test"]
model_nums = [len(split_dict[synset]) for synset in split_dict]
self.assertEqual(len(r2n2_dataset), sum(model_nums))
# Check the numbers of loaded instances for each category are correct.
for synset in split_dict:
split_synset_nums = sum(
len(split_dict[synset][model]) for model in split_dict[synset]
)
idx_start = r2n2_dataset.synset_start_idxs[synset]
idx_end = idx_start + r2n2_dataset.synset_num_models[synset]
synset_views_list = r2n2_dataset.views_per_model_list[idx_start:idx_end]
loaded_synset_views = sum(len(views) for views in synset_views_list)
self.assertEqual(loaded_synset_views, split_synset_nums)
# Retrieve an object from the dataset.
r2n2_obj = r2n2_dataset[39]
# Check that verts and faces returned by __getitem__ have the correct shapes and types.
verts, faces = r2n2_obj["verts"], r2n2_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)
# Check that image batch returned by __getitem__ has the correct shape.
self.assertEqual(r2n2_obj["images"].shape[0], 24)
self.assertEqual(r2n2_obj["images"].shape[1], 137)
self.assertEqual(r2n2_obj["images"].shape[2], 137)
self.assertEqual(r2n2_obj["images"].shape[-1], 3)
self.assertEqual(r2n2_dataset[39, [21]]["images"].shape[0], 1)
self.assertEqual(r2n2_dataset[39, torch.tensor([12, 21])]["images"].shape[0], 2)
# Check models with total view counts less than 24 return image batches
# of the correct shapes.
self.assertEqual(r2n2_dataset[635]["images"].shape[0], 5)
self.assertEqual(r2n2_dataset[8369]["images"].shape[0], 10)
def test_catch_render_arg_errors(self):
"""
Test rendering R2N2 with an invalid model_id, category or index, and
catch corresponding errors.
"""
# Load dataset in the train split.
r2n2_dataset = R2N2("train", SHAPENET_PATH, R2N2_PATH, SPLITS_PATH)
# Try loading with an invalid model_id and catch error.
with self.assertRaises(ValueError) as err:
r2n2_dataset.render(model_ids=["lamp0"])
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:
r2n2_dataset.render(idxs=[1000000])
self.assertTrue("are out of bounds" in str(err.exception))
def test_render_by_r2n2_calibration(self):
"""
Test rendering R2N2 models with calibration matrices from R2N2's own Blender
in batches.
"""
# Set up device and seed for random selections.
device = torch.device("cuda:0")
torch.manual_seed(39)
# Load dataset in the train split.
r2n2_dataset = R2N2("train", SHAPENET_PATH, R2N2_PATH, SPLITS_PATH)
model_idxs = torch.randint(1000, (2, )).tolist()
view_idxs = torch.randint(24, (2, )).tolist()
raster_settings = RasterizationSettings(image_size=512)
lights = PointLights(
location=torch.tensor([0.0, 1.0, -2.0], device=device)[None],
# TODO(nikhilar): 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,
)
r2n2_batch = r2n2_dataset.render(
idxs=model_idxs,
view_idxs=view_idxs,
device=device,
raster_settings=raster_settings,
lights=lights,
)
for idx in range(4):
r2n2_batch_rgb = r2n2_batch[idx, ..., :3].squeeze().cpu()
if DEBUG:
Image.fromarray(
(r2n2_batch_rgb.numpy() * 255).astype(np.uint8)).save(
DATA_DIR /
("DEBUG_r2n2_render_with_blender_calibrations_%s.png" %
idx))
image_ref = load_rgb_image(
"test_r2n2_render_with_blender_calibrations_%s.png" % idx,
DATA_DIR)
self.assertClose(r2n2_batch_rgb, image_ref, atol=0.05)
def test_render_r2n2(self):
"""
Test rendering objects from R2N2 selected both by indices and model_ids.
"""
# Set up device and seed for random selections.
device = torch.device("cuda:0")
torch.manual_seed(39)
# Load dataset in the train split.
r2n2_dataset = R2N2("train", SHAPENET_PATH, R2N2_PATH, SPLITS_PATH)
# Render first three models in the dataset.
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,
)
r2n2_by_idxs = r2n2_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(r2n2_by_idxs.shape[0], 3)
# Compare the rendered models to the reference images.
for idx in range(3):
r2n2_by_idxs_rgb = r2n2_by_idxs[idx, ..., :3].squeeze().cpu()
if DEBUG:
Image.fromarray((r2n2_by_idxs_rgb.numpy() * 255).astype(np.uint8)).save(
DATA_DIR / ("DEBUG_r2n2_render_by_idxs_%s.png" % idx)
)
image_ref = load_rgb_image(
"test_r2n2_render_by_idxs_and_ids_%s.png" % idx, DATA_DIR
)
self.assertClose(r2n2_by_idxs_rgb, image_ref, atol=0.05)
# Render the same models but by model_ids this time.
r2n2_by_model_ids = r2n2_dataset.render(
model_ids=[
"1a4a8592046253ab5ff61a3a2a0e2484",
"1a04dcce7027357ab540cc4083acfa57",
"1a9d0480b74d782698f5bccb3529a48d",
],
device=device,
cameras=cameras,
raster_settings=raster_settings,
lights=lights,
)
# Compare the rendered models to the reference images.
for idx in range(3):
r2n2_by_model_ids_rgb = r2n2_by_model_ids[idx, ..., :3].squeeze().cpu()
if DEBUG:
Image.fromarray(
(r2n2_by_model_ids_rgb.numpy() * 255).astype(np.uint8)
).save(DATA_DIR / ("DEBUG_r2n2_render_by_model_ids_%s.png" % idx))
image_ref = load_rgb_image(
"test_r2n2_render_by_idxs_and_ids_%s.png" % idx, DATA_DIR
)
self.assertClose(r2n2_by_model_ids_rgb, image_ref, atol=0.05)
###############################
# Test rendering by categories
###############################
# Render a mixture of categories.
categories = ["chair", "lamp"]
mixed_objs = r2n2_dataset.render(
categories=categories,
sample_nums=[1, 2],
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_r2n2_render_by_categories_%s.png" % idx)
)
image_ref = load_rgb_image(
"test_r2n2_render_by_categories_%s.png" % idx, DATA_DIR
)
self.assertClose(mixed_rgb, image_ref, atol=0.05)
