cifar10数据集下载及图片格式解析

全栈程序员站长

发布于 2022-06-25 07:34:20

1.4K00

代码可运行

文章被收录于专栏：全栈程序员必看全栈程序员必看

运行总次数：0

代码可运行

大家好，又见面了，我是你们的朋友全栈君。

CIFAR-10 是由 Hinton 的学生 Alex Krizhevsky 和 Ilya Sutskever 整理的一个用于识别普适物体的小型数据集。一共包含 10 个类别的 RGB 彩色图片：飞机（ a叩lane ）、汽车（ automobile ）、鸟类（ bird ）、猫（ cat ）、鹿（ deer ）、狗（ dog ）、蛙类（ frog ）、马（ horse ）、船（ ship ）和卡车（ truck ）。

图片的尺寸为 32×32×3 ，数据集中一共有 50000 张训练圄片和 10000 张测试图片。CIFAR-10 的图片样例如图所示。

cifar10数据集是一个常用的小型物体数据集，很多模型拿该数据集进行跑流程的简单测试。下面是通过pytorch官方代码自动下载cifar10数据集，并且python解析保存为img格式。

注意解析为img时候，root_dir中原始cifar-10-batches-py文件路径-不识别，需要修改一下。

保存的img路径提前建好文件夹。

import torchvision
from torch.utils.data import DataLoader

import numpy as np
import imageio  # 引入imageio包

train_data = torchvision.datasets.CIFAR10(root="dataset", train=True, transform=torchvision.transforms.ToTensor(), download=True)
test_data = torchvision.datasets.CIFAR10(root="dataset", train=False, transform=torchvision.transforms.ToTensor(), download=True)


# 解压 返回解压后的字典
def unpickle(file):
    import pickle as pk
    fo = open(file, 'rb')
    dict = pk.load(fo, encoding='iso-8859-1')
    fo.close()
    return dict



#begin unpickle
root_dir="./dataset/"
# 生成训练集图片
for j in range(1, 6):
    dataName = root_dir+"/data_batch_" + str(j)  # 读取当前目录下的data_batch1~5文件。
    Xtr = unpickle(dataName)
    print(dataName + " is loading...")

    for i in range(0, 10000):
        img = np.reshape(Xtr['data'][i], (3, 32, 32))  # Xtr['data']为图片二进制数据
        img = img.transpose(1, 2, 0)  # 读取image
        picName = root_dir+'/train/' + str(Xtr['labels'][i]) + '_' + str(i + (j - 1) * 10000) + '.jpg'
        imageio.imsave(picName, img)  # 使用的imageio的imsave类
    print(dataName + " loaded.")

print("test_batch is loading...")

# 生成测试集图片
testXtr = unpickle(root_dir+"/test_batch")
for i in range(0, 10000):
    img = np.reshape(testXtr['data'][i], (3, 32, 32))
    img = img.transpose(1, 2, 0)
    picName = root_dir+'/test/' + str(testXtr['labels'][i]) + '_' + str(i) + '.jpg'
    imageio.imsave(picName, img)
print("test_batch loaded.")