data_loaders/generate_tfr/lsun.py

# 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)