Refactor the Shapenet dataset to support different BUILDER_CONFIGS.

PiperOrigin-RevId: 311760835

Author: Aviv Segall

tensorflow_graphics/datasets/shapenet/shapenet.py

@@ -50,55 +50,49 @@
 
 _TAXONOMY_FILE_NAME = 'taxonomy.json'
 
-_MODEL_SUBPATH = os.path.join('models', 'model_normalized.obj')
-
 _SPLIT_FILE_URL = \
     'http://shapenet.cs.stanford.edu/shapenet/obj-zip/SHREC16/all.csv'
 
 _CHECKSUMS_DIR = os.path.normpath(
     os.path.join(os.path.dirname(__file__), 'checksums/'))
 
 
-class Shapenet(tfds.core.GeneratorBasedBuilder):
-  """ShapeNetCore V2.
-
-  Example usage of the dataset:
+class ShapenetConfig(tfds.core.BuilderConfig):
+  """Base class for Shapenet BuilderConfigs.
 
-  import tensorflow_datasets as tfds
-  from tensorflow_graphics.datasets.shapenet import Shapenet
+  The Shapenet database builder delegates the implementation of info,
+  split_generators and generate_examples to the specified ShapenetConfig. This
+  is done to allow multiple versions of the dataset.
+  """
 
-  data_set = Shapenet.load(
-      split='train',
-      download_and_prepare_kwargs={
-          'download_config':
-              tfds.download.DownloadConfig(manual_dir='~/shapenet_base')
-      })
+  def info(self, dataset_builder):
+    """Delegated Shapenet._info."""
+    raise NotImplementedError('Abstract method')
 
-  for example in data_set.take(1):
-    trimesh, label, model_id = example['trimesh'], example['label'],
-                               example['model_id']
+  def split_generators(self, dl_manager, dataset_builder):
+    """Delegated Shapenet._split_generators."""
+    raise NotImplementedError('Abstract method')
 
-  """
+  def generate_examples(self, **kwargs):
+    """Delegated Shapenet._generate_examples."""
+    raise NotImplementedError('Abstract method')
 
-  VERSION = tfds.core.Version('1.0.0')
 
-  @staticmethod
-  def load(*args, **kwargs):
-    return tfds.load('shapenet', *args, **kwargs)
+class MeshConfig(ShapenetConfig):
+  """A Shapenet config for loading the original .obj files."""
 
-  MANUAL_DOWNLOAD_INSTRUCTIONS = textwrap.dedent("""\
-  manual_dir should contain the extracted ShapeNetCore.v2.zip archive.
-  You need to register on https://shapenet.org/download/shapenetcore in order
-  to get the link to download the dataset.
-  """)
+  _MODEL_SUBPATH = os.path.join('models', 'model_normalized.obj')
 
-  def __init__(self, *args, **kwargs):
-    super(Shapenet, self).__init__(*args, **kwargs)
-    tfds.download.add_checksums_dir(_CHECKSUMS_DIR)
+  def __init__(self, model_subpath=_MODEL_SUBPATH):
+    super(MeshConfig, self).__init__(
+        name='shapenet_trimesh',
+        description=_DESCRIPTION,
+        version=tfds.core.Version('1.0.0'))
+    self.model_subpath = model_subpath
 
-  def _info(self):
+  def info(self, dataset_builder):
     return tfds.core.DatasetInfo(
-        builder=self,
+        builder=dataset_builder,
         description=_DESCRIPTION,
         features=tfds_features.FeaturesDict({
             'trimesh': tfg_features.TriangleMesh(),
@@ -111,8 +105,7 @@ def _info(self):
         citation=_CITATION,
     )
 
-  def _split_generators(self, dl_manager):
-    """Returns SplitGenerators."""
+  def split_generators(self, dl_manager, dataset_builder):
     # Extract the synset ids from the taxonomy file and update the ClassLabel
     # feature.
     with tf.io.gfile.GFile(
@@ -122,7 +115,7 @@ def _split_generators(self, dl_manager):
       # Remove duplicate labels (the json file contains two identical entries
       # for synset '04591713').
       labels = list(collections.OrderedDict.fromkeys(labels))
-    self.info.features['label'].names = labels
+    dataset_builder.info.features['label'].names = labels
 
     split_file = dl_manager.download(_SPLIT_FILE_URL)
     fieldnames = ['id', 'synset', 'sub_synset', 'model_id', 'split']
@@ -155,7 +148,7 @@ def _split_generators(self, dl_manager):
         ),
     ]
 
-  def _generate_examples(self, base_dir, models):
+  def generate_examples(self, base_dir, models):
     """Yields examples.
 
     The structure of the examples:
@@ -172,7 +165,8 @@ def _generate_examples(self, base_dir, models):
     for model in models:
       synset = model['synset']
       model_id = model['model_id']
-      model_filepath = os.path.join(base_dir, synset, model_id, _MODEL_SUBPATH)
+      model_filepath = os.path.join(base_dir, synset, model_id,
+                                    self.model_subpath)
       # If the model doesn't exist, skip it.
       if not tf.io.gfile.exists(model_filepath):
         continue
@@ -181,3 +175,54 @@ def _generate_examples(self, base_dir, models):
           'label': synset,
           'model_id': model_id,
       }
+
+
+class Shapenet(tfds.core.GeneratorBasedBuilder):
+  """ShapeNetCore V2.
+
+  Example usage of the dataset:
+
+  import tensorflow_datasets as tfds
+  from tensorflow_graphics.datasets.shapenet import Shapenet
+
+  data_set = Shapenet.load(
+      split='train',
+      download_and_prepare_kwargs={
+          'download_config':
+              tfds.download.DownloadConfig(manual_dir='~/shapenet_base')
+      })
+
+  for example in data_set.take(1):
+    trimesh, label, model_id = example['trimesh'], example['label'],
+                               example['model_id']
+
+  """
+
+  BUILDER_CONFIGS = [MeshConfig()]
+
+  VERSION = tfds.core.Version('1.0.0')
+
+  @staticmethod
+  def load(*args, **kwargs):
+    return tfds.load('shapenet', *args, **kwargs)
+
+  MANUAL_DOWNLOAD_INSTRUCTIONS = textwrap.dedent("""\
+  manual_dir should contain the extracted ShapeNetCore.v2.zip archive.
+  You need to register on https://shapenet.org/download/shapenetcore in order
+  to get the link to download the dataset.
+  """)
+
+  def __init__(self, *args, **kwargs):
+    super(Shapenet, self).__init__(*args, **kwargs)
+    tfds.download.add_checksums_dir(_CHECKSUMS_DIR)
+
+  def _info(self):
+    return self.builder_config.info(self)
+
+  def _split_generators(self, dl_manager):
+    """Returns SplitGenerators."""
+    return self.builder_config.split_generators(dl_manager, self)
+
+  def _generate_examples(self, **kwargs):
+    """Yields examples."""
+    return self.builder_config.generate_examples(**kwargs)