Fix ShapeNetDataset (#593)

Talmaj · facebook-github-bot · commit 029a9da00b40 · 2021-06-18T07:02:20.000-07:00
Summary: - Add MANIFEST.in and `include_package_data=True` to include dataset .json files in the installation Fix #435 - Fix `load_textures=False` for ShapeNetDataset with a test Fix #450, partly fix #444. I've set the textures to `None`, if they should be all white instead, let me know. Pull Request resolved: #593 Reviewed By: patricklabatut Differential Revision: D29116264 Pulled By: nikhilaravi fbshipit-source-id: 1fb0198e616b7f834dfeaf7168bb5e6e530810d1
diff --git a/MANIFEST.in b/MANIFEST.in
@@ -0,0 +1,3 @@
+include pytorch3d/datasets/shapenet/shapenet_synset_dict_v1.json
+include pytorch3d/datasets/shapenet/shapenet_synset_dict_v2.json
+include pytorch3d/datasets/r2n2/r2n2_synset_dict.json
diff --git a/pytorch3d/datasets/shapenet_base.py b/pytorch3d/datasets/shapenet_base.py
@@ -96,6 +96,8 @@ def _load_mesh(self, model_path) -> Tuple:
                     self.texture_resolution,
                     3,
                 )
+        else:
+            textures = None
 
         return verts, faces.verts_idx, textures
 
diff --git a/setup.py b/setup.py
@@ -142,4 +142,5 @@ def __init__(self, *args, **kwargs):
     },
     ext_modules=get_extensions(),
     cmdclass={"build_ext": BuildExtension},
+    include_package_data=True,
 )
diff --git a/tests/test_shapenet_core.py b/tests/test_shapenet_core.py
@@ -282,3 +282,10 @@ def test_render_shapenet_core(self):
                 "test_shapenet_core_render_without_sample_nums_%s.png" % idx, DATA_DIR
             )
             self.assertClose(mixed_rgb_2, image_ref, atol=0.05)
+
+    def test_load_textures_false(self):
+        shapenet_dataset = ShapeNetCore(
+            SHAPENET_PATH, load_textures=False, version=VERSION
+        )
+        model = shapenet_dataset[0]
+        self.assertIsNone(model["textures"])

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,3 @@`
	`1`	`+include pytorch3d/datasets/shapenet/shapenet_synset_dict_v1.json`
	`2`	`+include pytorch3d/datasets/shapenet/shapenet_synset_dict_v2.json`
	`3`	`+include pytorch3d/datasets/r2n2/r2n2_synset_dict.json`
Original file line number	Diff line number	Diff line change
`@@ -96,6 +96,8 @@ def _load_mesh(self, model_path) -> Tuple:`
`96`	`96`	`self.texture_resolution,`
`97`	`97`	`3,`
`98`	`98`	`)`
	`99`	`+ else:`
	`100`	`+ textures = None`
`99`	`101`
`100`	`102`	`return verts, faces.verts_idx, textures`
`101`	`103`
Original file line number	Diff line number	Diff line change
`@@ -142,4 +142,5 @@ def __init__(self, args, *kwargs):`
`142`	`142`	`},`
`143`	`143`	`ext_modules=get_extensions(),`
`144`	`144`	`cmdclass={"build_ext": BuildExtension},`
	`145`	`+ include_package_data=True,`
`145`	`146`	`)`