import os
import numpy as np
from PIL import Image

def unpickle(file):
    import pickle
    with open(file, 'rb') as fo:
        dict = pickle.load(fo, encoding='bytes')
    return dict

def rowToFile(row, label, filename):
    label = str(label)
    if not os.path.exists(label):
        os.mkdir(label)
    dim = 32
    wholeImg = list()
    for i in range(dim):
        oneRow = list()
        for j in range(dim):
            oneRGB = list()
            for k in range(3): # RGB
                pix = row[1024*k + 32*i + j]
                oneRGB.append(pix)
            oneRow.append(oneRGB)
        wholeImg.append(oneRow)
    wholeImg = np.array(wholeImg)
    img = Image.fromarray(wholeImg)
    img.save(os.path.join(label, filename.decode('ascii')))
    img.close()

trainFiles = ['data_batch_1', 'data_batch_2', 'data_batch_3', 'data_batch_4', 'data_batch_5']
testFiles = ['test_batch']

files = testFiles

for oneFile in files:
    oneBatch = unpickle(oneFile)
    labels = oneBatch[b'labels']
    data = oneBatch[b'data']
    filenames = oneBatch[b'filenames']
    for i in range(len(labels)):
        rowToFile(data[i], labels[i], filenames[i])
    print(oneFile + " is finished")

print("Everything is finished")