Repository: mila-udem/fuel
Pull request #300 (2016-01-28 17:32:27)

CelebA dataset
For the time being, it's only the aligned and cropped version (as well as the scaled and cropped 64x64 version used in e.g. [Larsen et. al. Autoencoding beyond pixels using a learned similarity metric](http://arxiv.org/pdf/1512.09300.pdf).

Someone can later integrate the non-cropped version of the dataset should the need arise.

@bartvm @dwf ?

bartvm: LGTM

Patch:
diff --git a/fuel/converters/__init__.py b/fuel/converters/__init__.py
index bca7841..2602bbe 100644
--- a/fuel/converters/__init__.py
+++ b/fuel/converters/__init__.py
@@ -13,6 +13,7 @@ command-line arguments, and is expected to download the required files.
 from fuel.converters import adult
 from fuel.converters import binarized_mnist
 from fuel.converters import caltech101_silhouettes
+from fuel.converters import celeba
 from fuel.converters import cifar10
 from fuel.converters import cifar100
 from fuel.converters import dogs_vs_cats
@@ -26,6 +27,7 @@ all_converters = (
     ('adult', adult.fill_subparser),
     ('binarized_mnist', binarized_mnist.fill_subparser),
     ('caltech101_silhouettes', caltech101_silhouettes.fill_subparser),
+    ('celeba', celeba.fill_subparser),
     ('cifar10', cifar10.fill_subparser),
     ('cifar100', cifar100.fill_subparser),
     ('dogs_vs_cats', dogs_vs_cats.fill_subparser),
diff --git a/fuel/converters/base.py b/fuel/converters/base.py
index 4c6a4ff..efade27 100644
--- a/fuel/converters/base.py
+++ b/fuel/converters/base.py
@@ -117,7 +117,7 @@ def progress_bar(name, maxval, prefix='Converting'):
     """
     widgets = ['{} {}: '.format(prefix, name), Percentage(), ' ',
                Bar(marker='=', left='[', right=']'), ' ', ETA()]
-    bar = ProgressBar(widgets=widgets, maxval=maxval, fd=sys.stdout).start()
+    bar = ProgressBar(widgets=widgets, max_value=maxval, fd=sys.stdout).start()
     try:
         yield bar
     finally:
diff --git a/fuel/converters/celeba.py b/fuel/converters/celeba.py
new file mode 100644
index 0000000..b7529d5
--- /dev/null
+++ b/fuel/converters/celeba.py
@@ -0,0 +1,213 @@
+import os
+import zipfile
+
+import h5py
+import numpy
+from six.moves import range
+from PIL import Image
+
+from fuel.converters.base import check_exists, progress_bar
+from fuel.datasets import H5PYDataset
+
+IMAGE_FILE = 'img_align_celeba.zip'
+ATTRIBUTES_FILE = 'list_attr_celeba.txt'
+DATASET_FILES = [IMAGE_FILE, ATTRIBUTES_FILE]
+NUM_EXAMPLES = 202599
+TRAIN_STOP = 162770
+VALID_STOP = 182637
+OUTPUT_FILENAME = 'celeba_aligned_cropped.hdf5'
+
+
+def _initialize_conversion(directory, output_path, image_shape):
+    h5file = h5py.File(output_path, mode='w')
+    split_dict = {
+        'train': {
+            'features': (0, TRAIN_STOP),
+            'targets': (0, TRAIN_STOP)},
+        'valid': {
+            'features': (TRAIN_STOP, VALID_STOP),
+            'targets': (TRAIN_STOP, VALID_STOP)},
+        'test': {
+            'features': (VALID_STOP, NUM_EXAMPLES),
+            'targets': (VALID_STOP, NUM_EXAMPLES)}}
+    h5file.attrs['split'] = H5PYDataset.create_split_array(split_dict)
+
+    targets_dataset = h5file.create_dataset(
+        'targets', (NUM_EXAMPLES, 40), dtype='uint8')
+    targets_dataset.dims[0].label = 'batch'
+    targets_dataset.dims[1].label = 'target'
+    targets_dataset[...] = (
+        numpy.loadtxt(os.path.join(directory, ATTRIBUTES_FILE), dtype='int32',
+                      skiprows=2, usecols=tuple(range(1, 41))) +
+        1) / 2
+
+    features_dataset = h5file.create_dataset(
+        'features', (NUM_EXAMPLES, 3) + image_shape, dtype='uint8')
+    features_dataset.dims[0].label = 'batch'
+    features_dataset.dims[1].label = 'channel'
+    features_dataset.dims[2].label = 'height'
+    features_dataset.dims[3].label = 'width'
+
+    return h5file
+
+
+@check_exists(required_files=DATASET_FILES)
+def convert_celeba_aligned_cropped(directory, output_directory,
+                                   output_filename=OUTPUT_FILENAME):
+    """Converts the aligned and cropped CelebA dataset to HDF5.
+
+    Converts the CelebA dataset to an HDF5 dataset compatible with
+    :class:`fuel.datasets.CelebA`. The converted dataset is saved as
+    'celeba_aligned_cropped.hdf5'.
+
+    It assumes the existence of the following files:
+
+    * `img_align_celeba.zip`
+    * `list_attr_celeba.txt`
+
+    Parameters
+    ----------
+    directory : str
+        Directory in which input files reside.
+    output_directory : str
+        Directory in which to save the converted dataset.
+    output_filename : str, optional
+        Name of the saved dataset. Defaults to
+        'celeba_aligned_cropped.hdf5'.
+
+    Returns
+    -------
+    output_paths : tuple of str
+        Single-element tuple containing the path to the converted
+        dataset.
+
+    """
+    output_path = os.path.join(output_directory, output_filename)
+    h5file = _initialize_conversion(directory, output_path, (218, 178))
+
+    features_dataset = h5file['features']
+    image_file_path = os.path.join(directory, IMAGE_FILE)
+    with zipfile.ZipFile(image_file_path, 'r') as image_file:
+        with progress_bar('images', NUM_EXAMPLES) as bar:
+            for i in range(NUM_EXAMPLES):
+                image_name = 'img_align_celeba/{:06d}.jpg'.format(i + 1)
+                features_dataset[i] = numpy.asarray(
+                    Image.open(
+                        image_file.open(image_name, 'r'))).transpose(2, 0, 1)
+                bar.update(i + 1)
+
+    h5file.flush()
+    h5file.close()
+
+    return (output_path,)
+
+
+@check_exists(required_files=DATASET_FILES)
+def convert_celeba_64(directory, output_directory,
+                      output_filename='celeba_64.hdf5'):
+    """Converts the 64x64 version of the CelebA dataset to HDF5.
+
+    This converter takes the aligned and cropped version of the
+    CelebA dataset as input and produces a version that's been resized
+    to 78x64 pixels and then center cropped to 64x64 pixels.
+
+    Converts the CelebA dataset to an HDF5 dataset compatible with
+    :class:`fuel.datasets.CelebA`. The converted dataset is saved as
+    'celeba_64.hdf5'.
+
+    It assumes the existence of the following files:
+
+    * `img_align_celeba.zip`
+    * `list_attr_celeba.txt`
+
+    Parameters
+    ----------
+    directory : str
+        Directory in which input files reside.
+    output_directory : str
+        Directory in which to save the converted dataset.
+    output_filename : str, optional
+        Name of the saved dataset. Defaults to 'celeba_64.hdf5'.
+
+    Returns
+    -------
+    output_paths : tuple of str
+        Single-element tuple containing the path to the converted dataset.
+
+    """
+    output_path = os.path.join(output_directory, output_filename)
+    h5file = _initialize_conversion(directory, output_path, (64, 64))
+
+    features_dataset = h5file['features']
+    image_file_path = os.path.join(directory, IMAGE_FILE)
+    with zipfile.ZipFile(image_file_path, 'r') as image_file:
+        with progress_bar('images', NUM_EXAMPLES) as bar:
+            for i in range(NUM_EXAMPLES):
+                image_name = 'img_align_celeba/{:06d}.jpg'.format(i + 1)
+                image = Image.open(
+                    image_file.open(image_name, 'r')).resize(
+                        (64, 78), Image.ANTIALIAS).crop((0, 7, 64, 64 + 7))
+                features_dataset[i] = numpy.asarray(image).transpose(2, 0, 1)
+                bar.update(i + 1)
+
+    h5file.flush()
+    h5file.close()
+
+    return (output_path,)
+
+
+def convert_celeba(which_format, directory, output_directory,
+                   output_filename=None):
+    """Converts the CelebA dataset to HDF5.
+
+    Converts the CelebA dataset to an HDF5 dataset compatible with
+    :class:`fuel.datasets.CelebA`. The converted dataset is
+    saved as 'celeba_aligned_cropped.hdf5' or 'celeba_64.hdf5',
+    depending on the `which_format` argument.
+
+    Parameters
+    ----------
+    which_format : str
+        Either 'aligned_cropped' or '64'. Determines which format
+        to convert to.
+    directory : str
+        Directory in which input files reside.
+    output_directory : str
+        Directory in which to save the converted dataset.
+    output_filename : str, optional
+        Name of the saved dataset. Defaults to
+        'celeba_aligned_cropped.hdf5' or 'celeba_64.hdf5',
+        depending on `which_format`.
+
+    Returns
+    -------
+    output_paths : tuple of str
+        Single-element tuple containing the path to the converted dataset.
+
+    """
+    if which_format not in ('aligned_cropped', '64'):
+        raise ValueError("CelebA format needs to be either "
+                         "'aligned_cropped' or '64'.")
+    if not output_filename:
+        output_filename = 'celeba_{}.hdf5'.format(which_format)
+    if which_format == 'aligned_cropped':
+        return convert_celeba_aligned_cropped(
+            directory, output_directory, output_filename)
+    else:
+        return convert_celeba_64(
+            directory, output_directory, output_filename)
+
+
+def fill_subparser(subparser):
+    """Sets up a subparser to convert the CelebA dataset files.
+
+    Parameters
+    ----------
+    subparser : :class:`argparse.ArgumentParser`
+        Subparser handling the `celeba` command.
+
+    """
+    subparser.add_argument(
+        "which_format", help="which dataset format", type=str,
+        choices=('aligned_cropped', '64'))
+    return convert_celeba
diff --git a/fuel/converters/ilsvrc2010.py b/fuel/converters/ilsvrc2010.py
index e524551..a67cd01 100644
--- a/fuel/converters/ilsvrc2010.py
+++ b/fuel/converters/ilsvrc2010.py
@@ -10,13 +10,12 @@ import os.path
 import h5py
 import numpy
 from picklable_itertools.extras import equizip
-from progressbar import ProgressBar
 from PIL import Image
 from scipy.io.matlab import loadmat
 from six.moves import zip, xrange
 import zmq
 
-from fuel.converters.base import check_exists
+from fuel.converters.base import check_exists, progress_bar
 from fuel.datasets import H5PYDataset
 from fuel.utils.formats import tar_open
 from fuel.utils.parallel import producer_consumer
@@ -302,7 +301,7 @@ def image_consumer(socket, hdf5_file, num_expected, shuffle_seed=None,
         received examples. Defaults to 0.
 
     """
-    with ProgressBar(maxval=num_expected) as pb:
+    with progress_bar('images', maxval=num_expected) as pb:
         if shuffle_seed is None:
             index_gen = iter(xrange(num_expected))
         else:
diff --git a/fuel/datasets/__init__.py b/fuel/datasets/__init__.py
index 5ad5549..376ad05 100644
--- a/fuel/datasets/__init__.py
+++ b/fuel/datasets/__init__.py
@@ -5,6 +5,7 @@ from fuel.datasets.base import (Dataset, IterableDataset,
 from fuel.datasets.hdf5 import H5PYDataset
 from fuel.datasets.adult import Adult
 from fuel.datasets.binarized_mnist import BinarizedMNIST
+from fuel.datasets.celeba import CelebA
 from fuel.datasets.cifar10 import CIFAR10
 from fuel.datasets.cifar100 import CIFAR100
 from fuel.datasets.caltech101_silhouettes import CalTech101Silhouettes
diff --git a/fuel/datasets/celeba.py b/fuel/datasets/celeba.py
new file mode 100644
index 0000000..8e43a83
--- /dev/null
+++ b/fuel/datasets/celeba.py
@@ -0,0 +1,50 @@
+from fuel.datasets import H5PYDataset
+from fuel.transformers.defaults import uint8_pixels_to_floatX
+from fuel.utils import find_in_data_path
+
+
+class CelebA(H5PYDataset):
+    """The CelebFaces Attributes Dataset (CelebA) dataset.
+
+    CelebA is a large-scale face
+    attributes dataset with more than 200K celebrity images, each
+    with 40 attribute annotations. The images in this dataset cover
+    large pose variations and background clutter. CelebA has large
+    diversities, large quantities, and rich annotations, including:
+
+    * 10,177 number of identities
+    * 202,599 number of face images
+    * 5 landmark locations per image
+    * 40 binary attributes annotations per image.
+
+    The dataset can be employed as the training and test sets for
+    the following computer vision tasks:
+
+    * face attribute recognition
+    * face detection
+    * landmark (or facial part) localization
+
+    Parameters
+    ----------
+    which_format : {'aligned_cropped, '64'}
+        Either the aligned and cropped version of CelebA, or
+        a 64x64 version of it.
+    which_sets : tuple of str
+        Which split to load. Valid values are 'train', 'valid' and
+        'test' corresponding to the training set (162,770 examples), the
+        validation set (19,867 examples) and the test set (19,962
+        examples).
+
+    """
+    _filename = 'celeba_{}.hdf5'
+    default_transformers = uint8_pixels_to_floatX(('features',))
+
+    def __init__(self, which_format, which_sets, **kwargs):
+        self.which_format = which_format
+        super(CelebA, self).__init__(
+            file_or_path=find_in_data_path(self.filename),
+            which_sets=which_sets, **kwargs)
+
+    @property
+    def filename(self):
+        return self._filename.format(self.which_format)
diff --git a/fuel/downloaders/__init__.py b/fuel/downloaders/__init__.py
index d933b68..ab968cc 100644
--- a/fuel/downloaders/__init__.py
+++ b/fuel/downloaders/__init__.py
@@ -9,6 +9,7 @@ Download functions accept two arguments:
 from fuel.downloaders import adult
 from fuel.downloaders import binarized_mnist
 from fuel.downloaders import caltech101_silhouettes
+from fuel.downloaders import celeba
 from fuel.downloaders import cifar10
 from fuel.downloaders import cifar100
 from fuel.downloaders import dogs_vs_cats
@@ -21,6 +22,7 @@ all_downloaders = (
     ('adult', adult.fill_subparser),
     ('binarized_mnist', binarized_mnist.fill_subparser),
     ('caltech101_silhouettes', caltech101_silhouettes.fill_subparser),
+    ('celeba', celeba.fill_subparser),
     ('cifar10', cifar10.fill_subparser),
     ('cifar100', cifar100.fill_subparser),
     ('iris', iris.fill_subparser),
diff --git a/fuel/downloaders/base.py b/fuel/downloaders/base.py
index 8cb9e2e..0656c5c 100644
--- a/fuel/downloaders/base.py
+++ b/fuel/downloaders/base.py
@@ -28,7 +28,7 @@ def progress_bar(name, maxval):
                    FileTransferSpeed()]
     else:
         widgets = ['{}: '.format(name), ' ', Timer(), ' ', FileTransferSpeed()]
-    bar = ProgressBar(widgets=widgets, maxval=maxval, fd=sys.stdout).start()
+    bar = ProgressBar(widgets=widgets, max_value=maxval, fd=sys.stdout).start()
     try:
         yield bar
     finally:
diff --git a/fuel/downloaders/celeba.py b/fuel/downloaders/celeba.py
new file mode 100644
index 0000000..81123c4
--- /dev/null
+++ b/fuel/downloaders/celeba.py
@@ -0,0 +1,19 @@
+from fuel.downloaders.base import default_downloader
+
+
+def fill_subparser(subparser):
+    """Sets up a subparser to download the CelebA dataset file.
+
+    Parameters
+    ----------
+    subparser : :class:`argparse.ArgumentParser`
+        Subparser handling the `celeba` command.
+
+    """
+    urls = ['https://www.dropbox.com/sh/8oqt9vytwxb3s4r/'
+            'AAB7G69NLjRNqv_tyiULHSVUa/list_attr_celeba.txt?dl=1',
+            'https://www.dropbox.com/sh/8oqt9vytwxb3s4r/'
+            'AADVdnYbokd7TXhpvfWLL3sga/img_align_celeba.zip?dl=1']
+    filenames = ['list_attr_celeba.txt', 'img_align_celeba.zip']
+    subparser.set_defaults(urls=urls, filenames=filenames)
+    return default_downloader


Test patch:
diff --git a/tests/test_celeba.py b/tests/test_celeba.py
new file mode 100644
index 0000000..76551a9
--- /dev/null
+++ b/tests/test_celeba.py
@@ -0,0 +1,29 @@
+import os
+
+import h5py
+import numpy
+from numpy.testing import assert_equal
+
+from fuel import config
+from fuel.datasets import H5PYDataset, CelebA
+
+
+def test_celeba():
+    data_path = config.data_path
+    try:
+        config.data_path = '.'
+        f = h5py.File('celeba_64.hdf5', 'w')
+        f['features'] = numpy.arange(
+            10 * 3 * 64 * 64, dtype='uint8').reshape((10, 3, 64, 64))
+        f['targets'] = numpy.arange(
+            10 * 40, dtype='uint8').reshape((10, 40))
+        split_dict = {'train': {'features': (0, 6), 'targets': (0, 6)},
+                      'valid': {'features': (6, 8), 'targets': (6, 8)},
+                      'test': {'features': (8, 10), 'targets': (8, 10)}}
+        f.attrs['split'] = H5PYDataset.create_split_array(split_dict)
+        f.close()
+        dataset = CelebA(which_format='64', which_sets=('train',))
+        assert_equal(dataset.filename, 'celeba_64.hdf5')
+    finally:
+        config.data_path = data_path
+        os.remove('celeba_64.hdf5')
diff --git a/tests/test_converters.py b/tests/test_converters.py
index 68e1bcd..6eeaffc 100644
--- a/tests/test_converters.py
+++ b/tests/test_converters.py
@@ -7,6 +7,7 @@ import shutil
 import struct
 import tarfile
 import tempfile
+import zipfile
 
 import h5py
 import numpy
@@ -19,7 +20,7 @@ from six.moves import range, zip, cPickle
 from fuel.converters.base import (fill_hdf5_file, check_exists,
                                   MissingInputFiles)
 from fuel.converters import (adult, binarized_mnist, caltech101_silhouettes,
-                             iris, cifar10, cifar100, mnist, svhn)
+                             celeba, iris, cifar10, cifar100, mnist, svhn)
 from fuel.downloaders.caltech101_silhouettes import silhouettes_downloader
 from fuel.downloaders.base import default_downloader
 from fuel.utils import remember_cwd
@@ -290,6 +291,70 @@ class TestBinarizedMNIST(object):
                      ('batch', 'channel', 'height', 'width'))
 
 
+class TestCelebA(object):
+    def setUp(self):
+        numpy.random.seed(21 + 1 + 2016)
+        self.tempdir = tempfile.mkdtemp()
+        with remember_cwd():
+            os.chdir(self.tempdir)
+            self.images = numpy.random.randint(
+                0, 256, (10, 218, 178, 3)).astype('uint8')
+            with zipfile.ZipFile('img_align_celeba.zip', 'w') as image_file:
+                for i, image in enumerate(self.images):
+                    Image.fromarray(image).save('img.jpeg')
+                    image_file.write(
+                        'img.jpeg',
+                        'img_align_celeba/{:06d}.jpg'.format(i + 1))
+            with open('list_attr_celeba.txt', 'w') as attr_file:
+                attr_file.write('mock\nmock')
+                for i in range(1, 11):
+                    attr_file.write('\n{:06d}.jpg'.format(i) + (' 1' * 40))
+
+    def tearDown(self):
+        shutil.rmtree(self.tempdir)
+
+    @mock.patch('fuel.converters.celeba.NUM_EXAMPLES', 10)
+    def test_64_converter(self):
+        parser = argparse.ArgumentParser()
+        subparsers = parser.add_subparsers()
+        subparser = subparsers.add_parser('celeba')
+        convert_function = celeba.fill_subparser(subparser)
+        subparser.set_defaults(
+            directory=self.tempdir, output_directory=self.tempdir,
+            output_filename='celeba_64.hdf5')
+        args = parser.parse_args(['celeba', '64'])
+        args_dict = vars(args)
+        filename, = convert_function(**args_dict)
+        h5file = h5py.File(filename, mode='r')
+        assert_equal(h5file['features'].shape, (10, 3, 64, 64))
+        assert_equal(h5file['targets'].shape, (10, 40))
+
+    @mock.patch('fuel.converters.celeba.NUM_EXAMPLES', 10)
+    def test_aligned_cropped_converter(self):
+        parser = argparse.ArgumentParser()
+        subparsers = parser.add_subparsers()
+        subparser = subparsers.add_parser('celeba')
+        convert_function = celeba.fill_subparser(subparser)
+        subparser.set_defaults(
+            directory=self.tempdir, output_directory=self.tempdir,
+            output_filename='celeba_aligned_cropped.hdf5')
+        args = parser.parse_args(['celeba', 'aligned_cropped'])
+        args_dict = vars(args)
+        filename, = convert_function(**args_dict)
+        h5file = h5py.File(filename, mode='r')
+        assert_equal(h5file['features'].shape, (10, 3, 218, 178))
+        assert_equal(h5file['targets'].shape, (10, 40))
+
+    @mock.patch('fuel.converters.celeba.convert_celeba_64')
+    def test_converter_default_filename(self, mock_converter_64):
+        celeba.convert_celeba('64', './', './')
+        mock_converter_64.assert_called_with('./', './', 'celeba_64.hdf5')
+
+    def test_converter_error_wrong_format(self):
+        assert_raises(
+            ValueError, celeba.convert_celeba, 'a', './', 'mock.hdf5')
+
+
 class TestCIFAR10(object):
     def setUp(self):
         numpy.random.seed(9 + 5 + 2015)
diff --git a/tests/test_downloaders.py b/tests/test_downloaders.py
index 4eae4bd..e41a131 100644
--- a/tests/test_downloaders.py
+++ b/tests/test_downloaders.py
@@ -7,9 +7,9 @@ from functools import wraps
 
 from numpy.testing import assert_equal, assert_raises
 
-from fuel.downloaders import (adult, binarized_mnist,
-                              caltech101_silhouettes, cifar10, cifar100,
-                              dogs_vs_cats, iris, mnist, svhn)
+from fuel.downloaders import (adult, binarized_mnist, caltech101_silhouettes,
+                              celeba, cifar10, cifar100, dogs_vs_cats, iris,
+                              mnist, svhn)
 from fuel.downloaders.base import (download, default_downloader,
                                    filename_from_url, NeedURLPrefix,
                                    ensure_directory_exists)
@@ -158,6 +158,22 @@ def test_iris():
     assert download_function is default_downloader
 
 
+def test_celeba():
+    parser = argparse.ArgumentParser()
+    subparsers = parser.add_subparsers()
+    download_function = celeba.fill_subparser(
+        subparsers.add_parser('celeba'))
+    args = parser.parse_args(['celeba'])
+    filenames = ['list_attr_celeba.txt', 'img_align_celeba.zip']
+    urls = ['https://www.dropbox.com/sh/8oqt9vytwxb3s4r/'
+            'AAB7G69NLjRNqv_tyiULHSVUa/list_attr_celeba.txt?dl=1',
+            'https://www.dropbox.com/sh/8oqt9vytwxb3s4r/'
+            'AADVdnYbokd7TXhpvfWLL3sga/img_align_celeba.zip?dl=1']
+    assert_equal(args.filenames, filenames)
+    assert_equal(args.urls, urls)
+    assert download_function is default_downloader
+
+
 def test_cifar10():
     parser = argparse.ArgumentParser()
     subparsers = parser.add_subparsers()