Texture loading and rendering in ShapeNetCore and R2N2 data loaders

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

tests/test_r2n2.py

@@ -33,6 +33,8 @@ from torch.utils.data import DataLoader
 R2N2_PATH = None
 SHAPENET_PATH = None
 SPLITS_PATH = None
+VOXELS_REL_PATH = "ShapeNetVox"
+
 
 DEBUG = False
 DATA_DIR = Path(__file__).resolve().parent / "data"
@@ -69,7 +71,12 @@ class TestR2N2(TestCaseMixin, unittest.TestCase):
         """
         # Load dataset in the train split.
         r2n2_dataset = R2N2(
-            "test", SHAPENET_PATH, R2N2_PATH, SPLITS_PATH, return_voxels=True
+            "test",
+            SHAPENET_PATH,
+            R2N2_PATH,
+            SPLITS_PATH,
+            return_voxels=True,
+            voxels_rel_path=VOXELS_REL_PATH,
         )
 
         # Check total number of objects in the dataset is correct.
@@ -133,7 +140,12 @@ class TestR2N2(TestCaseMixin, unittest.TestCase):
         """
         # Load dataset in the train split.
         r2n2_dataset = R2N2(
-            "val", SHAPENET_PATH, R2N2_PATH, SPLITS_PATH, return_voxels=True
+            "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.
@@ -362,7 +374,12 @@ class TestR2N2(TestCaseMixin, unittest.TestCase):
 
         # 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
+            "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)

tests/test_shapenet_core.py

@@ -22,6 +22,7 @@ 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
@@ -55,7 +56,7 @@ class TestShapenetCore(TestCaseMixin, unittest.TestCase):
         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)
+        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