text: # 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.
  # ==============================================================================
  
  """"
  LSUN dataset
  
  # Get image files
  Download the LSUN dataset as follows:
  git clone https://github.com/fyu/lsun.git
  cd lsun
  python2.7 download.py -c [CATEGORY]
  Unzip the downloaded .zip files and execute:
  python2.7 data.py export [IMAGE_DB_PATH] --out_dir [LSUN_FOLDER] --flat
  
  # Get tfr file from images
  Use this script to generate the tfr file.
  python lsun.py --res [RES] --category [CATEGORY] --lsun_dir [LSUN_FOLDER]
      --tfrecord_dir [OUTPUT_FOLDER] --write [--realnvp]
  Without realnvp flag you get 256x256 centre cropped area downsampled images,
  with flag you get 96x96 images with realnvp preprocessing.
  """
  
  from __future__ import print_function
  
  import os
  import os.path
  
  import numpy
  import skimage.transform
  from PIL import Image
  import tensorflow as tf
  import numpy as np
  from tqdm import tqdm
  
  from typing import Iterable
  
  
  def _int64_feature(value):
      if not isinstance(value, Iterable):
          value = [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 centre_crop(img):
      h, w = img.shape[:2]
      crop = min(h, w)
      return img[(h - crop) // 2: (h + crop) // 2, (w - crop) // 2: (w + crop) // 2]
  
  
  def dump(fn_root, tfrecord_dir, max_res, expected_images, shards, write, realnvp=False):
      """Main converter function."""
      resolution_log2 = int(np.log2(max_res))
      tfr_prefix = os.path.join(tfrecord_dir, os.path.basename(tfrecord_dir))
  
      print("Checking in", 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)
      print("Found", num_examples)
      assert num_examples == expected_images
  
      tfr_opt = tf.python_io.TFRecordOptions(
          tf.python_io.TFRecordCompressionType.NONE)
      p_shard = np.array_split(np.random.permutation(expected_images), shards)
      img_to_shard = np.zeros(expected_images, dtype=np.int)
      writers = []
      for shard in tqdm(range(shards)):
          img_to_shard[p_shard[shard]] = shard
          tfr_file = tfr_prefix + \
              '-r%02d-s-%04d-of-%04d.tfrecords' % (
                  resolution_log2, shard, shards)
          writers.append(tf.python_io.TFRecordWriter(tfr_file, tfr_opt))
  
      # print(np.unique(img_to_shard, return_counts=True))
      counts = np.unique(img_to_shard, return_counts=True)[1]
      assert len(counts) == shards
      print("Smallest and largest shards have size",
            np.min(counts), np.max(counts))
  
      for example_idx, img_fn in enumerate(tqdm(img_fn_list)):
          shard = img_to_shard[example_idx]
          img = numpy.array(Image.open(os.path.join(fn_root, img_fn)))
          rows = img.shape[0]
          cols = img.shape[1]
          if realnvp:
              downscale = min(rows / 96., cols / 96.)
              img = skimage.transform.pyramid_reduce(img, downscale)
              img *= 255.
              img = img.astype("uint8")
          else:
              img = centre_crop(img)
              img = Image.fromarray(img, 'RGB')
              img = img.resize((max_res, max_res), Image.ANTIALIAS)
              img = np.asarray(img)
          rows = img.shape[0]
          cols = img.shape[1]
          depth = img.shape[2]
          shape = (rows, cols, depth)
          img = img.tostring()
          example = tf.train.Example(
              features=tf.train.Features(
                  feature={
                      "shape": _int64_feature(shape),
                      "data": _bytes_feature(img),
                      "label": _int64_feature(0)
                  }
              )
          )
          if write:
              writers[shard].write(example.SerializeToString())
  
      print('%-40s\r' % 'Flushing data...', end='', flush=True)
      for writer in writers:
          writer.close()
  
      print('%-40s\r' % '', end='', flush=True)
      print('Added %d images.' % num_examples)
  
  
  if __name__ == "__main__":
      import argparse
  
      parser = argparse.ArgumentParser()
      parser.add_argument("--category", type=str, help="LSUN category")
      parser.add_argument("--realnvp", action='store_true',
                          help="Use this flag to do realnvp preprocessing instead of our centre-crops")
      parser.add_argument("--res", type=int, default=256, help="Image size")
      parser.add_argument("--lsun_dir", type=str,
                          required=True, help="place of lsun dir")
      parser.add_argument("--tfrecord_dir", type=str,
                          required=True, help='place to dump')
      parser.add_argument("--write", action='store_true',
                          help="Whether to write")
      hps = parser.parse_args()
  
      # LSUN
      # CATEGORIES = ["bedroom", "bridge", "church_outdoor", "classroom", "conference_room",
      #               "dining_room", "kitchen", "living"]
      base_tfr = hps.tfrecord_dir
      res = hps.res
      for realnvp in [False, True]:
          for category in ["tower", "church_outdoor", "bedroom"]:
              hps.realnvp = realnvp
              hps.category = category
              if realnvp:
                  hps.tfrecord_dir = "%s_%s/%s" % (base_tfr,
                                                   "realnvp", hps.category)
              else:
                  hps.tfrecord_dir = "%s/%s" % (base_tfr, hps.category)
              print(hps.realnvp, hps.category, hps.lsun_dir, hps.tfrecord_dir)
              imgs = {
                  'bedroom': 3033042,
                  'bridge': 818687,
                  'church_outdoor': 126227,
                  'classroom': 168103,
                  'conference_room': 229069,
                  'dining_room': 657571,
                  'kitchen': 2212277,
                  'living_room': 1315802,
                  'restaurant': 626331,
                  'tower': 708264
              }
  
              _NUM_IMAGES = {
                  'train': imgs[hps.category],
                  'validation': 300,
              }
  
              _NUM_SHARDS = {
                  'train': 2560,
                  'validation': 1,
              }
  
              _FILE = {
                  'train': os.path.join(hps.lsun_dir, '%s_train' % hps.category),
                  'validation': os.path.join(hps.lsun_dir, '%s_val' % hps.category)
  
              }
  
              if hps.realnvp:
                  res = 96
              else:
                  res = hps.res
  
              for split in ['validation', 'train']:
                  fn_root = _FILE[split]
                  tfrecord_dir = os.path.join(hps.tfrecord_dir, split)
                  total_imgs = _NUM_IMAGES[split]
                  shards = _NUM_SHARDS[split]
                  if not os.path.exists(tfrecord_dir):
                      os.mkdir(tfrecord_dir)
                  dump(fn_root, tfrecord_dir, res, total_imgs,
                       shards, hps.write, hps.realnvp)
  