AIStoryBoard/python/PaddleOCR-main/test_tipc/supplementary/load_cifar.py

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


def load_CIFAR_batch(filename):
    """load single batch of cifar"""
    with open(filename, "rb") as f:
        datadict = p.load(f, encoding="bytes")
        # 以字典的形式取出数据
        X = datadict[b"data"]
        Y = datadict[b"fine_labels"]
        try:
            X = X.reshape(10000, 3, 32, 32)
        except:
            X = X.reshape(50000, 3, 32, 32)
        Y = np.array(Y)
        print(Y.shape)
        return X, Y


if __name__ == "__main__":
    mode = "train"
    imgX, imgY = load_CIFAR_batch(f"./cifar-100-python/{mode}")
    with open(f"./cifar-100-python/{mode}_imgs/img_label.txt", "a+") as f:
        for i in range(imgY.shape[0]):
            f.write("img" + str(i) + " " + str(imgY[i]) + "\n")

    for i in range(imgX.shape[0]):
        imgs = imgX[i]
        img0 = imgs[0]
        img1 = imgs[1]
        img2 = imgs[2]
        i0 = Image.fromarray(img0)
        i1 = Image.fromarray(img1)
        i2 = Image.fromarray(img2)
        img = Image.merge("RGB", (i0, i1, i2))
        name = "img" + str(i) + ".png"
        img.save(f"./cifar-100-python/{mode}_imgs/" + name, "png")
    print("save successfully!")