pytorch之ImageFolder使用详解

ImageFolder(root, transform=None, target_transform=None, loader=default_loader)

from torchvision import transforms as T
import matplotlib.pyplot as plt
from torchvision.datasets import ImageFolder


dataset = ImageFolder('data/dogcat_2/')

# cat文件夹的图片对应label 0，dog对应1
print(dataset.class_to_idx)

# 所有图片的路径和对应的label
print(dataset.imgs)

# 没有任何的transform，所以返回的还是PIL Image对象
#print(dataset[0][1])# 第一维是第几张图，第二维为1返回label
#print(dataset[0][0]) # 为0返回图片数据
plt.imshow(dataset[0][0])
plt.axis('off')
plt.show()

normalize = T.Normalize(mean=[0.4, 0.4, 0.4], std=[0.2, 0.2, 0.2])
transform = T.Compose([
     T.RandomResizedCrop(224),
     T.RandomHorizontalFlip(),
     T.ToTensor(),
     normalize,
])
dataset = ImageFolder('data1/dogcat_2/', transform=transform)

# 深度学习中图片数据一般保存成CxHxW，即通道数x图片高x图片宽
#print(dataset[0][0].size())

to_img = T.ToPILImage()
# 0.2和0.4是标准差和均值的近似
a=to_img(dataset[0][0]*0.2+0.4)
plt.imshow(a)
plt.axis('off')
plt.show()