-
Notifications
You must be signed in to change notification settings - Fork 0
/
mnist.py
42 lines (34 loc) · 1.42 KB
/
mnist.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
#SOURCE: http://g.sweyla.com/blog/2012/mnist-numpy/
#Code given in the source was not working, adjusted a little bit.
import os, struct
from array import array as pyarray
import numpy as np
def load_mnist(dataset="training", digits=np.arange(10), path="."):
"""
Loads MNIST files into 3D numpy arrays
Adapted from: http://abel.ee.ucla.edu/cvxopt/_downloads/mnist.py
"""
if dataset == "training":
fname_img = os.path.join(path, 'train-images-idx3-ubyte')
fname_lbl = os.path.join(path, 'train-labels-idx1-ubyte')
elif dataset == "testing":
fname_img = os.path.join(path, 't10k-images-idx3-ubyte')
fname_lbl = os.path.join(path, 't10k-labels-idx1-ubyte')
else:
raise ValueError("dataset must be 'testing' or 'training'")
flbl = open(fname_lbl, 'rb')
magic_nr, size = struct.unpack(">II", flbl.read(8))
lbl = pyarray("b", flbl.read())
flbl.close()
fimg = open(fname_img, 'rb')
magic_nr, size, rows, cols = struct.unpack(">IIII", fimg.read(16))
img = pyarray("B", fimg.read())
fimg.close()
ind = [ k for k in range(size) if lbl[k] in digits ]
N = len(ind)
images = np.zeros((N, rows, cols), dtype=np.float16)
labels = np.zeros((N, 1), dtype=np.int8)
for i in range(len(ind)):
images[i] = np.array(img[ ind[i]*rows*cols : (ind[i]+1)*rows*cols ]).reshape((rows, cols))
labels[i] = lbl[ind[i]]
return images, labels