Code for Diode Depth dataset for TFDS

diode_depth.py

import os

import numpy as np
import tensorflow as tf
import tensorflow_datasets.public_api as tfds

class DiodeDepth(tfds.core.GeneratorBasedBuilder):
    """Short description of my dataset."""

    VERSION = tfds.core.Version('0.1.0')

    def _info(self):
        return tfds.core.DatasetInfo(
            builder=self,
            # This is the description that will appear on the datasets page.
            description=("This is the dataset for DIODE(Dense Indoor/Outdoor DEpth). It contains images as .png, depth information per pixel as .npy, and a depth mask as .npy. The images are kept at their original dimensions(all are 768x1024)."),
            # tfds.features.FeatureConnectors
            features=tfds.features.FeaturesDict({
                "image": tfds.features.Image(encoding_format="png"),
                "depth": tfds.features.Image(shape=(768,1024,1)),
                "depth_mask": tfds.features.Image(shape=(768,1024,1)),
            }),
            # 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=("image", "depth"),
            # Homepage of the dataset for documentation
            homepage="https://diode-dataset.org/",
            # Bibtex citation for the dataset
            citation=r"""@article{diode_dataset,
                        title={{DIODE}: {A} {D}ense {I}ndoor and {O}utdoor {DE}pth {D}ataset},
                        author={Igor Vasiljevic and Nick Kolkin and Shanyi Zhang and Ruotian Luo and 
                                Haochen Wang and Falcon Z. Dai and Andrea F. Daniele and Mohammadreza Mostajabi and 
                                Steven Basart and Matthew R. Walter and Gregory Shakhnarovich},
                        year = {2019}
                        journal={CoRR},
                        volume={abs/1908.00463},
                        year = {2019},
                        url={http://arxiv.org/abs/1908.00463}
                        }""",
        )

    def _split_generators(self, dl_manager):
        # Downloads the data and defines the splits
        # dl_manager is a tfds.download.DownloadManager that can be used to
        # download and extract URLs
        dl_paths = dl_manager.download_and_extract({
            'train': 'http://diode-dataset.s3.amazonaws.com/train.tar.gz',
            'validation': 'http://diode-dataset.s3.amazonaws.com/val.tar.gz'
        })

        # Specify the splits
        return [
            tfds.core.SplitGenerator(
                name=tfds.Split.TRAIN,
                gen_kwargs={
                    "data_directory": dl_paths['train'],
                    "indoors_outdoor_all": "all"
                },
            ),
            tfds.core.SplitGenerator(
                name=tfds.Split.VALIDATION,
                gen_kwargs={
                    "data_directory": dl_paths['validation'],
                    "indoors_outdoor_all": "all"
                },
            ),
        ]

    def _generate_examples(self, data_directory, indoors_outdoor_all):
        image_paths = []
        if indoors_outdoor_all == "all":
            image_paths = tf.io.gfile.glob(os.path.join(data_directory,"*/*/*/*.png"))
        elif indoors_outdoor_all == "indoors":
            image_paths = tf.io.gfile.glob(os.path.join(data_directory,"indoors/*/*/*.png"))
        elif indoors_outdoor_all == "outdoor":
            image_paths = tf.io.gfile.glob(os.path.join(data_directory,"outdoor/*/*/*.png"))
        else:
            raise ValueError(f"{indoors_outdoor_all} is not a valid choice for `indoors_outdoor_all`. Choose `indoors`, `outdoor` or `all`.")
        
        for image_path in image_paths:
            depth_path = image_path.replace(".png","_depth.npy")
            depth_mask_path = image_path.replace(".png","_depth_mask.npy")
            yield image_path.replace(".png",""), {
                "image": image_path,
                "depth": np.load(depth_path),
                "depth_mask": np.load(depth_mask_path)
            }