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"""LSUN dataset formatting. | |
| Download and format the LSUN dataset as follow: | |
| git clone https://github.com/fyu/lsun.git | |
| cd lsun | |
| python2.7 download.py -c [CATEGORY] | |
| Then unzip the downloaded .zip files before executing: | |
| python2.7 data.py export [IMAGE_DB_PATH] --out_dir [LSUN_FOLDER] --flat | |
| Then use the script as follow: | |
| python lsun_formatting.py \ | |
| --file_out [OUTPUT_FILE_PATH_PREFIX] \ | |
| --fn_root [LSUN_FOLDER] | |
| """ | |
| from __future__ import print_function | |
| import os | |
| import os.path | |
| import numpy | |
| import skimage.transform | |
| from PIL import Image | |
| 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.") | |
| 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.""" | |
| fn_root = FLAGS.fn_root | |
| img_fn_list = os.listdir(fn_root) | |
| img_fn_list = [img_fn for img_fn in img_fn_list | |
| if img_fn.endswith('.webp')] | |
| num_examples = len(img_fn_list) | |
| n_examples_per_file = 10000 | |
| for example_idx, img_fn in enumerate(img_fn_list): | |
| if example_idx % n_examples_per_file == 0: | |
| file_out = "%s_%05d.tfrecords" | |
| file_out = file_out % (FLAGS.file_out, | |
| example_idx // n_examples_per_file) | |
| print("Writing on:", file_out) | |
| writer = tf.python_io.TFRecordWriter(file_out) | |
| if example_idx % 1000 == 0: | |
| print(example_idx, "/", num_examples) | |
| image_raw = numpy.array(Image.open(os.path.join(fn_root, img_fn))) | |
| rows = image_raw.shape[0] | |
| cols = image_raw.shape[1] | |
| depth = image_raw.shape[2] | |
| downscale = min(rows / 96., cols / 96.) | |
| image_raw = skimage.transform.pyramid_reduce(image_raw, downscale) | |
| image_raw *= 255. | |
| image_raw = image_raw.astype("uint8") | |
| 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()) | |
| if example_idx % n_examples_per_file == (n_examples_per_file - 1): | |
| writer.close() | |
| writer.close() | |
| if __name__ == "__main__": | |
| main() | |