Spaces:
Running
Running
| # Copyright 2016 Google Inc. 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. | |
| # ============================================================================== | |
| r"""CelebA dataset formating. | |
| Download img_align_celeba.zip from | |
| http://mmlab.ie.cuhk.edu.hk/projects/CelebA.html under the | |
| link "Align&Cropped Images" in the "Img" directory and list_eval_partition.txt | |
| under the link "Train/Val/Test Partitions" in the "Eval" directory. Then do: | |
| unzip img_align_celeba.zip | |
| Use the script as follow: | |
| python celeba_formatting.py \ | |
| --partition_fn [PARTITION_FILE_PATH] \ | |
| --file_out [OUTPUT_FILE_PATH_PREFIX] \ | |
| --fn_root [CELEBA_FOLDER] \ | |
| --set [SUBSET_INDEX] | |
| """ | |
| from __future__ import print_function | |
| import os | |
| import os.path | |
| import scipy.io | |
| import scipy.io.wavfile | |
| import scipy.ndimage | |
| import tensorflow as tf | |
| tf.flags.DEFINE_string("file_out", "", | |
| "Filename of the output .tfrecords file.") | |
| tf.flags.DEFINE_string("fn_root", "", "Name of root file path.") | |
| tf.flags.DEFINE_string("partition_fn", "", "Partition file path.") | |
| tf.flags.DEFINE_string("set", "", "Name of subset.") | |
| FLAGS = tf.flags.FLAGS | |
| def _int64_feature(value): | |
| return tf.train.Feature(int64_list=tf.train.Int64List(value=[value])) | |
| def _bytes_feature(value): | |
| return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value])) | |
| def main(): | |
| """Main converter function.""" | |
| # Celeb A | |
| with open(FLAGS.partition_fn, "r") as infile: | |
| img_fn_list = infile.readlines() | |
| img_fn_list = [elem.strip().split() for elem in img_fn_list] | |
| img_fn_list = [elem[0] for elem in img_fn_list if elem[1] == FLAGS.set] | |
| fn_root = FLAGS.fn_root | |
| num_examples = len(img_fn_list) | |
| file_out = "%s.tfrecords" % FLAGS.file_out | |
| writer = tf.python_io.TFRecordWriter(file_out) | |
| for example_idx, img_fn in enumerate(img_fn_list): | |
| if example_idx % 1000 == 0: | |
| print(example_idx, "/", num_examples) | |
| image_raw = scipy.ndimage.imread(os.path.join(fn_root, img_fn)) | |
| rows = image_raw.shape[0] | |
| cols = image_raw.shape[1] | |
| depth = image_raw.shape[2] | |
| image_raw = image_raw.tostring() | |
| example = tf.train.Example( | |
| features=tf.train.Features( | |
| feature={ | |
| "height": _int64_feature(rows), | |
| "width": _int64_feature(cols), | |
| "depth": _int64_feature(depth), | |
| "image_raw": _bytes_feature(image_raw) | |
| } | |
| ) | |
| ) | |
| writer.write(example.SerializeToString()) | |
| writer.close() | |
| if __name__ == "__main__": | |
| main() | |