Spaces:
Running
Running
| # Copyright 2017 The TensorFlow Authors 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. | |
| # ============================================================================== | |
| """Provides dataset dictionaries as used in our network models.""" | |
| from __future__ import absolute_import | |
| from __future__ import division | |
| from __future__ import print_function | |
| import os | |
| import tensorflow as tf | |
| import tensorflow.contrib.slim as slim | |
| from tensorflow.contrib.slim.python.slim.data import dataset | |
| from tensorflow.contrib.slim.python.slim.data import dataset_data_provider | |
| from tensorflow.contrib.slim.python.slim.data import tfexample_decoder | |
| _ITEMS_TO_DESCRIPTIONS = { | |
| 'image': 'Images', | |
| 'mask': 'Masks', | |
| 'vox': 'Voxels' | |
| } | |
| def _get_split(file_pattern, num_samples, num_views, image_size, vox_size): | |
| """Get dataset.Dataset for the given dataset file pattern and properties.""" | |
| # A dictionary from TF-Example keys to tf.FixedLenFeature instance. | |
| keys_to_features = { | |
| 'image': tf.FixedLenFeature( | |
| shape=[num_views, image_size, image_size, 3], | |
| dtype=tf.float32, default_value=None), | |
| 'mask': tf.FixedLenFeature( | |
| shape=[num_views, image_size, image_size, 1], | |
| dtype=tf.float32, default_value=None), | |
| 'vox': tf.FixedLenFeature( | |
| shape=[vox_size, vox_size, vox_size, 1], | |
| dtype=tf.float32, default_value=None), | |
| } | |
| items_to_handler = { | |
| 'image': tfexample_decoder.Tensor( | |
| 'image', shape=[num_views, image_size, image_size, 3]), | |
| 'mask': tfexample_decoder.Tensor( | |
| 'mask', shape=[num_views, image_size, image_size, 1]), | |
| 'vox': tfexample_decoder.Tensor( | |
| 'vox', shape=[vox_size, vox_size, vox_size, 1]) | |
| } | |
| decoder = tfexample_decoder.TFExampleDecoder( | |
| keys_to_features, items_to_handler) | |
| return dataset.Dataset( | |
| data_sources=file_pattern, | |
| reader=tf.TFRecordReader, | |
| decoder=decoder, | |
| num_samples=num_samples, | |
| items_to_descriptions=_ITEMS_TO_DESCRIPTIONS) | |
| def get(dataset_dir, | |
| dataset_name, | |
| split_name, | |
| shuffle=True, | |
| num_readers=1, | |
| common_queue_capacity=64, | |
| common_queue_min=50): | |
| """Provides input data for a specified dataset and split.""" | |
| dataset_to_kwargs = { | |
| 'shapenet_chair': { | |
| 'file_pattern': '03001627_%s.tfrecords' % split_name, | |
| 'num_views': 24, | |
| 'image_size': 64, | |
| 'vox_size': 32, | |
| }, 'shapenet_all': { | |
| 'file_pattern': '*_%s.tfrecords' % split_name, | |
| 'num_views': 24, | |
| 'image_size': 64, | |
| 'vox_size': 32, | |
| }, | |
| } | |
| split_sizes = { | |
| 'shapenet_chair': { | |
| 'train': 4744, | |
| 'val': 678, | |
| 'test': 1356, | |
| }, | |
| 'shapenet_all': { | |
| 'train': 30643, | |
| 'val': 4378, | |
| 'test': 8762, | |
| } | |
| } | |
| kwargs = dataset_to_kwargs[dataset_name] | |
| kwargs['file_pattern'] = os.path.join(dataset_dir, kwargs['file_pattern']) | |
| kwargs['num_samples'] = split_sizes[dataset_name][split_name] | |
| dataset_split = _get_split(**kwargs) | |
| data_provider = dataset_data_provider.DatasetDataProvider( | |
| dataset_split, | |
| num_readers=num_readers, | |
| common_queue_capacity=common_queue_capacity, | |
| common_queue_min=common_queue_min, | |
| shuffle=shuffle) | |
| inputs = { | |
| 'num_samples': dataset_split.num_samples, | |
| } | |
| [image, mask, vox] = data_provider.get(['image', 'mask', 'vox']) | |
| inputs['image'] = image | |
| inputs['mask'] = mask | |
| inputs['voxel'] = vox | |
| return inputs | |