"""pix3d dataset."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import tensorflow_datasets.public_api as tfds

# TODO(pix3d): BibTeX citation
_CITATION = """@inproceedings{pix3d,
  title={Pix3D: Dataset and Methods for Single-Image 3D Shape Modeling},
  author={Sun, Xingyuan and Wu, Jiajun and Zhang, Xiuming and Zhang, Zhoutong and Zhang, Chengkai and Xue, Tianfan and Tenenbaum, Joshua B and Freeman, William T},
  booktitle={IEEE Conference on Computer Vision and Pattern Recognition (CVPR)},
  year={2018}
}
"""

# TODO(pix3d):
_DESCRIPTION = """Pix3D is a large-scale dataset of diverse image-shape pairs with pixel-level 2D-3D alignment. 
It has wide applications in shape-related tasks including reconstruction, retrieval, viewpoint estimation, etc.

Pix3D contains 10,069 2D-3D pairs of 395 distinct 3D shapes, categorised into nine object categories. 
Each sample comprises of an image, 3D shape represented as (non-watertight) triangle mesh and voxel grid, 
bounding-box, segmentation mask, intrinsic and extrinsic camera parameters and 2D and 3D key points. 
"""


class Pix3d(tfds.core.GeneratorBasedBuilder):
  """TODO(pix3d): Short description of my dataset."""

  # TODO(pix3d): Set up version.
  VERSION = tfds.core.Version('0.1.0')

  def _info(self):
    # TODO(pix3d): Specifies the tfds.core.DatasetInfo object
    return tfds.core.DatasetInfo(
        builder=self,
        # This is the description that will appear on the datasets page.
        description=_DESCRIPTION,
        # tfds.features.FeatureConnectors
        features=tfds.features.FeaturesDict({
            # These are the features of your dataset like images, labels ...
        }),
        # If there's a common (input, target) tuple from the features,
        # specify them here. They'll be used if as_supervised=True in
        # builder.as_dataset.
        supervised_keys=(),
        # Homepage of the dataset for documentation
        homepage='https://dataset-homepage/',
        citation=_CITATION,
    )

  def _split_generators(self, dl_manager):
    """Returns SplitGenerators."""
    # TODO(pix3d): Downloads the data and defines the splits
    # dl_manager is a tfds.download.DownloadManager that can be used to
    # download and extract URLs
    return [
        tfds.core.SplitGenerator(
            name=tfds.Split.TRAIN,
            # These kwargs will be passed to _generate_examples
            gen_kwargs={},
        ),
    ]

  def _generate_examples(self):
    """Yields examples."""
    # TODO(pix3d): Yields (key, example) tuples from the dataset
    yield 'key', {}


@@ -146,68 +146,196 @@
 flow
-_datasets.public_api as tfds%0A%0A# TODO(pix3d): BibTeX citation
+.compat.v2 as tf%0A%0Aimport tensorflow_datasets.public_api as tfds%0Afrom tensorflow_datasets import features as tfds_features%0Afrom tensorflow_graphics.datasets import features as tfg_features%0A
 %0A_CI
@@ -711,23 +711,8 @@
 %22%22%0A%0A
-# TODO(pix3d):%0A
 _DES
@@ -1306,51 +1306,99 @@
  %22%22%22
-TODO(pix3d): Short description of my datase
+Pix3D is a large-scale dataset of diverse image-shape pairs with pixel-level 2D-3D alignmen
 t.%22%22
@@ -1819,74 +1819,908 @@
     
-  # These are the features of your dataset like images, labels ...
+'image': tfds_features.Image(shape=(None, None, 3), dtype=tf.uint8),%0A          'image/filename': tfds_features.Text(),%0A          'image/source': tfds_features.Text(),%0A          '2d_keypoints': tfds_features.Tensor(shape=(None, None, 2), dtype=tf.float32),%0A          'mask': tfds_features.Image(shape=(None, None, 1), dtype=tf.float32),%0A          'model': tfg_features.TriangleMesh(),%0A          'model/source': tfds_features.Text(),%0A          '3d_keypoints': tfds_features.Tensor(shape=(None, 3), dtype=tf.float32),%0A          'voxel': tfg_features.VoxelGrid(shape=(128, 128, 128)),%0A          'pose': tfg_features.Pose(),%0A          'camera': tfg_features.Camera(),%0A          'category': tfds_features.ClassLabel(shape=(), dtype=tf.int64, num_classes=9),%0A          'bbox': tfds_features.BBoxFeature(),%0A          'truncated': tf.bool,%0A          'occluded': tf.bool,%0A          'slightly_occluded': tf.bool
 %0A   
@@ -3005,28 +3005,30 @@
 http
-s
 ://
-dataset-homepage
+pix3d.csail.mit.edu
 /',%0A
@@ -3305,16 +3305,17 @@
 ct URLs%0A
+%0A
     retu

add Pix3D DatasetInfo Spec