dataloader.py

# -*- coding: utf-8 -*-
from __future__ import unicode_literals
"""
    This downloading script is modified from:
    https://github.com/sorki/python-mnist

"""

import os
import struct
from array import array
import numpy as np

class MNIST(object):
    def __init__(self, path='.'):
        self.path = path

        self.test_img_fname = 't10k-images.idx3-ubyte'
        self.test_lbl_fname = 't10k-labels.idx1-ubyte'

        self.train_img_fname = 'train-images.idx3-ubyte'
        self.train_lbl_fname = 'train-labels.idx1-ubyte'

        self.test_images = []
        self.test_labels = []

        self.train_images = []
        self.train_labels = []

    def load_testing(self):
        ims, labels = self.load(os.path.join(self.path, self.test_img_fname),
                                os.path.join(self.path, self.test_lbl_fname))

        self.test_images = np.array(ims)
        self.test_labels = np.array(labels)

        return self.test_images, self.test_labels

    def load_training(self):
        ims, labels = self.load(os.path.join(self.path, self.train_img_fname),
                                os.path.join(self.path, self.train_lbl_fname))

        self.train_images = np.array(ims)
        self.train_labels = np.array(labels)

        return self.train_images, self.train_labels

    @classmethod
    def load(cls, path_img, path_lbl):
        with open(path_lbl, 'rb') as file:
            magic, size = struct.unpack(">II", file.read(8))
            if magic != 2049:
                raise ValueError('Magic number mismatch, expected 2049,'
                                 'got {}'.format(magic))

            labels = array("B", file.read())

        with open(path_img, 'rb') as file:
            magic, size, rows, cols = struct.unpack(">IIII", file.read(16))
            if magic != 2051:
                raise ValueError('Magic number mismatch, expected 2051,'
                                 'got {}'.format(magic))

            image_data = array("B", file.read())

        images = []
        for i in range(size):
            images.append([0] * rows * cols)

        for i in range(size):
            images[i][:] = image_data[i * rows * cols:(i + 1) * rows * cols]

        return images, labels

    @classmethod
    def display(cls, img, width=28, threshold=200):
        render = ''
        for i in range(len(img)):
            if i % width == 0:
                render += '\n'
            if img[i] > threshold:
                render += '@'
            else:
                render += '.'
        return render


def load_data():
    d = np.load('data/data.npz')
    return d['x_train'], d['y_train'], d['x_test'], d['y_test']


if __name__ == '__main__':
    if not os.path.exists('data/data.npz'):
        data = MNIST('./data')
        data.load_training()
        data.load_testing()
        print("save data in data/data.npz!")
        np.savez('data/data.npz', x_train=data.train_images, y_train=data.train_labels,
                 x_test=data.test_images, y_test=data.test_labels)
    x_train, y_train, x_test, y_test = load_data()
    print x_test
    print 'success!'