# coding=utf-8

# coding=utf-8
import struct
import numpy as np
#import pylab as pl


testDataSetPath = "t10k-images.idx3-ubyte"
testLabelSetPath = "t10k-labels.idx1-ubyte"
trainDataSetPath = "train-images.idx3-ubyte"
trainLabelSetPath = "train-labels.idx1-ubyte"

 # 读取一张图片 (numRows * numCols byte) 28 * 28 = 784 byte
def loadImage(buf, index, numRows, numCols):
    #  '>784B'的意思就是用大端法读取784个unsigned byte
    fmt = ">" + str(numRows * numCols) + "B"
    im = struct.unpack_from(fmt, buf, index)
    index += struct.calcsize(fmt)
    im = np.array(im).reshape((numRows, numCols))
    #pl.imshow(im)
    #pl.show()
    return im, index

def loadDataSet(path):
    binFile = open(path, 'rb')
    buf = binFile.read()

    # '>IIII'是说使用大端法读取4个unsinged int32
    index = 0
    magic, numImages, numRows, numCols = struct.unpack_from(">IIII", buf, index)
    index += struct.calcsize(">IIII")

    images = []
    for i in range(numImages):
        im, index = loadImage(buf, index, numRows, numCols)
        images.append(im)
    return images

def loadLabelSet(path):
    buf = open(path,'rb').read()
    index = 0

    magic, num = struct.unpack_from('>II', buf, index)
    index += struct.calcsize('>II')

    labels = []
    for i in range(num):
        label, = struct.unpack_from('>B', buf, index)       # 返回的是tuple, 后面加逗号表明只需要第一个元素
        index += struct.calcsize('>B')
        labels.append(label)

    return labels

trainDataSet = loadDataSet(trainDataSetPath)
trainLabelSe = loadLabelSet(trainLabelSetPath)
testDataSet = loadDataSet(testDataSetPath)
testLabelSet = loadLabelSet(testLabelSetPath)