data_mnist.py

import array
import gzip
import os
import struct
from urllib.request import urlretrieve

import numpy as np


def download(url, filename):
    if not os.path.exists("data"):
        os.makedirs("data")
    out_file = os.path.join("data", filename)
    if not os.path.isfile(out_file):
        urlretrieve(url, out_file)


def mnist():
    base_url = "http://yann.lecun.com/exdb/mnist/"

    def parse_labels(filename):
        with gzip.open(filename, "rb") as fh:
            magic, num_data = struct.unpack(">II", fh.read(8))
            return np.array(array.array("B", fh.read()), dtype=np.uint8)

    def parse_images(filename):
        with gzip.open(filename, "rb") as fh:
            magic, num_data, rows, cols = struct.unpack(">IIII", fh.read(16))
            return np.array(array.array("B", fh.read()), dtype=np.uint8).reshape(num_data, rows, cols)

    for filename in [
        "train-images-idx3-ubyte.gz",
        "train-labels-idx1-ubyte.gz",
        "t10k-images-idx3-ubyte.gz",
        "t10k-labels-idx1-ubyte.gz",
    ]:
        download(base_url + filename, filename)

    train_images = parse_images("data/train-images-idx3-ubyte.gz")
    train_labels = parse_labels("data/train-labels-idx1-ubyte.gz")
    test_images = parse_images("data/t10k-images-idx3-ubyte.gz")
    test_labels = parse_labels("data/t10k-labels-idx1-ubyte.gz")

    return train_images, train_labels, test_images, test_labels