From 38534e1b23844c29b11d5d9cb10fcb028472a736 Mon Sep 17 00:00:00 2001
From: jiehanw <jiehanw@nvidia.com>
Date: Mon, 12 Apr 2021 18:24:33 -0400
Subject: [PATCH 01/13] temp

Signed-off-by: jiehanw <jiehanw@nvidia.com>
---
 kaolin/io/__init__.py |   1 +
 kaolin/io/dataset.py  |  35 +++++++++++
 kaolin/io/shrec.py    | 139 ++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 175 insertions(+)
 create mode 100644 kaolin/io/shrec.py

diff --git a/kaolin/io/__init__.py b/kaolin/io/__init__.py
index 0278efe2e..1fa0a236e 100644
--- a/kaolin/io/__init__.py
+++ b/kaolin/io/__init__.py
@@ -6,3 +6,4 @@
 from . import shapenet
 from . import usd
 from . import modelnet
+from . import shrec
diff --git a/kaolin/io/dataset.py b/kaolin/io/dataset.py
index 594269aaa..2d035c20c 100644
--- a/kaolin/io/dataset.py
+++ b/kaolin/io/dataset.py
@@ -126,6 +126,41 @@ def _get_cache_key(dataset, index):
 
 KaolinDatasetItem = namedtuple('KaolinDatasetItem', ['data', 'attributes'])
 
+class KaolinDataset(Dataset):
+    """A dataset supporting the separation of data and attributes, and combines
+    them in its `__getitem__`.
+    The return value of `__getitem__` will be a named tuple containing the
+    return value of both `get_data` and `get_attributes`.
+    The difference between `get_data` and `get_attributes` is that data are able
+    to be transformed or preprocessed (such as using `ProcessedDataset`), while
+    attributes are generally not.
+    """
+
+    def __getitem__(self, index):
+        """Returns the item at the given index.
+        Will contain a named tuple of both data and attributes.
+        """
+        attributes = self.get_attributes(index)
+        data = self.get_data(index)
+        return KaolinDatasetItem(data=data, attributes=attributes)
+
+    @abstractmethod
+    def get_data(self, index):
+        """Returns the data at the given index."""
+        pass
+
+    @abstractmethod
+    def get_attributes(self, index):
+        """Returns the attributes at the given index.
+        Attributes are usually not transformed by wrappers such as
+        `ProcessedDataset`.
+        """
+        pass
+
+    @abstractmethod
+    def __len__(self):
+        """Returns the number of entries."""
+        pass
 
 class KaolinDataset(Dataset):
     """A dataset supporting the separation of data and attributes, and combines
diff --git a/kaolin/io/shrec.py b/kaolin/io/shrec.py
new file mode 100644
index 000000000..67caa4a11
--- /dev/null
+++ b/kaolin/io/shrec.py
@@ -0,0 +1,139 @@
+# Copyright (c) 2019-2021, NVIDIA CORPORATION. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import warnings
+
+from pathlib import Path
+
+from kaolin.io.dataset import KaolinDataset
+from kaolin.io.obj import import_mesh
+
+VALID_CATEGORIES = [
+    "alien",
+    "ants",
+    "armadillo",
+    "bird1",
+    "bird2",
+    "camel",
+    "cat",
+    "centaur",
+    "dinosaur",
+    "dino_ske",
+    "dog1",
+    "dog2",
+    "flamingo",
+    "glasses",
+    "gorilla",
+    "hand",
+    "horse",
+    "lamp",
+    "laptop",
+    "man",
+    "myScissor",
+    "octopus",
+    "pliers",
+    "rabbit",
+    "santa",
+    "shark",
+    "snake",
+    "spiders",
+    "two_balls",
+    "woman",
+]
+
+class SHREC16(KaolinDataset):
+    r"""Dataset class for SHREC16, used for the "Large-scale 3D shape retrieval
+    from ShapeNet Core55" contest at Eurographics 2016.
+    More details about the challenge and the dataset are available
+    `here <https://shapenet.cs.stanford.edu/shrec16/>`_.
+
+    The `__getitem__` method will return a `KaolinDatasetItem`, with its `data`
+    field containing a `kaolin.io.obj.return_type`.
+
+    Args:
+        root (str): Path to the root directory of the dataset.
+        categories (list): List of categories to load (each class is
+                           specified as a string, and must be a valid `SHREC16`
+                           category). If this argument is not specified, all categories
+                           are loaded by default.
+        train (bool): If True, return the train split, else return the test.
+
+    Returns:
+        .. code-block::
+           dict: {
+                attributes: {path: str, category: str, label: int},
+                data: kaolin.rep.TriangleMesh
+           }
+        path: The filepath to the .obj file on disk.
+        name: The file name of the .obj file on disk.
+        label: A human-readable string describing the loaded sample.
+    
+    Example:
+        >>> dataset = SHREC16(root='/path/to/SHREC16/', categories=['alien', 'ants'], train=False)
+        >>> sample = dataset[0]
+        >>> sample["attributes"]["path"]
+        /path/to/SHREC16/alien/test/T411.obj
+        >>> sample["attributes"]["label"]
+        alien
+    """
+
+    def __init__(self, root: str, categories: Iterable = None, train: bool = True):
+
+        if not categories:
+            categories = VALID_CATEGORIES
+
+        self.root = Path(root)
+        self.paths = []
+        self.labels = []
+
+        for i, category in enumerate(categories):
+            if category not in VALID_CATEGORIES:
+                raise ValueError(
+                    f"Specified category {category} is not valid. "
+                    f"Valid categories are {VALID_CATEGORIES}"
+                )
+
+            clsdir = os.path.join(root, category, "train" if train else "test")
+            curr_models = glob.glob(clsdir + "/*.obj")
+
+            self.paths += curr_models
+            self.labels += [] * len(cur)
+
+            if len(cur) == 0:
+                raise RuntimeWarning(
+                    "No .obj files could be read " f"for category '{cl}'. Skipping..."
+                )
+        
+        self.names = [p.names for p in self.paths]
+
+    def __len__(self):
+        """Returns the length of the dataset. """
+        return len(self.paths)
+
+    def get_data(self, idx):
+        """Returns a Usd Mesh from the path."""
+        obj_location = self.paths[idx]
+        mesh = import_mesh(str(obj_location))
+        return mesh
+
+    def get_attributes(self, idx):
+        attributes = {
+            "path": self.paths[idx],
+            "name": self.names[idx],
+            "label": self.labels[idx],
+        }
+
+        return attributes
+
+    def get_cache_key(self, index):
+        return self.names[index]
\ No newline at end of file

From 9369006ffcbb101c9dd1c988763dc88a292194aa Mon Sep 17 00:00:00 2001
From: jiehanw <jiehanw@nvidia.com>
Date: Wed, 21 Apr 2021 16:51:42 -0400
Subject: [PATCH 02/13] fix shrec docstring

Signed-off-by: jiehanw <jiehanw@nvidia.com>
---
 kaolin/io/shrec.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/kaolin/io/shrec.py b/kaolin/io/shrec.py
index 67caa4a11..596674321 100644
--- a/kaolin/io/shrec.py
+++ b/kaolin/io/shrec.py
@@ -57,8 +57,8 @@ class SHREC16(KaolinDataset):
     More details about the challenge and the dataset are available
     `here <https://shapenet.cs.stanford.edu/shrec16/>`_.
 
-    The `__getitem__` method will return a `KaolinDatasetItem`, with its `data`
-    field containing a `kaolin.io.obj.return_type`.
+    The `__getitem__` method will return a `KaolinDatasetItem`, with its data field 
+    containing a namedtuple returned by :func:`kaolin.io.obj.import_mesh`.
 
     Args:
         root (str): Path to the root directory of the dataset.

From e597e872f785a4174910fddaed8e60ef1d4ed6d4 Mon Sep 17 00:00:00 2001
From: jiehanw <jiehanw@nvidia.com>
Date: Wed, 21 Apr 2021 16:59:08 -0400
Subject: [PATCH 03/13] fix flake8

Signed-off-by: jiehanw <jiehanw@nvidia.com>
---
 kaolin/io/dataset.py | 41 -----------------------------------------
 kaolin/io/shrec.py   |  8 +++-----
 2 files changed, 3 insertions(+), 46 deletions(-)

diff --git a/kaolin/io/dataset.py b/kaolin/io/dataset.py
index 2d035c20c..172d89e47 100644
--- a/kaolin/io/dataset.py
+++ b/kaolin/io/dataset.py
@@ -162,47 +162,6 @@ def __len__(self):
         """Returns the number of entries."""
         pass
 
-class KaolinDataset(Dataset):
-    """A dataset supporting the separation of data and attributes, and combines
-    them in its `__getitem__`.
-
-    The return value of `__getitem__` will be a named tuple containing the
-    return value of both `get_data` and `get_attributes`.
-
-    The difference between `get_data` and `get_attributes` is that data are able
-    to be transformed or preprocessed (such as using `ProcessedDataset`), while
-    attributes are generally not.
-    """
-
-    def __getitem__(self, index):
-        """Returns the item at the given index.
-
-        Will contain a named tuple of both data and attributes.
-        """
-        attributes = self.get_attributes(index)
-        data = self.get_data(index)
-        return KaolinDatasetItem(data=data, attributes=attributes)
-
-    @abstractmethod
-    def get_data(self, index):
-        """Returns the data at the given index."""
-        pass
-
-    @abstractmethod
-    def get_attributes(self, index):
-        """Returns the attributes at the given index.
-
-        Attributes are usually not transformed by wrappers such as
-        `ProcessedDataset`.
-        """
-        pass
-
-    @abstractmethod
-    def __len__(self):
-        """Returns the number of entries."""
-        pass
-
-
 class ProcessedDataset(KaolinDataset):
     def __init__(self, dataset, preprocessing_transform=None,
                  cache_dir=None, num_workers=None, transform=None,
diff --git a/kaolin/io/shrec.py b/kaolin/io/shrec.py
index 596674321..386594e91 100644
--- a/kaolin/io/shrec.py
+++ b/kaolin/io/shrec.py
@@ -11,8 +11,6 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-import warnings
-
 from pathlib import Path
 
 from kaolin.io.dataset import KaolinDataset
@@ -77,7 +75,7 @@ class SHREC16(KaolinDataset):
         path: The filepath to the .obj file on disk.
         name: The file name of the .obj file on disk.
         label: A human-readable string describing the loaded sample.
-    
+
     Example:
         >>> dataset = SHREC16(root='/path/to/SHREC16/', categories=['alien', 'ants'], train=False)
         >>> sample = dataset[0]
@@ -113,7 +111,7 @@ def __init__(self, root: str, categories: Iterable = None, train: bool = True):
                 raise RuntimeWarning(
                     "No .obj files could be read " f"for category '{cl}'. Skipping..."
                 )
-        
+
         self.names = [p.names for p in self.paths]
 
     def __len__(self):
@@ -136,4 +134,4 @@ def get_attributes(self, idx):
         return attributes
 
     def get_cache_key(self, index):
-        return self.names[index]
\ No newline at end of file
+        return self.names[index]

From 750f1aa77508f8ba4946ea89d62d3add8aae861f Mon Sep 17 00:00:00 2001
From: jiehanw <jiehanw@nvidia.com>
Date: Mon, 26 Apr 2021 17:13:17 -0400
Subject: [PATCH 04/13] fix docstring

Signed-off-by: jiehanw <jiehanw@nvidia.com>
---
 kaolin/io/shrec.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/kaolin/io/shrec.py b/kaolin/io/shrec.py
index 386594e91..be8ff6744 100644
--- a/kaolin/io/shrec.py
+++ b/kaolin/io/shrec.py
@@ -85,7 +85,7 @@ class SHREC16(KaolinDataset):
         alien
     """
 
-    def __init__(self, root: str, categories: Iterable = None, train: bool = True):
+    def __init__(self, root, categories=None, train=true):
 
         if not categories:
             categories = VALID_CATEGORIES

From 6d7a4fe96634442f1b250714c3eb617b981d5523 Mon Sep 17 00:00:00 2001
From: jiehanw <jiehanw@nvidia.com>
Date: Mon, 26 Apr 2021 17:22:53 -0400
Subject: [PATCH 05/13] fix docstring

Signed-off-by: jiehanw <jiehanw@nvidia.com>
---
 kaolin/io/shrec.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/kaolin/io/shrec.py b/kaolin/io/shrec.py
index be8ff6744..9011e5649 100644
--- a/kaolin/io/shrec.py
+++ b/kaolin/io/shrec.py
@@ -85,7 +85,7 @@ class SHREC16(KaolinDataset):
         alien
     """
 
-    def __init__(self, root, categories=None, train=true):
+    def __init__(self, root, categories=None, train=True):
 
         if not categories:
             categories = VALID_CATEGORIES

From 6a9290967e1ff1fb74d5c18ece4c284aabebff40 Mon Sep 17 00:00:00 2001
From: jiehanw <jiehanw@nvidia.com>
Date: Mon, 26 Apr 2021 17:48:04 -0400
Subject: [PATCH 06/13] fix docstring warning

Signed-off-by: jiehanw <jiehanw@nvidia.com>
---
 kaolin/io/shrec.py | 18 ------------------
 1 file changed, 18 deletions(-)

diff --git a/kaolin/io/shrec.py b/kaolin/io/shrec.py
index 9011e5649..d4693da60 100644
--- a/kaolin/io/shrec.py
+++ b/kaolin/io/shrec.py
@@ -65,24 +65,6 @@ class SHREC16(KaolinDataset):
                            category). If this argument is not specified, all categories
                            are loaded by default.
         train (bool): If True, return the train split, else return the test.
-
-    Returns:
-        .. code-block::
-           dict: {
-                attributes: {path: str, category: str, label: int},
-                data: kaolin.rep.TriangleMesh
-           }
-        path: The filepath to the .obj file on disk.
-        name: The file name of the .obj file on disk.
-        label: A human-readable string describing the loaded sample.
-
-    Example:
-        >>> dataset = SHREC16(root='/path/to/SHREC16/', categories=['alien', 'ants'], train=False)
-        >>> sample = dataset[0]
-        >>> sample["attributes"]["path"]
-        /path/to/SHREC16/alien/test/T411.obj
-        >>> sample["attributes"]["label"]
-        alien
     """
 
     def __init__(self, root, categories=None, train=True):

From 5cd3172611e82ba40088f5165896459d6b4b0e35 Mon Sep 17 00:00:00 2001
From: jiehanw <jiehanw@nvidia.com>
Date: Mon, 26 Apr 2021 18:08:23 -0400
Subject: [PATCH 07/13] add shrec.rst to toctree

Signed-off-by: jiehanw <jiehanw@nvidia.com>
---
 docs/modules/kaolin.io.rst | 1 +
 1 file changed, 1 insertion(+)

diff --git a/docs/modules/kaolin.io.rst b/docs/modules/kaolin.io.rst
index 5cd2dbc1f..4619fc186 100644
--- a/docs/modules/kaolin.io.rst
+++ b/docs/modules/kaolin.io.rst
@@ -21,3 +21,4 @@ and :ref:`materials module<kaolin.io.materials>` contains Materials definition t
    kaolin.io.shapenet
    kaolin.io.usd
    kaolin.io.modelnet
+   kaolin.io.shrec

From 2cca1820352dc4f2c96668bf7d3b0f5f69d4e3fc Mon Sep 17 00:00:00 2001
From: jiehanw <jiehanw@nvidia.com>
Date: Tue, 27 Apr 2021 15:24:17 -0400
Subject: [PATCH 08/13] remove with_materials since shrec16 does not have
 materials

Signed-off-by: jiehanw <jiehanw@nvidia.com>
---
 kaolin/io/shrec.py                   | 173 +++++++++++++++++----------
 tests/python/kaolin/io/test_shrec.py |  76 ++++++++++++
 2 files changed, 184 insertions(+), 65 deletions(-)
 create mode 100644 tests/python/kaolin/io/test_shrec.py

diff --git a/kaolin/io/shrec.py b/kaolin/io/shrec.py
index d4693da60..3917c907f 100644
--- a/kaolin/io/shrec.py
+++ b/kaolin/io/shrec.py
@@ -11,43 +11,84 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import warnings
+
 from pathlib import Path
 
 from kaolin.io.dataset import KaolinDataset
 from kaolin.io.obj import import_mesh
 
-VALID_CATEGORIES = [
-    "alien",
-    "ants",
-    "armadillo",
-    "bird1",
-    "bird2",
-    "camel",
-    "cat",
-    "centaur",
-    "dinosaur",
-    "dino_ske",
-    "dog1",
-    "dog2",
-    "flamingo",
-    "glasses",
-    "gorilla",
-    "hand",
-    "horse",
-    "lamp",
-    "laptop",
-    "man",
-    "myScissor",
-    "octopus",
-    "pliers",
-    "rabbit",
-    "santa",
-    "shark",
-    "snake",
-    "spiders",
-    "two_balls",
-    "woman",
-]
+synset_to_label = {
+    '02691156': 'airplane', 
+    '02747177': 'ashcan', 
+    '02773838': 'bag', 
+    '02801938': 'basket', 
+    '02808440': 'bathtub', 
+    '02818832': 'bed', 
+    '02828884': 'bench', 
+    '02834778': 'bicycle', 
+    '02843684': 'birdhouse', 
+    '02871439': 'bookshelf', 
+    '02876657': 'bottle', 
+    '02880940': 'bowl', 
+    '02924116': 'bus', 
+    '02933112': 'cabinet', 
+    '02942699': 'camera', 
+    '02946921': 'can', 
+    '02954340': 'cap', 
+    '02958343': 'car', 
+    '03001627': 'chair', 
+    '03046257': 'clock', 
+    '03085013': 'computer keyboard', 
+    '03207941': 'dishwasher', 
+    '03211117': 'display', 
+    '03261776': 'earphone', 
+    '03325088': 'faucet', 
+    '03337140': 'file', 
+    '03467517': 'guitar', 
+    '03513137': 'helmet', 
+    '03593526': 'jar', 
+    '03624134': 'knife', 
+    '03636649': 'lamp', 
+    '03642806': 'laptop', 
+    '03691459': 'loudspeaker', 
+    '03710193': 'mailbox', 
+    '03759954': 'microphone', 
+    '03761084': 'microwave', 
+    '03790512': 'motorcycle', 
+    '03797390': 'mug', 
+    '03928116': 'piano', 
+    '03938244': 'pillow', 
+    '03948459': 'pistol', 
+    '03991062': 'pot', 
+    '04004475': 'printer', 
+    '04074963': 'remote control', 
+    '04090263': 'rifle', 
+    '04099429': 'rocket', 
+    '04225987': 'skateboard', 
+    '04256520': 'sofa', 
+    '04330267': 'stove', 
+    '04379243': 'table', 
+    '04401088': 'telephone', 
+    '04460130': 'tower', 
+    '04468005': 'train', 
+    '04530566': 'vessel', 
+    '04554684': 'washer', 
+    '04591713': 'wine bottle'}
+# Label to Synset mapping (for ShapeNet core classes)
+label_to_synset = {v: k for k, v in synset_to_label.items()}
+
+def _convert_categories(categories):
+    if categories is None:
+        synset = [value for key, value in label_to_synset.items()]
+    else:
+        if not (c in synset_to_label.keys() + label_to_synset.keys()
+                for c in categories):
+            warnings.warn('Some or all of the categories requested are not part of \
+                ShapeNetCore. Data loading may fail if these categories are not avaliable.')
+        synsets = [label_to_synset[c] if c in label_to_synset.keys()
+                   else c for c in categories]
+    return synsets
 
 class SHREC16(KaolinDataset):
     r"""Dataset class for SHREC16, used for the "Large-scale 3D shape retrieval
@@ -55,8 +96,8 @@ class SHREC16(KaolinDataset):
     More details about the challenge and the dataset are available
     `here <https://shapenet.cs.stanford.edu/shrec16/>`_.
 
-    The `__getitem__` method will return a `KaolinDatasetItem`, with its data field 
-    containing a namedtuple returned by :func:`kaolin.io.obj.import_mesh`.
+    The `__getitem__` method will return a `KaolinDatasetItem`, with its `data`
+    field containing a `kaolin.io.obj.return_type`.
 
     Args:
         root (str): Path to the root directory of the dataset.
@@ -64,55 +105,57 @@ class SHREC16(KaolinDataset):
                            specified as a string, and must be a valid `SHREC16`
                            category). If this argument is not specified, all categories
                            are loaded by default.
-        train (bool): If True, return the train split, else return the test.
+        split (str): String to indicate whether to load train, test or val set.
     """
 
-    def __init__(self, root, categories=None, train=True):
-
-        if not categories:
-            categories = VALID_CATEGORIES
+    def __init__(self, root: str, categories: list = None, split: str = "train"):
 
         self.root = Path(root)
         self.paths = []
-        self.labels = []
-
-        for i, category in enumerate(categories):
-            if category not in VALID_CATEGORIES:
+        self.synset_idxs = []
+        self.synsets = _convert_categories(categories)
+        self.labels = [synset_to_label[s] for s in self.synsets]
+
+        # loops through desired classes
+        for i in range(len(self.synsets)):
+            syn = self.synsets[i]
+
+            if split == "train":
+                class_target = self.root / "train" / syn
+            elif split == "test":
+                class_target = self.root / "test" / syn
+            else:
+                class_target = self.root / "val" / syn
+
+            if not class_target.exists():
                 raise ValueError(
-                    f"Specified category {category} is not valid. "
-                    f"Valid categories are {VALID_CATEGORIES}"
-                )
+                    'Class {0} ({1}) was not found at location {2}.'.format(
+                        syn, self.labels[i], str(class_target)))
 
-            clsdir = os.path.join(root, category, "train" if train else "test")
-            curr_models = glob.glob(clsdir + "/*.obj")
+            # find all objects in the class
+            models = sorted(class_target.glob('*'))
 
-            self.paths += curr_models
-            self.labels += [] * len(cur)
+            self.paths += models
+            self.synset_idxs += [i] * len(models)
 
-            if len(cur) == 0:
-                raise RuntimeWarning(
-                    "No .obj files could be read " f"for category '{cl}'. Skipping..."
-                )
-
-        self.names = [p.names for p in self.paths]
+        self.names = [p.name for p in self.paths]
 
     def __len__(self):
-        """Returns the length of the dataset. """
         return len(self.paths)
 
-    def get_data(self, idx):
-        """Returns a Usd Mesh from the path."""
-        obj_location = self.paths[idx]
+    def get_data(self, index):
+        obj_location = self.paths[index]
         mesh = import_mesh(str(obj_location))
         return mesh
 
-    def get_attributes(self, idx):
+    def get_attributes(self, index):
+        synset_idx = self.synset_idxs[index]
         attributes = {
-            "path": self.paths[idx],
-            "name": self.names[idx],
-            "label": self.labels[idx],
+            'name': self.names[index],
+            'path': self.paths[index],
+            'synset': self.synsets[synset_idx],
+            'label': self.labels[synset_idx]
         }
-
         return attributes
 
     def get_cache_key(self, index):
diff --git a/tests/python/kaolin/io/test_shrec.py b/tests/python/kaolin/io/test_shrec.py
new file mode 100644
index 000000000..5db5437f5
--- /dev/null
+++ b/tests/python/kaolin/io/test_shrec.py
@@ -0,0 +1,76 @@
+# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from pathlib import Path
+
+import os
+
+import pytest
+import torch
+
+from kaolin.io.obj import return_type
+from kaolin.io.shrec import SHREC16
+
+SHREC16_PATH = '/home/jiehanw/Downloads/shrec16'
+SHREC16_TEST_CATEGORY_SYNSETS = ['02691156']
+SHREC16_TEST_CATEGORY_LABELS = ['airplane']
+SHREC16_TEST_CATEGORY_SYNSETS_2 = ['02958343']
+SHREC16_TEST_CATEGORY_LABELS_2 = ['car']
+SHREC16_TEST_CATEGORY_SYNSETS_MULTI = ['02691156', '02958343']
+SHREC16_TEST_CATEGORY_LABELS_MULTI = ['airplane', 'car']
+
+ALL_CATEGORIES = [
+    SHREC16_TEST_CATEGORY_SYNSETS,
+    SHREC16_TEST_CATEGORY_LABELS,
+    SHREC16_TEST_CATEGORY_SYNSETS_2,
+    SHREC16_TEST_CATEGORY_LABELS_2,
+    SHREC16_TEST_CATEGORY_SYNSETS_MULTI,
+    SHREC16_TEST_CATEGORY_LABELS_MULTI,
+]
+
+# Skip test in a CI environment
+@pytest.mark.skipif(os.getenv('CI') == 'true', reason="CI does not have dataset")
+@pytest.mark.parametrize('categories', ALL_CATEGORIES)
+@pytest.mark.parametrize('split', ['val'])
+@pytest.mark.parametrize('index', [0, -1])
+class TestSHREC16(object):
+
+    @pytest.fixture(autouse=True)
+    def shrec16_dataset(self, categories, split):
+        return SHREC16(root=SHREC16_PATH,
+                       categories=categories,
+                       split=split)
+
+    def test_basic_getitem(self, shrec16_dataset, index):
+        assert len(shrec16_dataset) > 0
+
+        if index == -1:
+            index = len(shrec16_dataset) - 1
+
+        item = shrec16_dataset[index]
+        data = item.data
+        attributes = item.attributes
+        assert isinstance(data, return_type)
+        assert isinstance(attributes, dict)
+
+        assert isinstance(data.vertices, torch.Tensor)
+        assert len(data.vertices.shape) == 2
+        assert data.vertices.shape[1] == 3
+        assert isinstance(data.faces, torch.Tensor)
+        assert len(data.faces.shape) == 2
+
+        assert isinstance(attributes['name'], str)
+        assert isinstance(attributes['path'], Path)
+        assert isinstance(attributes['synset'], str)
+        assert isinstance(attributes['label'], str)

From 4f9c241081b3cb1c7c494741baf2a183682b3da4 Mon Sep 17 00:00:00 2001
From: jiehanw <jiehanw@nvidia.com>
Date: Tue, 27 Apr 2021 15:36:21 -0400
Subject: [PATCH 09/13] change test path to absolute path

Signed-off-by: jiehanw <jiehanw@nvidia.com>
---
 tests/python/kaolin/io/test_shrec.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tests/python/kaolin/io/test_shrec.py b/tests/python/kaolin/io/test_shrec.py
index 5db5437f5..da585aac6 100644
--- a/tests/python/kaolin/io/test_shrec.py
+++ b/tests/python/kaolin/io/test_shrec.py
@@ -22,7 +22,7 @@
 from kaolin.io.obj import return_type
 from kaolin.io.shrec import SHREC16
 
-SHREC16_PATH = '/home/jiehanw/Downloads/shrec16'
+SHREC16_PATH = '/data/shrec16'
 SHREC16_TEST_CATEGORY_SYNSETS = ['02691156']
 SHREC16_TEST_CATEGORY_LABELS = ['airplane']
 SHREC16_TEST_CATEGORY_SYNSETS_2 = ['02958343']
@@ -43,7 +43,6 @@
 @pytest.mark.skipif(os.getenv('CI') == 'true', reason="CI does not have dataset")
 @pytest.mark.parametrize('categories', ALL_CATEGORIES)
 @pytest.mark.parametrize('split', ['val'])
-@pytest.mark.parametrize('index', [0, -1])
 class TestSHREC16(object):
 
     @pytest.fixture(autouse=True)
@@ -52,6 +51,7 @@ def shrec16_dataset(self, categories, split):
                        categories=categories,
                        split=split)
 
+    @pytest.mark.parametrize('index', [0, -1])
     def test_basic_getitem(self, shrec16_dataset, index):
         assert len(shrec16_dataset) > 0
 

From 03632884ec28eeacc216d7a952c0c5f82593c253 Mon Sep 17 00:00:00 2001
From: jiehanw <jiehanw@nvidia.com>
Date: Wed, 28 Apr 2021 10:33:53 -0400
Subject: [PATCH 10/13] empty

Signed-off-by: jiehanw <jiehanw@nvidia.com>
---
 tests/python/kaolin/io/test_shrec.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tests/python/kaolin/io/test_shrec.py b/tests/python/kaolin/io/test_shrec.py
index da585aac6..242d18b0c 100644
--- a/tests/python/kaolin/io/test_shrec.py
+++ b/tests/python/kaolin/io/test_shrec.py
@@ -39,6 +39,7 @@
     SHREC16_TEST_CATEGORY_LABELS_MULTI,
 ]
 
+
 # Skip test in a CI environment
 @pytest.mark.skipif(os.getenv('CI') == 'true', reason="CI does not have dataset")
 @pytest.mark.parametrize('categories', ALL_CATEGORIES)

From dd1b4379ccdeb417459f0d15c146868d825437c8 Mon Sep 17 00:00:00 2001
From: jiehanw <jiehanw@nvidia.com>
Date: Wed, 2 Jun 2021 15:00:42 -0400
Subject: [PATCH 11/13] change shrec16 synsetids

Signed-off-by: jiehanw <jiehanw@nvidia.com>
---
 kaolin/io/shrec.py                   | 191 +++++++++++++++------------
 tests/python/kaolin/io/test_shrec.py |  36 ++++-
 2 files changed, 135 insertions(+), 92 deletions(-)

diff --git a/kaolin/io/shrec.py b/kaolin/io/shrec.py
index 3917c907f..06697aa68 100644
--- a/kaolin/io/shrec.py
+++ b/kaolin/io/shrec.py
@@ -16,78 +16,77 @@
 from pathlib import Path
 
 from kaolin.io.dataset import KaolinDataset
-from kaolin.io.obj import import_mesh
-
-synset_to_label = {
-    '02691156': 'airplane', 
-    '02747177': 'ashcan', 
-    '02773838': 'bag', 
-    '02801938': 'basket', 
-    '02808440': 'bathtub', 
-    '02818832': 'bed', 
-    '02828884': 'bench', 
-    '02834778': 'bicycle', 
-    '02843684': 'birdhouse', 
-    '02871439': 'bookshelf', 
-    '02876657': 'bottle', 
-    '02880940': 'bowl', 
-    '02924116': 'bus', 
-    '02933112': 'cabinet', 
-    '02942699': 'camera', 
-    '02946921': 'can', 
-    '02954340': 'cap', 
-    '02958343': 'car', 
-    '03001627': 'chair', 
-    '03046257': 'clock', 
-    '03085013': 'computer keyboard', 
-    '03207941': 'dishwasher', 
-    '03211117': 'display', 
-    '03261776': 'earphone', 
-    '03325088': 'faucet', 
-    '03337140': 'file', 
-    '03467517': 'guitar', 
-    '03513137': 'helmet', 
-    '03593526': 'jar', 
-    '03624134': 'knife', 
-    '03636649': 'lamp', 
-    '03642806': 'laptop', 
-    '03691459': 'loudspeaker', 
-    '03710193': 'mailbox', 
-    '03759954': 'microphone', 
-    '03761084': 'microwave', 
-    '03790512': 'motorcycle', 
-    '03797390': 'mug', 
-    '03928116': 'piano', 
-    '03938244': 'pillow', 
-    '03948459': 'pistol', 
-    '03991062': 'pot', 
-    '04004475': 'printer', 
-    '04074963': 'remote control', 
-    '04090263': 'rifle', 
-    '04099429': 'rocket', 
-    '04225987': 'skateboard', 
-    '04256520': 'sofa', 
-    '04330267': 'stove', 
-    '04379243': 'table', 
-    '04401088': 'telephone', 
-    '04460130': 'tower', 
-    '04468005': 'train', 
-    '04530566': 'vessel', 
-    '04554684': 'washer', 
-    '04591713': 'wine bottle'}
+from kaolin.io.obj import import_mesh, ignore_error_handler
+
+synset_to_labels = {
+    '03790512': ['motorcycle', 'bike'],
+    '02808440': ['bathtub', 'bathing tub', 'bath', 'tub'],
+    '02871439': ['bookshelf'],
+    '03761084': ['microwave', 'microwave oven'],
+    '04530566': ['vessel', 'watercraft'],
+    '02691156': ['airplane', 'aeroplane', 'plane'],
+    '04379243': ['table'],
+    '03337140': ['file', 'file cabinet', 'filing cabinet'],
+    '04256520': ['sofa', 'couch', 'lounge'],
+    '03636649': ['lamp'],
+    '03928116': ['piano', 'pianoforte', 'forte-piano'],
+    '04004475': ['printer', 'printing machine'],
+    '03593526': ['jar'],
+    '04330267': ['stove'],
+    '04554684': ['washer', 'automatic washer', 'washing machine'],
+    '03948459': ['pistol', 'handgun', 'side arm', 'shooting iron'],
+    '03001627': ['chair'],
+    '03797390': ['mug'],
+    '02801938': ['basket', 'handbasket'],
+    '03710193': ['mailbox', 'letter box'],
+    '03938244': ['pillow'],
+    '03624134': ['knife'],
+    '02954340': ['cap'],
+    '02773838': ['bag', 'traveling bag', 'travelling bag', 'grip', 'suitcase'],
+    '02747177': ['ashcan', 'trash can', 'garbage can', 'wastebin', 
+                 'ash bin', 'ash-bin', 'ashbin', 'dustbin', 'trash barrel', 'trash bin'],
+    '04460130': ['tower'],
+    '02933112': ['cabinet'],
+    '02876657': ['bottle'],
+    '03991062': ['pot', 'flowerpot'],
+    '02843684': ['birdhouse'],
+    '02818832': ['bed'],
+    '02958343': ['car', 'auto', 'automobile', 'machine', 'motorcar'],
+    '03642806': ['laptop', 'laptop computer'],
+    '03085013': ['computer keyboard', 'keypad'],
+    '04074963': ['remote control', 'remote'],
+    '02924116': ['bus', 'autobus', 'coach', 'charabanc', 'double-decker', 
+                 'jitney', 'motorbus', 'motorcoach', 'omnibus', 'passenger vehi'],
+    '04225987': ['skateboard'],
+    '03261776': ['earphone', 'earpiece', 'headphone', 'phone'],
+    '02880940': ['bowl'],
+    '03325088': ['faucet', 'spigot'],
+    '03211117': ['display', 'video display'],
+    '04468005': ['train', 'railroad train'],
+    '03691459': ['loudspeaker', 'speaker', 'speaker unit', 'loudspeaker system', 'speaker system'],
+    '04090263': ['rifle'],
+    '02946921': ['can', 'tin', 'tin can'],
+    '04099429': ['rocket', 'projectile'],
+    '03467517': ['guitar'],
+    '04401088': ['telephone', 'phone', 'telephone set'],
+    '03046257': ['clock'],
+    '03759954': ['microphone', 'mike'],
+    '03513137': ['helmet'],
+    '02834778': ['bicycle', 'bike', 'wheel', 'cycle'],
+    '03207941': ['dishwasher', 'dish washer', 'dishwashing machine'],
+    '02828884': ['bench'],
+    '02942699': ['camera', 'photographic camera']}
+
 # Label to Synset mapping (for ShapeNet core classes)
-label_to_synset = {v: k for k, v in synset_to_label.items()}
+label_to_synset = {label: synset for synset, labels in synset_to_labels.items() for label in labels}
 
 def _convert_categories(categories):
-    if categories is None:
-        synset = [value for key, value in label_to_synset.items()]
-    else:
-        if not (c in synset_to_label.keys() + label_to_synset.keys()
-                for c in categories):
-            warnings.warn('Some or all of the categories requested are not part of \
-                ShapeNetCore. Data loading may fail if these categories are not avaliable.')
-        synsets = [label_to_synset[c] if c in label_to_synset.keys()
-                   else c for c in categories]
+    if not (c in synset_to_label.keys() + label_to_synset.keys()
+            for c in categories):
+        warnings.warn('Some or all of the categories requested are not part of \
+            Shrec16. Data loading may fail if these categories are not avaliable.')
+    synsets = [label_to_synset[c] if c in label_to_synset.keys()
+               else c for c in categories]
     return synsets
 
 class SHREC16(KaolinDataset):
@@ -113,30 +112,48 @@ def __init__(self, root: str, categories: list = None, split: str = "train"):
         self.root = Path(root)
         self.paths = []
         self.synset_idxs = []
-        self.synsets = _convert_categories(categories)
-        self.labels = [synset_to_label[s] for s in self.synsets]
-
-        # loops through desired classes
-        for i in range(len(self.synsets)):
-            syn = self.synsets[i]
 
-            if split == "train":
-                class_target = self.root / "train" / syn
-            elif split == "test":
-                class_target = self.root / "test" / syn
+        if split == "test":
+            # Setting synsets and labels to None if in test split
+            self.synsets = [None]
+            self.labels = [None]
+        else:
+            if categories is None:
+                self.synsets = list(synset_to_labels.keys())
             else:
-                class_target = self.root / "val" / syn
-
-            if not class_target.exists():
-                raise ValueError(
-                    'Class {0} ({1}) was not found at location {2}.'.format(
-                        syn, self.labels[i], str(class_target)))
+                self.synsets = _convert_categories(categories)
+            self.labels = [synset_to_labels[s] for s in self.synsets]
 
+        # loops through desired classes
+        if split == "test":
+            class_target = self.root / "test"
             # find all objects in the class
             models = sorted(class_target.glob('*'))
 
             self.paths += models
-            self.synset_idxs += [i] * len(models)
+            self.synset_idxs += [0] * len(models)
+
+        else:
+            for i in range(len(self.synsets)):
+                syn = self.synsets[i]
+
+                if split == "train":
+                    class_target = self.root / "train" / syn
+                elif split == "val":
+                    class_target = self.root / "val" / syn
+                else:
+                    raise ValueError(f'Split must be either train, test or val, got {split} instead.')
+
+                if not class_target.exists():
+                    raise ValueError(
+                        'Class {0} ({1}) was not found at location {2}.'.format(
+                            syn, self.labels[i], str(class_target)))
+
+                # find all objects in the class
+                models = sorted(class_target.glob('*'))
+
+                self.paths += models
+                self.synset_idxs += [i] * len(models)
 
         self.names = [p.name for p in self.paths]
 
@@ -145,7 +162,7 @@ def __len__(self):
 
     def get_data(self, index):
         obj_location = self.paths[index]
-        mesh = import_mesh(str(obj_location))
+        mesh = import_mesh(str(obj_location), error_handler=ignore_error_handler)
         return mesh
 
     def get_attributes(self, index):
diff --git a/tests/python/kaolin/io/test_shrec.py b/tests/python/kaolin/io/test_shrec.py
index 242d18b0c..5bd52a54b 100644
--- a/tests/python/kaolin/io/test_shrec.py
+++ b/tests/python/kaolin/io/test_shrec.py
@@ -22,7 +22,7 @@
 from kaolin.io.obj import return_type
 from kaolin.io.shrec import SHREC16
 
-SHREC16_PATH = '/data/shrec16'
+SHREC16_PATH = '/home/jiehanw/Downloads/shrec16'
 SHREC16_TEST_CATEGORY_SYNSETS = ['02691156']
 SHREC16_TEST_CATEGORY_LABELS = ['airplane']
 SHREC16_TEST_CATEGORY_SYNSETS_2 = ['02958343']
@@ -31,6 +31,7 @@
 SHREC16_TEST_CATEGORY_LABELS_MULTI = ['airplane', 'car']
 
 ALL_CATEGORIES = [
+    None,
     SHREC16_TEST_CATEGORY_SYNSETS,
     SHREC16_TEST_CATEGORY_LABELS,
     SHREC16_TEST_CATEGORY_SYNSETS_2,
@@ -43,7 +44,7 @@
 # Skip test in a CI environment
 @pytest.mark.skipif(os.getenv('CI') == 'true', reason="CI does not have dataset")
 @pytest.mark.parametrize('categories', ALL_CATEGORIES)
-@pytest.mark.parametrize('split', ['val'])
+@pytest.mark.parametrize('split', ['train', 'val', 'test'])
 class TestSHREC16(object):
 
     @pytest.fixture(autouse=True)
@@ -53,7 +54,7 @@ def shrec16_dataset(self, categories, split):
                        split=split)
 
     @pytest.mark.parametrize('index', [0, -1])
-    def test_basic_getitem(self, shrec16_dataset, index):
+    def test_basic_getitem(self, shrec16_dataset, index, split):
         assert len(shrec16_dataset) > 0
 
         if index == -1:
@@ -73,5 +74,30 @@ def test_basic_getitem(self, shrec16_dataset, index):
 
         assert isinstance(attributes['name'], str)
         assert isinstance(attributes['path'], Path)
-        assert isinstance(attributes['synset'], str)
-        assert isinstance(attributes['label'], str)
+        
+        if split == "test":
+            assert attributes['synset'] is None
+            assert attributes['label'] is None
+        else:
+            assert isinstance(attributes['synset'], str)
+            assert isinstance(attributes['label'], list)
+    
+    @pytest.mark.parametrize('index', [-1, -2])
+    def test_neg_index(self, shrec16_dataset, index):
+
+        assert len(shrec16_dataset) > 0
+
+        gt_item = shrec16_dataset[len(shrec16_dataset) + index]
+        gt_data = gt_item.data
+        gt_attributes = gt_item.attributes
+
+        item = shrec16_dataset[index]
+        data = item.data
+        attributes = item.attributes
+
+        assert torch.equal(data.vertices, gt_data.vertices)
+        assert torch.equal(data.faces, gt_data.faces)
+
+        assert attributes['name'] == gt_attributes['name']
+        assert attributes['path'] == gt_attributes['path']
+        assert attributes['synset'] == gt_attributes['synset']

From d07ad9c223bfb41c329192dd60020cba2e27e336 Mon Sep 17 00:00:00 2001
From: Jiehan Wang <jiehanw@nvidia.com>
Date: Mon, 14 Jun 2021 16:09:57 -0400
Subject: [PATCH 12/13] fix testing path

Signed-off-by: Jiehan Wang <jiehanw@nvidia.com>
---
 tests/python/kaolin/io/test_shrec.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/python/kaolin/io/test_shrec.py b/tests/python/kaolin/io/test_shrec.py
index 5bd52a54b..39dc79884 100644
--- a/tests/python/kaolin/io/test_shrec.py
+++ b/tests/python/kaolin/io/test_shrec.py
@@ -22,7 +22,7 @@
 from kaolin.io.obj import return_type
 from kaolin.io.shrec import SHREC16
 
-SHREC16_PATH = '/home/jiehanw/Downloads/shrec16'
+SHREC16_PATH = '/data/shrec16/'
 SHREC16_TEST_CATEGORY_SYNSETS = ['02691156']
 SHREC16_TEST_CATEGORY_LABELS = ['airplane']
 SHREC16_TEST_CATEGORY_SYNSETS_2 = ['02958343']

From b29b72ad696a126862e409505239d1a41e36c630 Mon Sep 17 00:00:00 2001
From: Jiehan Wang <jiehanw@nvidia.com>
Date: Wed, 16 Jun 2021 09:47:34 -0400
Subject: [PATCH 13/13] change test split path

Signed-off-by: Jiehan Wang <jiehanw@nvidia.com>
---
 kaolin/io/shrec.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/kaolin/io/shrec.py b/kaolin/io/shrec.py
index 06697aa68..f7e9310af 100644
--- a/kaolin/io/shrec.py
+++ b/kaolin/io/shrec.py
@@ -125,7 +125,7 @@ def __init__(self, root: str, categories: list = None, split: str = "train"):
             self.labels = [synset_to_labels[s] for s in self.synsets]
 
         # loops through desired classes
-        if split == "test":
+        if split == "test_allinone":
             class_target = self.root / "test"
             # find all objects in the class
             models = sorted(class_target.glob('*'))