提交 976a6982 编写于 作者: Y Yu Yang

Add cifar dataset

上级 befc3e06
"""
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)
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册
新手
引导
客服 返回
顶部