import numpy
from PIL import Image


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

def picProcessor(pic):
    processed_pic=numpy.zeros((32,32,3), dtype=numpy.uint8)

    #colors are stored continously
    step=0
    for i in range(32):
        for j in range(32):
            processed_pic[i][j][0]=pic[step*32+j]
            processed_pic[i][j][1]=pic[step*32+j+1024]
            processed_pic[i][j][2]=pic[step*32+j+2048]
        step=step+1

    return processed_pic

def generatePicture(file,pic):
    img=Image.fromarray(pic,'RGB')
    img.save(file)



def main():
    #randomly opening a batch
    batch_no=numpy.random.randint(5)+1
    batch_location="D:\Programok\CIFAR\cifar-10-python\cifar-10-batches-py\data_batch_"+str(batch_no)
    batch=unpickle(batch_location)

    #randomly getting a picture from the randomly chosen batch
    picnum=numpy.random.randint(0,10000)


    test_pic=picProcessor(batch[b'data'][picnum])
    generatePicture('teszt.png',test_pic)
    return 0

main()