2019年4月29日 星期一

CIFAR load and save to picture

import pickle as p
import numpy as np
from PIL import Image
import os

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

'''
        batch file key:
        b'batch_label'
        b'labels'
        b'data'
        b'filenames'
'''
CIFAR10PATH = "E:\\project\\python\\cifar10\\cifar-10-batches-py\\"
DATABATCH = "data_batch_"
OUTPUTFOLDER = ".\\images\\"
#E:\project\python\cifar10\cifar-10-batches-py
if __name__ == "__main__":
    for i in range(1,2):
        batchName = CIFAR10PATH + DATABATCH+str(i)
        print("batch name:"+str(batchName))
        cifar = unpickle(batchName)
        for d in range(10000):
            labelNum = cifar[b'labels'][d]
            imgData = cifar[b'data'][d]
            #remove utf-8 name with b'
            filenames = cifar[b'filenames'][d].decode("utf-8")

            #If the dimension with length 3 is the RGB axis, then you want it at the end.
            #If your array has shape (3, R, C), then np.transpose(img, (1, 2, 0))
            #will have shape (R, C, 3).
            imgData = imgData.reshape(3, 32, 32) #CIFAR default format 3 channels with RxC
            #(0, 1, 2) no change (3, R, C) ==> transpose to (1, 2, 0) => (R, C, 3)
            imgData = imgData.transpose(1, 2, 0)
            img = Image.fromarray(imgData, 'RGB')
            folderName = OUTPUTFOLDER + str(labelNum)
            if not os.path.exists(folderName):
                os.makedirs(folderName)
            img.save(folderName+"\\"+str(filenames).strip()+'.jpg')

沒有留言:

張貼留言