# 数据集导入（公开数据集）
# 1.导入依赖包
import torchvision

# 2.获取数据集
# 2.1 训练集
train_data = torchvision.datasets.CIFAR10(root='./dataset', train=True, download=True)
test_data = torchvision.datasets.CIFAR10(root='./dataset', train=False, download=True)
# 2.2 查看数据量
# print(f'训练集数据量：{len(train_data)}')
# print(f'测试集数据量：{len(test_data)}')
# print(f'类别：{train_data.class_to_idx}')

# 数据集导入（自定义数据集）
# 1.导入依赖包
from torchvision.datasets import ImageFolder
from torchvision.transforms import ToTensor, Resize, Compose  # 组合

# 2.读取数据集（数据路径 +数据增强）
train_data = ImageFolder(root='dataset/flower_datas/train', transform=Compose([ToTensor(),
                                                                               Resize((224, 224))]))  # H W
test_data = ImageFolder(root='dataset/flower_datas/val', transform=Compose([ToTensor(),
                                                                               Resize((224, 224))]))
# 查看数据情况
print(f'训练集数据量：{len(train_data)}')
print(f'测试集数据量：{len(test_data)}')
print(f'类别：{train_data.class_to_idx}')