Spaces:
Running
Running
| # coding=utf-8 | |
| # Copyright 2023 The Google Research Authors. | |
| # | |
| # 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. | |
| """Input pipeline for TFDS datasets.""" | |
| import functools | |
| import os | |
| from typing import Dict, List, Tuple | |
| from clu import deterministic_data | |
| from clu import preprocess_spec | |
| import jax | |
| import jax.numpy as jnp | |
| import ml_collections | |
| import sunds | |
| import tensorflow as tf | |
| import tensorflow_datasets as tfds | |
| from invariant_slot_attention.lib import preprocessing | |
| Array = jnp.ndarray | |
| PRNGKey = Array | |
| PATH_CLEVR_WITH_MASKS = "gs://multi-object-datasets/clevr_with_masks/clevr_with_masks_train.tfrecords" | |
| FEATURES_CLEVR_WITH_MASKS = { | |
| "image": tf.io.FixedLenFeature([240, 320, 3], tf.string), | |
| "mask": tf.io.FixedLenFeature([11, 240, 320, 1], tf.string), | |
| "x": tf.io.FixedLenFeature([11], tf.float32), | |
| "y": tf.io.FixedLenFeature([11], tf.float32), | |
| "z": tf.io.FixedLenFeature([11], tf.float32), | |
| "pixel_coords": tf.io.FixedLenFeature([11, 3], tf.float32), | |
| "rotation": tf.io.FixedLenFeature([11], tf.float32), | |
| "size": tf.io.FixedLenFeature([11], tf.string), | |
| "material": tf.io.FixedLenFeature([11], tf.string), | |
| "shape": tf.io.FixedLenFeature([11], tf.string), | |
| "color": tf.io.FixedLenFeature([11], tf.string), | |
| "visibility": tf.io.FixedLenFeature([11], tf.float32), | |
| } | |
| PATH_TETROMINOES = "gs://multi-object-datasets/tetrominoes/tetrominoes_train.tfrecords" | |
| FEATURES_TETROMINOES = { | |
| "image": tf.io.FixedLenFeature([35, 35, 3], tf.string), | |
| "mask": tf.io.FixedLenFeature([4, 35, 35, 1], tf.string), | |
| "x": tf.io.FixedLenFeature([4], tf.float32), | |
| "y": tf.io.FixedLenFeature([4], tf.float32), | |
| "shape": tf.io.FixedLenFeature([4], tf.float32), | |
| "color": tf.io.FixedLenFeature([4, 3], tf.float32), | |
| "visibility": tf.io.FixedLenFeature([4], tf.float32), | |
| } | |
| PATH_OBJECTS_ROOM = "gs://multi-object-datasets/objects_room/objects_room_train.tfrecords" | |
| FEATURES_OBJECTS_ROOM = { | |
| "image": tf.io.FixedLenFeature([64, 64, 3], tf.string), | |
| "mask": tf.io.FixedLenFeature([7, 64, 64, 1], tf.string), | |
| } | |
| PATH_WAYMO_OPEN = "datasets/waymo_v_1_4_0_images/tfrecords" | |
| FEATURES_WAYMO_OPEN = { | |
| "image": tf.io.FixedLenFeature([128, 192, 3], tf.string), | |
| "segmentations": tf.io.FixedLenFeature([128, 192], tf.string), | |
| "depth": tf.io.FixedLenFeature([128, 192], tf.float32), | |
| "num_objects": tf.io.FixedLenFeature([1], tf.int64), | |
| "has_mask": tf.io.FixedLenFeature([1], tf.int64), | |
| "camera": tf.io.FixedLenFeature([1], tf.int64), | |
| } | |
| def _decode_tetrominoes(example_proto): | |
| single_example = tf.io.parse_single_example( | |
| example_proto, FEATURES_TETROMINOES) | |
| for k in ["mask", "image"]: | |
| single_example[k] = tf.squeeze( | |
| tf.io.decode_raw(single_example[k], tf.uint8), axis=-1) | |
| return single_example | |
| def _decode_objects_room(example_proto): | |
| single_example = tf.io.parse_single_example( | |
| example_proto, FEATURES_OBJECTS_ROOM) | |
| for k in ["mask", "image"]: | |
| single_example[k] = tf.squeeze( | |
| tf.io.decode_raw(single_example[k], tf.uint8), axis=-1) | |
| return single_example | |
| def _decode_clevr_with_masks(example_proto): | |
| single_example = tf.io.parse_single_example( | |
| example_proto, FEATURES_CLEVR_WITH_MASKS) | |
| for k in ["mask", "image", "color", "material", "shape", "size"]: | |
| single_example[k] = tf.squeeze( | |
| tf.io.decode_raw(single_example[k], tf.uint8), axis=-1) | |
| return single_example | |
| def _decode_waymo_open(example_proto): | |
| """Unserializes a serialized tf.train.Example sample.""" | |
| single_example = tf.io.parse_single_example( | |
| example_proto, FEATURES_WAYMO_OPEN) | |
| for k in ["image", "segmentations"]: | |
| single_example[k] = tf.squeeze( | |
| tf.io.decode_raw(single_example[k], tf.uint8), axis=-1) | |
| single_example["segmentations"] = tf.expand_dims( | |
| single_example["segmentations"], axis=-1) | |
| single_example["depth"] = tf.expand_dims( | |
| single_example["depth"], axis=-1) | |
| return single_example | |
| def _preprocess_minimal(example): | |
| return { | |
| "image": example["image"], | |
| "segmentations": tf.cast(tf.argmax(example["mask"], axis=0), tf.uint8), | |
| } | |
| def _sunds_create_task(): | |
| """Create a sunds task to return images and instance segmentation.""" | |
| return sunds.tasks.Nerf( | |
| yield_mode=sunds.tasks.YieldMode.IMAGE, | |
| additional_camera_specs={ | |
| "depth_image": False, # Not available in the dataset. | |
| "category_image": False, # Not available in the dataset. | |
| "instance_image": True, | |
| "extrinsics": True, | |
| }, | |
| additional_frame_specs={"pose": True}, | |
| add_name=True | |
| ) | |
| def preprocess_example(features, | |
| preprocess_strs): | |
| """Processes a single data example. | |
| Args: | |
| features: A dictionary containing the tensors of a single data example. | |
| preprocess_strs: List of strings, describing one preprocessing operation | |
| each, in clu.preprocess_spec format. | |
| Returns: | |
| Dictionary containing the preprocessed tensors of a single data example. | |
| """ | |
| all_ops = preprocessing.all_ops() | |
| preprocess_fn = preprocess_spec.parse("|".join(preprocess_strs), all_ops) | |
| return preprocess_fn(features) # pytype: disable=bad-return-type # allow-recursive-types | |
| def get_batch_dims(global_batch_size): | |
| """Gets the first two axis sizes for data batches. | |
| Args: | |
| global_batch_size: Integer, the global batch size (across all devices). | |
| Returns: | |
| List of batch dimensions | |
| Raises: | |
| ValueError if the requested dimensions don't make sense with the | |
| number of devices. | |
| """ | |
| num_local_devices = jax.local_device_count() | |
| if global_batch_size % jax.host_count() != 0: | |
| raise ValueError(f"Global batch size {global_batch_size} not evenly " | |
| f"divisble with {jax.host_count()}.") | |
| per_host_batch_size = global_batch_size // jax.host_count() | |
| if per_host_batch_size % num_local_devices != 0: | |
| raise ValueError(f"Global batch size {global_batch_size} not evenly " | |
| f"divisible with {jax.host_count()} hosts with a per host " | |
| f"batch size of {per_host_batch_size} and " | |
| f"{num_local_devices} local devices. ") | |
| return [num_local_devices, per_host_batch_size // num_local_devices] | |
| def create_datasets( | |
| config, | |
| data_rng): | |
| """Create datasets for training and evaluation. | |
| For the same data_rng and config this will return the same datasets. The | |
| datasets only contain stateless operations. | |
| Args: | |
| config: Configuration to use. | |
| data_rng: JAX PRNGKey for dataset pipeline. | |
| Returns: | |
| A tuple with the training dataset and the evaluation dataset. | |
| """ | |
| if config.data.dataset_name == "tetrominoes": | |
| ds = tf.data.TFRecordDataset( | |
| PATH_TETROMINOES, | |
| compression_type="GZIP", buffer_size=2*(2**20)) | |
| ds = ds.map(_decode_tetrominoes, | |
| num_parallel_calls=tf.data.experimental.AUTOTUNE) | |
| ds = ds.map(_preprocess_minimal, | |
| num_parallel_calls=tf.data.experimental.AUTOTUNE) | |
| class TetrominoesBuilder: | |
| """Builder for tentrominoes dataset.""" | |
| def as_dataset(self, split, *unused_args, ds=ds, **unused_kwargs): | |
| """Simple function to conform to the builder api.""" | |
| if split == "train": | |
| # We use 512 training examples. | |
| ds = ds.skip(100) | |
| ds = ds.take(512) | |
| return tf.data.experimental.assert_cardinality(512)(ds) | |
| elif split == "validation": | |
| # 100 validation examples. | |
| ds = ds.take(100) | |
| return tf.data.experimental.assert_cardinality(100)(ds) | |
| else: | |
| raise ValueError("Invalid split.") | |
| dataset_builder = TetrominoesBuilder() | |
| elif config.data.dataset_name == "objects_room": | |
| ds = tf.data.TFRecordDataset( | |
| PATH_OBJECTS_ROOM, | |
| compression_type="GZIP", buffer_size=2*(2**20)) | |
| ds = ds.map(_decode_objects_room, | |
| num_parallel_calls=tf.data.experimental.AUTOTUNE) | |
| ds = ds.map(_preprocess_minimal, | |
| num_parallel_calls=tf.data.experimental.AUTOTUNE) | |
| class ObjectsRoomBuilder: | |
| """Builder for objects room dataset.""" | |
| def as_dataset(self, split, *unused_args, ds=ds, **unused_kwargs): | |
| """Simple function to conform to the builder api.""" | |
| if split == "train": | |
| # 1M - 100 training examples. | |
| ds = ds.skip(100) | |
| return tf.data.experimental.assert_cardinality(999900)(ds) | |
| elif split == "validation": | |
| # 100 validation examples. | |
| ds = ds.take(100) | |
| return tf.data.experimental.assert_cardinality(100)(ds) | |
| else: | |
| raise ValueError("Invalid split.") | |
| dataset_builder = ObjectsRoomBuilder() | |
| elif config.data.dataset_name == "clevr_with_masks": | |
| ds = tf.data.TFRecordDataset( | |
| PATH_CLEVR_WITH_MASKS, | |
| compression_type="GZIP", buffer_size=2*(2**20)) | |
| ds = ds.map(_decode_clevr_with_masks, | |
| num_parallel_calls=tf.data.experimental.AUTOTUNE) | |
| ds = ds.map(_preprocess_minimal, | |
| num_parallel_calls=tf.data.experimental.AUTOTUNE) | |
| class CLEVRWithMasksBuilder: | |
| def as_dataset(self, split, *unused_args, ds=ds, **unused_kwargs): | |
| if split == "train": | |
| ds = ds.skip(100) | |
| return tf.data.experimental.assert_cardinality(99900)(ds) | |
| elif split == "validation": | |
| ds = ds.take(100) | |
| return tf.data.experimental.assert_cardinality(100)(ds) | |
| else: | |
| raise ValueError("Invalid split.") | |
| dataset_builder = CLEVRWithMasksBuilder() | |
| elif config.data.dataset_name == "waymo_open": | |
| train_path = os.path.join( | |
| PATH_WAYMO_OPEN, "training/camera_1/*tfrecords*") | |
| eval_path = os.path.join( | |
| PATH_WAYMO_OPEN, "validation/camera_1/*tfrecords*") | |
| train_files = tf.data.Dataset.list_files(train_path) | |
| eval_files = tf.data.Dataset.list_files(eval_path) | |
| train_data_reader = functools.partial( | |
| tf.data.TFRecordDataset, | |
| compression_type="ZLIB", buffer_size=2*(2**20)) | |
| eval_data_reader = functools.partial( | |
| tf.data.TFRecordDataset, | |
| compression_type="ZLIB", buffer_size=2*(2**20)) | |
| train_dataset = train_files.interleave( | |
| train_data_reader, num_parallel_calls=tf.data.experimental.AUTOTUNE) | |
| eval_dataset = eval_files.interleave( | |
| eval_data_reader, num_parallel_calls=tf.data.experimental.AUTOTUNE) | |
| train_dataset = train_dataset.map( | |
| _decode_waymo_open, num_parallel_calls=tf.data.experimental.AUTOTUNE) | |
| eval_dataset = eval_dataset.map( | |
| _decode_waymo_open, num_parallel_calls=tf.data.experimental.AUTOTUNE) | |
| # We need to set the dataset cardinality. We assume we have | |
| # the full dataset. | |
| train_dataset = train_dataset.apply( | |
| tf.data.experimental.assert_cardinality(158081)) | |
| class WaymoOpenBuilder: | |
| def as_dataset(self, split, *unused_args, **unused_kwargs): | |
| if split == "train": | |
| return train_dataset | |
| elif split == "validation": | |
| return eval_dataset | |
| else: | |
| raise ValueError("Invalid split.") | |
| dataset_builder = WaymoOpenBuilder() | |
| elif config.data.dataset_name == "multishapenet_easy": | |
| dataset_builder = sunds.builder( | |
| name=config.get("tfds_name", "msn_easy"), | |
| data_dir=config.get( | |
| "data_dir", "gs://kubric-public/tfds"), | |
| try_gcs=True) | |
| dataset_builder.as_dataset = functools.partial( | |
| dataset_builder.as_dataset, task=_sunds_create_task()) | |
| elif config.data.dataset_name == "tfds": | |
| dataset_builder = tfds.builder( | |
| config.data.tfds_name, data_dir=config.data.data_dir) | |
| else: | |
| raise ValueError("Please specify a valid dataset name.") | |
| batch_dims = get_batch_dims(config.batch_size) | |
| train_preprocess_fn = functools.partial( | |
| preprocess_example, preprocess_strs=config.preproc_train) | |
| eval_preprocess_fn = functools.partial( | |
| preprocess_example, preprocess_strs=config.preproc_eval) | |
| train_split_name = config.get("train_split", "train") | |
| eval_split_name = config.get("validation_split", "validation") | |
| train_ds = deterministic_data.create_dataset( | |
| dataset_builder, | |
| split=train_split_name, | |
| rng=data_rng, | |
| preprocess_fn=train_preprocess_fn, | |
| cache=False, | |
| shuffle_buffer_size=config.data.shuffle_buffer_size, | |
| batch_dims=batch_dims, | |
| num_epochs=None, | |
| shuffle=True) | |
| if config.data.dataset_name == "waymo_open": | |
| # We filter Waymo Open for empty segmentation masks. | |
| def filter_fn(features): | |
| unique_instances = tf.unique( | |
| tf.reshape(features[preprocessing.SEGMENTATIONS], (-1,)))[0] | |
| n_instances = tf.size(unique_instances, tf.int32) | |
| # n_instances == 1 means we only have the background. | |
| return 2 <= n_instances | |
| else: | |
| filter_fn = None | |
| eval_ds = deterministic_data.create_dataset( | |
| dataset_builder, | |
| split=eval_split_name, | |
| rng=None, | |
| preprocess_fn=eval_preprocess_fn, | |
| filter_fn=filter_fn, | |
| cache=False, | |
| batch_dims=batch_dims, | |
| num_epochs=1, | |
| shuffle=False, | |
| pad_up_to_batches=None) | |
| if config.data.dataset_name == "waymo_open": | |
| # We filter Waymo Open for empty segmentation masks after preprocessing. | |
| # For the full dataset, we know how many we will end up with. | |
| eval_batch_size = batch_dims[0] * batch_dims[1] | |
| # We don't pad the last batch => floor. | |
| eval_num_batches = int( | |
| jnp.floor(1872 / eval_batch_size / jax.host_count())) | |
| eval_ds = eval_ds.apply( | |
| tf.data.experimental.assert_cardinality( | |
| eval_num_batches)) | |
| return train_ds, eval_ds | |