Add cifar dataset

python/paddle/v2/data_set/cifar.py

@@ -0,0 +1,173 @@
+"""
+CIFAR Dataset.
+
+URL: https://www.cs.toronto.edu/~kriz/cifar.html
+
+the default train_creator, test_creator used for CIFAR-10 dataset.
+"""
+from config import DATA_HOME
+import os
+import hashlib
+import urllib2
+import shutil
+import tarfile
+import cPickle
+import itertools
+import numpy
+
+__all__ = ['CIFAR10', 'CIFAR100', 'train_creator', 'test_creator']
+
+
+def __download_file__(filename, url, md5):
+    def __file_ok__():
+        if not os.path.exists(filename):
+            return False
+        md5_hash = hashlib.md5()
+        with open(filename, 'rb') as f:
+            for chunk in iter(lambda: f.read(4096), b""):
+                md5_hash.update(chunk)
+
+        return md5_hash.hexdigest() == md5
+
+    while not __file_ok__():
+        response = urllib2.urlopen(url)
+        with open(filename, mode='wb') as of:
+            shutil.copyfileobj(fsrc=response, fdst=of)
+
+
+def __read_one_batch__(batch):
+    data = batch['data']
+    labels = batch.get('labels', batch.get('fine_labels', None))
+    assert labels is not None
+    for sample, label in itertools.izip(data, labels):
+        yield (sample / 255.0).astype(numpy.float32), int(label)
+
+
+CIFAR10_URL = 'https://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz'
+CIFAR10_MD5 = 'c58f30108f718f92721af3b95e74349a'
+CIFAR100_URL = 'https://www.cs.toronto.edu/~kriz/cifar-100-python.tar.gz'
+CIFAR100_MD5 = 'eb9058c3a382ffc7106e4002c42a8d85'
+
+
+class CIFAR(object):
+    """
+    CIFAR dataset reader. The base class for CIFAR-10 and CIFAR-100
+
+    :param url: Download url.
+    :param md5: File md5sum
+    :param meta_filename: Meta file name in package.
+    :param train_filename: Train file name in package.
+    :param test_filename: Test file name in package.
+    """
+
+    def __init__(self, url, md5, meta_filename, train_filename, test_filename):
+        filename = os.path.split(url)[-1]
+        assert DATA_HOME is not None
+        filepath = os.path.join(DATA_HOME, md5)
+        if not os.path.exists(filepath):
+            os.makedirs(filepath)
+
+        self.__full_file__ = os.path.join(filepath, filename)
+        self.__meta_filename__ = meta_filename
+        self.__train_filename__ = train_filename
+        self.__test_filename__ = test_filename
+        __download_file__(filename=self.__full_file__, url=url, md5=md5)
+
+    def labels(self):
+        """
+        labels get all dataset label in order.
+        :return: a list of label.
+        :rtype: list[string]
+        """
+        with tarfile.open(self.__full_file__, mode='r') as f:
+            name = [
+                each_item.name for each_item in f
+                if self.__meta_filename__ in each_item.name
+            ][0]
+            meta_f = f.extractfile(name)
+            meta = cPickle.load(meta_f)
+        for key in meta:
+            if 'label' in key:
+                return meta[key]
+        else:
+            raise RuntimeError("Unexpected branch.")
+
+    def train(self):
+        """
+        Train Reader
+        """
+        return self.__read_batch__(self.__train_filename__)
+
+    def test(self):
+        """
+        Test Reader
+        """
+        return self.__read_batch__(self.__test_filename__)
+
+    def __read_batch__(self, sub_name):
+        with tarfile.open(self.__full_file__, mode='r') as f:
+            names = (each_item.name for each_item in f
+                     if sub_name in each_item.name)
+
+            for name in names:
+                batch = cPickle.load(f.extractfile(name))
+                for item in __read_one_batch__(batch):
+                    yield item
+
+
+class CIFAR10(CIFAR):
+    """
+    CIFAR-10 dataset, images are classified in 10 classes.
+    """
+
+    def __init__(self):
+        super(CIFAR10, self).__init__(
+            CIFAR10_URL,
+            CIFAR10_MD5,
+            meta_filename='batches.meta',
+            train_filename='data_batch',
+            test_filename='test_batch')
+
+
+class CIFAR100(CIFAR):
+    """
+    CIFAR-100 dataset, images are classified in 100 classes.
+    """
+
+    def __init__(self):
+        super(CIFAR100, self).__init__(
+            CIFAR100_URL,
+            CIFAR100_MD5,
+            meta_filename='meta',
+            train_filename='train',
+            test_filename='test')
+
+
+def train_creator():
+    """
+    Default train reader creator. Use CIFAR-10 dataset.
+    """
+    cifar = CIFAR10()
+    return cifar.train
+
+
+def test_creator():
+    """
+    Default test reader creator. Use CIFAR-10 dataset.
+    """
+    cifar = CIFAR10()
+    return cifar.test
+
+
+def unittest(label_count=100):
+    cifar = globals()["CIFAR%d" % label_count]()
+    assert len(cifar.labels()) == label_count
+    for _ in cifar.test():
+        pass
+    for _ in cifar.train():
+        pass
+
+
+if __name__ == '__main__':
+    unittest(10)
+    unittest(100)