"""
    使用torchvision提供的数据集
"""
from PIL import Image
from torchvision import transforms
import torchvision
from torch.utils.tensorboard import SummaryWriter

# 构建一个Compose完成指定处理流程
dataset_transform = transforms.Compose([
    transforms.ToTensor()
])
writer = SummaryWriter("logs/006")

# 下载数据集
# https://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz
train_set = torchvision.datasets.CIFAR10(root="./visionData", train=True, transform=dataset_transform, download=True)
test_set = torchvision.datasets.CIFAR10(root="./visionData", train=False, transform=dataset_transform, download=True)

print(test_set[0])  # 打印第一个数据
print(test_set.classes)  # 打印数据集的标签

for i in range(10):
    # target表示该图像是标签集中的第几个类别
    img, target = test_set[i]
    print(f'第 {i} 个图像的类别是 {test_set.classes[target]}')
    writer.add_image("逐个记录图像", img, i)

writer.close()