dataset_tool.py

# Copyright (c) 2019, NVIDIA Corporation. All rights reserved.
#
# This work is made available under the Nvidia Source Code License-NC.
# To view a copy of this license, visit
# https://nvlabs.github.io/stylegan2/license.html
"""Tool for creating multi-resolution TFRecords datasets."""

# pylint: disable=too-many-lines
import os
import sys
import glob
import argparse
import threading
import six.moves.queue as Queue  # pylint: disable=import-error
import traceback
import numpy as np
import tensorflow as tf
import PIL.Image
import dnnlib.tflib as tflib

from training import dataset

#----------------------------------------------------------------------------


def error(msg):
    print('Error: ' + msg)
    exit(1)


#----------------------------------------------------------------------------


class TFRecordExporter:
    def __init__(self,
                 tfrecord_dir,
                 expected_images,
                 print_progress=True,
                 progress_interval=10):
        self.tfrecord_dir = tfrecord_dir
        self.tfr_prefix = os.path.join(self.tfrecord_dir,
                                       os.path.basename(self.tfrecord_dir))
        self.expected_images = expected_images
        self.cur_images = 0
        self.shape = None
        self.resolution_log2 = None
        self.tfr_writers = []
        self.print_progress = print_progress
        self.progress_interval = progress_interval

        if self.print_progress:
            print('Creating dataset "%s"' % tfrecord_dir)
        if not os.path.isdir(self.tfrecord_dir):
            os.makedirs(self.tfrecord_dir)
        assert os.path.isdir(self.tfrecord_dir)

    def close(self):
        if self.print_progress:
            print('%-40s\r' % 'Flushing data...', end='', flush=True)
        for tfr_writer in self.tfr_writers:
            tfr_writer.close()
        self.tfr_writers = []
        if self.print_progress:
            print('%-40s\r' % '', end='', flush=True)
            print('Added %d images.' % self.cur_images)

    def choose_shuffled_order(
            self):  # Note: Images and labels must be added in shuffled order.
        order = np.arange(self.expected_images)
        np.random.RandomState(123).shuffle(order)
        return order

    def add_image(self, img):
        if self.print_progress and self.cur_images % self.progress_interval == 0:
            print('%d / %d\r' % (self.cur_images, self.expected_images),
                  end='',
                  flush=True)
        if self.shape is None:
            self.shape = img.shape
            self.resolution_log2 = int(np.log2(self.shape[1]))
            assert self.shape[0] in [1, 3]
            assert self.shape[1] == self.shape[2]
            assert self.shape[1] == 2**self.resolution_log2
            tfr_opt = tf.python_io.TFRecordOptions(
                tf.python_io.TFRecordCompressionType.NONE)
            for lod in range(self.resolution_log2 - 1):
                tfr_file = self.tfr_prefix + '-r%02d.tfrecords' % (
                    self.resolution_log2 - lod)
                self.tfr_writers.append(
                    tf.python_io.TFRecordWriter(tfr_file, tfr_opt))
        assert img.shape == self.shape
        for lod, tfr_writer in enumerate(self.tfr_writers):
            if lod:
                img = img.astype(np.float32)
                img = (img[:, 0::2, 0::2] + img[:, 0::2, 1::2] +
                       img[:, 1::2, 0::2] + img[:, 1::2, 1::2]) * 0.25
            quant = np.rint(img).clip(0, 255).astype(np.uint8)
            ex = tf.train.Example(features=tf.train.Features(
                feature={
                    'shape':
                    tf.train.Feature(int64_list=tf.train.Int64List(
                        value=quant.shape)),
                    'data':
                    tf.train.Feature(bytes_list=tf.train.BytesList(
                        value=[quant.tostring()]))
                }))
            tfr_writer.write(ex.SerializeToString())
        self.cur_images += 1

    def add_labels(self, labels):
        if self.print_progress:
            print('%-40s\r' % 'Saving labels...', end='', flush=True)
        assert labels.shape[0] == self.cur_images
        with open(self.tfr_prefix + '-rxx.labels', 'wb') as f:
            np.save(f, labels.astype(np.float32))

    def __enter__(self):
        return self

    def __exit__(self, *args):
        self.close()


#----------------------------------------------------------------------------


class ExceptionInfo(object):
    def __init__(self):
        self.value = sys.exc_info()[1]
        self.traceback = traceback.format_exc()


#----------------------------------------------------------------------------


class WorkerThread(threading.Thread):
    def __init__(self, task_queue):
        threading.Thread.__init__(self)
        self.task_queue = task_queue

    def run(self):
        while True:
            func, args, result_queue = self.task_queue.get()
            if func is None:
                break
            try:
                result = func(*args)
            except:
                result = ExceptionInfo()
            result_queue.put((result, args))


#----------------------------------------------------------------------------


class ThreadPool(object):
    def __init__(self, num_threads):
        assert num_threads >= 1
        self.task_queue = Queue.Queue()
        self.result_queues = dict()
        self.num_threads = num_threads
        for _idx in range(self.num_threads):
            thread = WorkerThread(self.task_queue)
            thread.daemon = True
            thread.start()

    def add_task(self, func, args=()):
        assert hasattr(func, '__call__')  # must be a function
        if func not in self.result_queues:
            self.result_queues[func] = Queue.Queue()
        self.task_queue.put((func, args, self.result_queues[func]))

    def get_result(self, func):  # returns (result, args)
        result, args = self.result_queues[func].get()
        if isinstance(result, ExceptionInfo):
            print('\n\nWorker thread caught an exception:\n' +
                  result.traceback)
            raise result.value
        return result, args

    def finish(self):
        for _idx in range(self.num_threads):
            self.task_queue.put((None, (), None))

    def __enter__(self):  # for 'with' statement
        return self

    def __exit__(self, *excinfo):
        self.finish()

    def process_items_concurrently(self,
                                   item_iterator,
                                   process_func=lambda x: x,
                                   pre_func=lambda x: x,
                                   post_func=lambda x: x,
                                   max_items_in_flight=None):
        if max_items_in_flight is None:
            max_items_in_flight = self.num_threads * 4
        assert max_items_in_flight >= 1
        results = []
        retire_idx = [0]

        def task_func(prepared, _idx):
            return process_func(prepared)

        def retire_result():
            processed, (_prepared, idx) = self.get_result(task_func)
            results[idx] = processed
            while retire_idx[0] < len(results) and results[
                    retire_idx[0]] is not None:
                yield post_func(results[retire_idx[0]])
                results[retire_idx[0]] = None
                retire_idx[0] += 1

        for idx, item in enumerate(item_iterator):
            prepared = pre_func(item)
            results.append(None)
            self.add_task(func=task_func, args=(prepared, idx))
            while retire_idx[0] < idx - max_items_in_flight + 2:
                for res in retire_result():
                    yield res
        while retire_idx[0] < len(results):
            for res in retire_result():
                yield res


#----------------------------------------------------------------------------


def display(tfrecord_dir):
    print('Loading dataset "%s"' % tfrecord_dir)
    tflib.init_tf({'gpu_options.allow_growth': True})
    dset = dataset.TFRecordDataset(tfrecord_dir,
                                   max_label_size='full',
                                   repeat=False,
                                   shuffle_mb=0)
    tflib.init_uninitialized_vars()
    import cv2  # pip install opencv-python

    idx = 0
    while True:
        try:
            images, labels = dset.get_minibatch_np(1)
        except tf.errors.OutOfRangeError:
            break
        if idx == 0:
            print('Displaying images')
            cv2.namedWindow('dataset_tool')
            print('Press SPACE or ENTER to advance, ESC to exit')
        print('\nidx = %-8d\nlabel = %s' % (idx, labels[0].tolist()))
        cv2.imshow('dataset_tool', images[0].transpose(
            1, 2, 0)[:, :, ::-1])  # CHW => HWC, RGB => BGR
        idx += 1
        if cv2.waitKey() == 27:
            break
    print('\nDisplayed %d images.' % idx)


#----------------------------------------------------------------------------


def extract(tfrecord_dir, output_dir):
    print('Loading dataset "%s"' % tfrecord_dir)
    tflib.init_tf({'gpu_options.allow_growth': True})
    dset = dataset.TFRecordDataset(tfrecord_dir,
                                   max_label_size=0,
                                   repeat=False,
                                   shuffle_mb=0)
    tflib.init_uninitialized_vars()

    print('Extracting images to "%s"' % output_dir)
    if not os.path.isdir(output_dir):
        os.makedirs(output_dir)
    idx = 0
    while True:
        if idx % 10 == 0:
            print('%d\r' % idx, end='', flush=True)
        try:
            images, _labels = dset.get_minibatch_np(1)
        except tf.errors.OutOfRangeError:
            break
        if images.shape[1] == 1:
            img = PIL.Image.fromarray(images[0][0], 'L')
        else:
            img = PIL.Image.fromarray(images[0].transpose(1, 2, 0), 'RGB')
        img.save(os.path.join(output_dir, 'img%08d.png' % idx))
        idx += 1
    print('Extracted %d images.' % idx)


#----------------------------------------------------------------------------


def compare(tfrecord_dir_a, tfrecord_dir_b, ignore_labels):
    max_label_size = 0 if ignore_labels else 'full'
    print('Loading dataset "%s"' % tfrecord_dir_a)
    tflib.init_tf({'gpu_options.allow_growth': True})
    dset_a = dataset.TFRecordDataset(tfrecord_dir_a,
                                     max_label_size=max_label_size,
                                     repeat=False,
                                     shuffle_mb=0)
    print('Loading dataset "%s"' % tfrecord_dir_b)
    dset_b = dataset.TFRecordDataset(tfrecord_dir_b,
                                     max_label_size=max_label_size,
                                     repeat=False,
                                     shuffle_mb=0)
    tflib.init_uninitialized_vars()

    print('Comparing datasets')
    idx = 0
    identical_images = 0
    identical_labels = 0
    while True:
        if idx % 100 == 0:
            print('%d\r' % idx, end='', flush=True)
        try:
            images_a, labels_a = dset_a.get_minibatch_np(1)
        except tf.errors.OutOfRangeError:
            images_a, labels_a = None, None
        try:
            images_b, labels_b = dset_b.get_minibatch_np(1)
        except tf.errors.OutOfRangeError:
            images_b, labels_b = None, None
        if images_a is None or images_b is None:
            if images_a is not None or images_b is not None:
                print('Datasets contain different number of images')
            break
        if images_a.shape == images_b.shape and np.all(images_a == images_b):
            identical_images += 1
        else:
            print('Image %d is different' % idx)
        if labels_a.shape == labels_b.shape and np.all(labels_a == labels_b):
            identical_labels += 1
        else:
            print('Label %d is different' % idx)
        idx += 1
    print('Identical images: %d / %d' % (identical_images, idx))
    if not ignore_labels:
        print('Identical labels: %d / %d' % (identical_labels, idx))


#----------------------------------------------------------------------------


def create_mnist(tfrecord_dir, mnist_dir):
    print('Loading MNIST from "%s"' % mnist_dir)
    import gzip
    with gzip.open(os.path.join(mnist_dir, 'train-images-idx3-ubyte.gz'),
                   'rb') as file:
        images = np.frombuffer(file.read(), np.uint8, offset=16)
    with gzip.open(os.path.join(mnist_dir, 'train-labels-idx1-ubyte.gz'),
                   'rb') as file:
        labels = np.frombuffer(file.read(), np.uint8, offset=8)
    images = images.reshape(-1, 1, 28, 28)
    images = np.pad(images, [(0, 0), (0, 0), (2, 2), (2, 2)],
                    'constant',
                    constant_values=0)
    assert images.shape == (60000, 1, 32, 32) and images.dtype == np.uint8
    assert labels.shape == (60000, ) and labels.dtype == np.uint8
    assert np.min(images) == 0 and np.max(images) == 255
    assert np.min(labels) == 0 and np.max(labels) == 9
    onehot = np.zeros((labels.size, np.max(labels) + 1), dtype=np.float32)
    onehot[np.arange(labels.size), labels] = 1.0

    with TFRecordExporter(tfrecord_dir, images.shape[0]) as tfr:
        order = tfr.choose_shuffled_order()
        for idx in range(order.size):
            tfr.add_image(images[order[idx]])
        tfr.add_labels(onehot[order])

def create_test_mnist(tfrecord_dir, mnist_dir):
    print('Loading MNIST from "%s"' % mnist_dir)
    import gzip
    with gzip.open(os.path.join(mnist_dir, 't10k-images-idx3-ubyte.gz'),
                   'rb') as file:
        images = np.frombuffer(file.read(), np.uint8, offset=16)
    with gzip.open(os.path.join(mnist_dir, 't10k-labels-idx1-ubyte.gz'),
                   'rb') as file:
        labels = np.frombuffer(file.read(), np.uint8, offset=8)
    images = images.reshape(-1, 1, 28, 28)
    images = np.pad(images, [(0, 0), (0, 0), (2, 2), (2, 2)],
                    'constant',
                    constant_values=0)
    assert images.shape == (10000, 1, 32, 32) and images.dtype == np.uint8
    assert labels.shape == (10000, ) and labels.dtype == np.uint8
    assert np.min(images) == 0 and np.max(images) == 255
    assert np.min(labels) == 0 and np.max(labels) == 9
    onehot = np.zeros((labels.size, np.max(labels) + 1), dtype=np.float32)
    onehot[np.arange(labels.size), labels] = 1.0

    with TFRecordExporter(tfrecord_dir, images.shape[0]) as tfr:
        order = tfr.choose_shuffled_order()
        for idx in range(order.size):
            tfr.add_image(images[order[idx]])
        tfr.add_labels(onehot[order])


#----------------------------------------------------------------------------


def create_mnistrgb(tfrecord_dir,
                    mnist_dir,
                    num_images=1000000,
                    random_seed=123):
    print('Loading MNIST from "%s"' % mnist_dir)
    import gzip
    with gzip.open(os.path.join(mnist_dir, 'train-images-idx3-ubyte.gz'),
                   'rb') as file:
        images = np.frombuffer(file.read(), np.uint8, offset=16)
    images = images.reshape(-1, 28, 28)
    images = np.pad(images, [(0, 0), (2, 2), (2, 2)],
                    'constant',
                    constant_values=0)
    assert images.shape == (60000, 32, 32) and images.dtype == np.uint8
    assert np.min(images) == 0 and np.max(images) == 255

    with TFRecordExporter(tfrecord_dir, num_images) as tfr:
        rnd = np.random.RandomState(random_seed)
        for _idx in range(num_images):
            tfr.add_image(images[rnd.randint(images.shape[0], size=3)])


#----------------------------------------------------------------------------


def create_cifar10(tfrecord_dir, cifar10_dir):
    print('Loading CIFAR-10 from "%s"' % cifar10_dir)
    import pickle
    images = []
    labels = []
    for batch in range(1, 6):
        with open(os.path.join(cifar10_dir, 'data_batch_%d' % batch),
                  'rb') as file:
            data = pickle.load(file, encoding='latin1')
        images.append(data['data'].reshape(-1, 3, 32, 32))
        labels.append(data['labels'])
    images = np.concatenate(images)
    labels = np.concatenate(labels)
    assert images.shape == (50000, 3, 32, 32) and images.dtype == np.uint8
    assert labels.shape == (50000, ) and labels.dtype == np.int32
    assert np.min(images) == 0 and np.max(images) == 255
    assert np.min(labels) == 0 and np.max(labels) == 9
    onehot = np.zeros((labels.size, np.max(labels) + 1), dtype=np.float32)
    onehot[np.arange(labels.size), labels] = 1.0

    with TFRecordExporter(tfrecord_dir, images.shape[0]) as tfr:
        order = tfr.choose_shuffled_order()
        for idx in range(order.size):
            tfr.add_image(images[order[idx]])
        tfr.add_labels(onehot[order])


#----------------------------------------------------------------------------


def create_cifar100(tfrecord_dir, cifar100_dir):
    print('Loading CIFAR-100 from "%s"' % cifar100_dir)
    import pickle
    with open(os.path.join(cifar100_dir, 'train'), 'rb') as file:
        data = pickle.load(file, encoding='latin1')
    images = data['data'].reshape(-1, 3, 32, 32)
    labels = np.array(data['fine_labels'])
    assert images.shape == (50000, 3, 32, 32) and images.dtype == np.uint8
    assert labels.shape == (50000, ) and labels.dtype == np.int32
    assert np.min(images) == 0 and np.max(images) == 255
    assert np.min(labels) == 0 and np.max(labels) == 99
    onehot = np.zeros((labels.size, np.max(labels) + 1), dtype=np.float32)
    onehot[np.arange(labels.size), labels] = 1.0

    with TFRecordExporter(tfrecord_dir, images.shape[0]) as tfr:
        order = tfr.choose_shuffled_order()
        for idx in range(order.size):
            tfr.add_image(images[order[idx]])
        tfr.add_labels(onehot[order])


#----------------------------------------------------------------------------


def create_svhn(tfrecord_dir, svhn_dir):
    print('Loading SVHN from "%s"' % svhn_dir)
    import pickle
    images = []
    labels = []
    for batch in range(1, 4):
        with open(os.path.join(svhn_dir, 'train_%d.pkl' % batch),
                  'rb') as file:
            data = pickle.load(file, encoding='latin1')
        images.append(data[0])
        labels.append(data[1])
    images = np.concatenate(images)
    labels = np.concatenate(labels)
    assert images.shape == (73257, 3, 32, 32) and images.dtype == np.uint8
    assert labels.shape == (73257, ) and labels.dtype == np.uint8
    assert np.min(images) == 0 and np.max(images) == 255
    assert np.min(labels) == 0 and np.max(labels) == 9
    onehot = np.zeros((labels.size, np.max(labels) + 1), dtype=np.float32)
    onehot[np.arange(labels.size), labels] = 1.0

    with TFRecordExporter(tfrecord_dir, images.shape[0]) as tfr:
        order = tfr.choose_shuffled_order()
        for idx in range(order.size):
            tfr.add_image(images[order[idx]])
        tfr.add_labels(onehot[order])


#----------------------------------------------------------------------------


def create_lsun(tfrecord_dir, lmdb_dir, resolution=256, max_images=None):
    print('Loading LSUN dataset from "%s"' % lmdb_dir)
    import lmdb  # pip install lmdb # pylint: disable=import-error
    import cv2  # pip install opencv-python
    import io
    with lmdb.open(lmdb_dir, readonly=True).begin(write=False) as txn:
        total_images = txn.stat()['entries']  # pylint: disable=no-value-for-parameter
        if max_images is None:
            max_images = total_images
        with TFRecordExporter(tfrecord_dir, max_images) as tfr:
            for _idx, (_key, value) in enumerate(txn.cursor()):
                try:
                    try:
                        img = cv2.imdecode(
                            np.fromstring(value, dtype=np.uint8), 1)
                        if img is None:
                            raise IOError('cv2.imdecode failed')
                        img = img[:, :, ::-1]  # BGR => RGB
                    except IOError:
                        img = np.asarray(PIL.Image.open(io.BytesIO(value)))
                    crop = np.min(img.shape[:2])
                    img = img[(img.shape[0] - crop) //
                              2:(img.shape[0] + crop) //
                              2, (img.shape[1] - crop) //
                              2:(img.shape[1] + crop) // 2]
                    img = PIL.Image.fromarray(img, 'RGB')
                    img = img.resize((resolution, resolution),
                                     PIL.Image.ANTIALIAS)
                    img = np.asarray(img)
                    img = img.transpose([2, 0, 1])  # HWC => CHW
                    tfr.add_image(img)
                except:
                    print(sys.exc_info()[1])
                if tfr.cur_images == max_images:
                    break


#----------------------------------------------------------------------------


def create_lsun_wide(tfrecord_dir,
                     lmdb_dir,
                     width=512,
                     height=384,
                     max_images=None):
    assert width == 2**int(np.round(np.log2(width)))
    assert height <= width
    print('Loading LSUN dataset from "%s"' % lmdb_dir)
    import lmdb  # pip install lmdb # pylint: disable=import-error
    import cv2  # pip install opencv-python
    import io
    with lmdb.open(lmdb_dir, readonly=True).begin(write=False) as txn:
        total_images = txn.stat()['entries']  # pylint: disable=no-value-for-parameter
        if max_images is None:
            max_images = total_images
        with TFRecordExporter(tfrecord_dir, max_images,
                              print_progress=False) as tfr:
            for idx, (_key, value) in enumerate(txn.cursor()):
                try:
                    try:
                        img = cv2.imdecode(
                            np.fromstring(value, dtype=np.uint8), 1)
                        if img is None:
                            raise IOError('cv2.imdecode failed')
                        img = img[:, :, ::-1]  # BGR => RGB
                    except IOError:
                        img = np.asarray(PIL.Image.open(io.BytesIO(value)))

                    ch = int(np.round(width * img.shape[0] / img.shape[1]))
                    if img.shape[1] < width or ch < height:
                        continue

                    img = img[(img.shape[0] - ch) // 2:(img.shape[0] + ch) //
                              2]
                    img = PIL.Image.fromarray(img, 'RGB')
                    img = img.resize((width, height), PIL.Image.ANTIALIAS)
                    img = np.asarray(img)
                    img = img.transpose([2, 0, 1])  # HWC => CHW

                    canvas = np.zeros([3, width, width], dtype=np.uint8)
                    canvas[:, (width - height) // 2:(width + height) //
                           2] = img
                    tfr.add_image(canvas)
                    print('\r%d / %d => %d ' %
                          (idx + 1, total_images, tfr.cur_images),
                          end='')

                except:
                    print(sys.exc_info()[1])
                if tfr.cur_images == max_images:
                    break
    print()


#----------------------------------------------------------------------------


def create_celeba(tfrecord_dir, celeba_dir, cx=89, cy=121):
    print('Loading CelebA from "%s"' % celeba_dir)
    glob_pattern = os.path.join(celeba_dir, 'img_align_celeba_png', '*.png')
    image_filenames = sorted(glob.glob(glob_pattern))
    expected_images = 202599
    if len(image_filenames) != expected_images:
        error('Expected to find %d images' % expected_images)

    with TFRecordExporter(tfrecord_dir, len(image_filenames)) as tfr:
        order = tfr.choose_shuffled_order()
        for idx in range(order.size):
            img = np.asarray(PIL.Image.open(image_filenames[order[idx]]))
            assert img.shape == (218, 178, 3)
            img = img[cy - 64:cy + 64, cx - 64:cx + 64]
            img = img.transpose(2, 0, 1)  # HWC => CHW
            tfr.add_image(img)


#----------------------------------------------------------------------------


def create_from_images(tfrecord_dir, image_dir, shuffle):
    print('Loading images from "%s"' % image_dir)
    image_filenames = sorted(glob.glob(os.path.join(image_dir, '*')))
    if len(image_filenames) == 0:
        error('No input images found')

    img = np.asarray(PIL.Image.open(image_filenames[0]))
    resolution = img.shape[0]
    channels = img.shape[2] if img.ndim == 3 else 1
    if img.shape[1] != resolution:
        error('Input images must have the same width and height')
    if resolution != 2**int(np.floor(np.log2(resolution))):
        error('Input image resolution must be a power-of-two')
    if channels not in [1, 3]:
        error('Input images must be stored as RGB or grayscale')

    with TFRecordExporter(tfrecord_dir, len(image_filenames)) as tfr:
        order = tfr.choose_shuffled_order() if shuffle else np.arange(
            len(image_filenames))
        for idx in range(order.size):
            img = np.asarray(PIL.Image.open(image_filenames[order[idx]]))
            if channels == 1:
                img = img[np.newaxis, :, :]  # HW => CHW
            else:
                img = img.transpose([2, 0, 1])  # HWC => CHW
            tfr.add_image(img)


#----------------------------------------------------------------------------


def create_from_dsprites_npz(tfrecord_dir, dsprites_filename, shuffle,
                             shape_only, sample_color):
    print('Loading images from "%s"' % dsprites_filename)
    data = np.load(dsprites_filename, encoding='latin1', allow_pickle=True)
    images = data['imgs'] * 255
    labels = data['latents_classes']
    if shape_only:
        labels = convert_to_shape(labels)
    img = images[0]
    resolution = img.shape[0]
    channels = img.shape[2] if img.ndim == 3 else 1
    if img.shape[1] != resolution:
        error('Input images must have the same width and height')
    if resolution != 2**int(np.floor(np.log2(resolution))):
        error('Input image resolution must be a power-of-two')
    if channels not in [1, 3]:
        error('Input images must be stored as RGB or grayscale')

    with TFRecordExporter(tfrecord_dir, images.shape[0]) as tfr:
        order = tfr.choose_shuffled_order() if shuffle else np.arange(
            images.shape[0])
        for idx in range(order.size):
            img = images[order[idx]]
            if channels == 1:
                img = img[np.newaxis, :, :]  # HW => CHW
            else:
                img = img.transpose([2, 0, 1])  # HWC => CHW

            if sample_color:
                color = np.repeat(np.repeat(np.random.uniform(
                    0.5, 1, [3, 1, 1]),
                                            img.shape[1],
                                            axis=1),
                                  img.shape[2],
                                  axis=2)
                img = img * color
            tfr.add_image(img)
        tfr.add_labels(labels[order])


#----------------------------------------------------------------------------


def convert_to_shape(labels):
    labels_shape = labels[:, 1]
    # labels_shape = labels_shape.astype(np.int)
    labels_onehot = np.zeros((labels_shape.size, labels_shape.max() + 1))
    labels_onehot[np.arange(labels_shape.size), labels_shape] = 1
    return labels_onehot


#----------------------------------------------------------------------------


def create_from_hdf5(tfrecord_dir, hdf5_filename, shuffle):
    print('Loading HDF5 archive from "%s"' % hdf5_filename)
    import h5py  # conda install h5py
    with h5py.File(hdf5_filename, 'r') as hdf5_file:
        hdf5_data = max([
            value for key, value in hdf5_file.items() if key.startswith('data')
        ],
                        key=lambda lod: lod.shape[3])
        with TFRecordExporter(tfrecord_dir, hdf5_data.shape[0]) as tfr:
            order = tfr.choose_shuffled_order() if shuffle else np.arange(
                hdf5_data.shape[0])
            for idx in range(order.size):
                tfr.add_image(hdf5_data[order[idx]])
            npy_filename = os.path.splitext(hdf5_filename)[0] + '-labels.npy'
            if os.path.isfile(npy_filename):
                tfr.add_labels(np.load(npy_filename)[order])


#----------------------------------------------------------------------------


def execute_cmdline(argv):
    prog = argv[0]
    parser = argparse.ArgumentParser(
        prog=prog,
        description=
        'Tool for creating multi-resolution TFRecords datasets for StyleGAN and ProGAN.',
        epilog='Type "%s <command> -h" for more information.' % prog)

    subparsers = parser.add_subparsers(dest='command')
    subparsers.required = True

    def add_command(cmd, desc, example=None):
        epilog = 'Example: %s %s' % (prog,
                                     example) if example is not None else None
        return subparsers.add_parser(cmd,
                                     description=desc,
                                     help=desc,
                                     epilog=epilog)

    p = add_command('display', 'Display images in dataset.',
                    'display datasets/mnist')
    p.add_argument('tfrecord_dir', help='Directory containing dataset')

    p = add_command('extract', 'Extract images from dataset.',
                    'extract datasets/mnist mnist-images')
    p.add_argument('tfrecord_dir', help='Directory containing dataset')
    p.add_argument('output_dir', help='Directory to extract the images into')

    p = add_command('compare', 'Compare two datasets.',
                    'compare datasets/mydataset datasets/mnist')
    p.add_argument('tfrecord_dir_a', help='Directory containing first dataset')
    p.add_argument('tfrecord_dir_b',
                   help='Directory containing second dataset')
    p.add_argument('--ignore_labels',
                   help='Ignore labels (default: 0)',
                   type=int,
                   default=0)

    p = add_command('create_mnist', 'Create dataset for MNIST.',
                    'create_mnist datasets/mnist ~/downloads/mnist')
    p.add_argument('tfrecord_dir', help='New dataset directory to be created')
    p.add_argument('mnist_dir', help='Directory containing MNIST')

    p = add_command('create_test_mnist', 'Create testdataset for MNIST.',
                    'create_test_mnist datasets/mnist ~/downloads/mnist')
    p.add_argument('tfrecord_dir', help='New dataset directory to be created')
    p.add_argument('mnist_dir', help='Directory containing test MNIST')

    p = add_command('create_mnistrgb', 'Create dataset for MNIST-RGB.',
                    'create_mnistrgb datasets/mnistrgb ~/downloads/mnist')
    p.add_argument('tfrecord_dir', help='New dataset directory to be created')
    p.add_argument('mnist_dir', help='Directory containing MNIST')
    p.add_argument(
        '--num_images',
        help='Number of composite images to create (default: 1000000)',
        type=int,
        default=1000000)
    p.add_argument('--random_seed',
                   help='Random seed (default: 123)',
                   type=int,
                   default=123)

    p = add_command('create_cifar10', 'Create dataset for CIFAR-10.',
                    'create_cifar10 datasets/cifar10 ~/downloads/cifar10')
    p.add_argument('tfrecord_dir', help='New dataset directory to be created')
    p.add_argument('cifar10_dir', help='Directory containing CIFAR-10')

    p = add_command('create_cifar100', 'Create dataset for CIFAR-100.',
                    'create_cifar100 datasets/cifar100 ~/downloads/cifar100')
    p.add_argument('tfrecord_dir', help='New dataset directory to be created')
    p.add_argument('cifar100_dir', help='Directory containing CIFAR-100')

    p = add_command('create_svhn', 'Create dataset for SVHN.',
                    'create_svhn datasets/svhn ~/downloads/svhn')
    p.add_argument('tfrecord_dir', help='New dataset directory to be created')
    p.add_argument('svhn_dir', help='Directory containing SVHN')

    p = add_command(
        'create_lsun', 'Create dataset for single LSUN category.',
        'create_lsun datasets/lsun-car-100k ~/downloads/lsun/car_lmdb --resolution 256 --max_images 100000'
    )
    p.add_argument('tfrecord_dir', help='New dataset directory to be created')
    p.add_argument('lmdb_dir', help='Directory containing LMDB database')
    p.add_argument('--resolution',
                   help='Output resolution (default: 256)',
                   type=int,
                   default=256)
    p.add_argument('--max_images',
                   help='Maximum number of images (default: none)',
                   type=int,
                   default=None)

    p = add_command(
        'create_lsun_wide',
        'Create LSUN dataset with non-square aspect ratio.',
        'create_lsun_wide datasets/lsun-car-512x384 ~/downloads/lsun/car_lmdb --width 512 --height 384'
    )
    p.add_argument('tfrecord_dir', help='New dataset directory to be created')
    p.add_argument('lmdb_dir', help='Directory containing LMDB database')
    p.add_argument('--width',
                   help='Output width (default: 512)',
                   type=int,
                   default=512)
    p.add_argument('--height',
                   help='Output height (default: 384)',
                   type=int,
                   default=384)
    p.add_argument('--max_images',
                   help='Maximum number of images (default: none)',
                   type=int,
                   default=None)

    p = add_command('create_celeba', 'Create dataset for CelebA.',
                    'create_celeba datasets/celeba ~/downloads/celeba')
    p.add_argument('tfrecord_dir', help='New dataset directory to be created')
    p.add_argument('celeba_dir', help='Directory containing CelebA')
    p.add_argument('--cx',
                   help='Center X coordinate (default: 89)',
                   type=int,
                   default=89)
    p.add_argument('--cy',
                   help='Center Y coordinate (default: 121)',
                   type=int,
                   default=121)

    p = add_command('create_from_images',
                    'Create dataset from a directory full of images.',
                    'create_from_images datasets/mydataset myimagedir')
    p.add_argument('tfrecord_dir', help='New dataset directory to be created')
    p.add_argument('image_dir', help='Directory containing the images')
    p.add_argument('--shuffle',
                   help='Randomize image order (default: 1)',
                   type=int,
                   default=1)

    p = add_command(
        'create_from_dsprites_npz', 'Create dataset from a dsprites_filename.',
        'create_from_dsprites_npz datasets/mydataset dsprites_py3.npz')
    p.add_argument('tfrecord_dir', help='New dataset directory to be created')
    p.add_argument('dsprites_filename',
                   help='dsprites_filename containing the images')
    p.add_argument('--shuffle',
                   help='Randomize image order (default: 1)',
                   type=int,
                   default=0)
    p.add_argument('--shape_only',
                   help='Use only shape as label',
                   type=int,
                   default=1)
    p.add_argument('--sample_color',
                   help='Build dsprites color dataset.',
                   type=int,
                   default=0)

    # p = add_command(    'create_dsprites_shape_labels_from_tfr', 'Create shape labels from a dsprites_tfr_label.',
    # 'create_from_dsprites_npz datasets/mydataset dsprites_py3.npz')
    # p.add_argument(     'tfrecord_dir',     help='New dataset directory to be created')
    # p.add_argument(     'tfr_filename',        help='dsprites_label_filename containing the label')

    p = add_command(
        'create_from_hdf5', 'Create dataset from legacy HDF5 archive.',
        'create_from_hdf5 datasets/celebahq ~/downloads/celeba-hq-1024x1024.h5'
    )
    p.add_argument('tfrecord_dir', help='New dataset directory to be created')
    p.add_argument('hdf5_filename', help='HDF5 archive containing the images')
    p.add_argument('--shuffle',
                   help='Randomize image order (default: 1)',
                   type=int,
                   default=1)

    args = parser.parse_args(argv[1:] if len(argv) > 1 else ['-h'])
    func = globals()[args.command]
    del args.command
    func(**vars(args))


#----------------------------------------------------------------------------

if __name__ == "__main__":
    execute_cmdline(sys.argv)

#----------------------------------------------------------------------------