torchvision

torchvision

from torchvision import transforms as T
from torchvision.utils import make_grid, save_image
from torch.utils import data
from torch.utils.data import DataLoader
from PIL import Image
import os

transform = T.Compose([T.Resize(224),
                       T.RandomCrop(224),
                       T.ToTensor(),
                       T.Normalize(mean=[0.5,0.5,0.5], std=[0.5,0.5,0.5])])

class DogCat(data.Dataset):
    def __init__(self, root):
        imgs = os.listdir(root)
        self.imgs = [os.path.join(root, img) for img in imgs]
        self.transform = transform

    def __getitem__(self, idx):
        img_pth = self.imgs[idx]
        img = Image.open(img_pth)
        label = 1 if 'dog' == img_pth.split('\\')[-1].split('.')[0] else 0
        if self.transform:
            img = self.transform(img)
        return img, label

    def __len__(self):
        return len(self.imgs)


if __name__ == '__main__':
    dataset = DogCat(r'D:\\03.Data\\01.CatDog\\image\\train')

    dataloader = DataLoader(dataset, batch_size=13,shuffle=True)
    dataiter = iter(dataloader)

    img, label = next(dataiter)
    T.ToPILImage()(img[0]*0.5+0.5).show()
    print(label)

    # make_grid将多张图片拼接在一个网格中
    img_grid = make_grid(next(dataiter)[0]*0.5+0.5, 4, 4)
    T.ToPILImage()(img_grid).show()
    # save_image将Tensor保存成图片
    print(type(img_grid))  # <class 'torch.Tensor'>
    save_image(img_grid, 'a.png')
    print()