From 029a9da00b4015b6988a28e353fc389572ac6254 Mon Sep 17 00:00:00 2001
From: Talmaj Marinc <talmaj@gmail.com>
Date: Fri, 18 Jun 2021 07:01:37 -0700
Subject: [PATCH] Fix ShapeNetDataset (#593)

Summary:
- Add MANIFEST.in and `include_package_data=True` to include dataset .json files in the installation
Fix https://github.com/facebookresearch/pytorch3d/issues/435
- Fix `load_textures=False` for ShapeNetDataset with a test
Fix https://github.com/facebookresearch/pytorch3d/issues/450, partly fix https://github.com/facebookresearch/pytorch3d/issues/444. I've set the textures to `None`, if they should be all white instead, let me know.

Pull Request resolved: https://github.com/facebookresearch/pytorch3d/pull/593

Reviewed By: patricklabatut

Differential Revision: D29116264

Pulled By: nikhilaravi

fbshipit-source-id: 1fb0198e616b7f834dfeaf7168bb5e6e530810d1
---
 MANIFEST.in                         | 3 +++
 pytorch3d/datasets/shapenet_base.py | 2 ++
 setup.py                            | 1 +
 tests/test_shapenet_core.py         | 7 +++++++
 4 files changed, 13 insertions(+)
 create mode 100644 MANIFEST.in

diff --git a/MANIFEST.in b/MANIFEST.in
new file mode 100644
index 00000000..a0d6ffee
--- /dev/null
+++ 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
index 91b82448..24810e1f 100644
--- a/pytorch3d/datasets/shapenet_base.py
+++ b/pytorch3d/datasets/shapenet_base.py
@@ -96,6 +96,8 @@ class ShapeNetBase(torch.utils.data.Dataset):  # pragma: no cover
                     self.texture_resolution,
                     3,
                 )
+        else:
+            textures = None
 
         return verts, faces.verts_idx, textures
 
diff --git a/setup.py b/setup.py
index 28c4081d..a12abb14 100755
--- a/setup.py
+++ b/setup.py
@@ -142,4 +142,5 @@ setup(
     },
     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
index 0268459f..260c95e9 100644
--- a/tests/test_shapenet_core.py
+++ b/tests/test_shapenet_core.py
@@ -282,3 +282,10 @@ class TestShapenetCore(TestCaseMixin, unittest.TestCase):
                 "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"])