image_output_1.png
from PIL import Image

from torch.utils.data import Dataset, DataLoader

from torchvision import datasets, transforms

import matplotlib.pyplot as plt

import os

import torch

batch_size = 32
latent_dim = 256
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

class DogDataset(Dataset):
    def __init__(self, img_dir, transform1=None, transform2=None):
        self.img_dir = img_dir
        self.img_names = os.listdir(img_dir)
        self.transform1 = transform1
        self.transform2 = transform2
        self.imgs = []
        for img_name in self.img_names:
            img = Image.open(os.path.join(img_dir, img_name))
            if self.transform1 is not None:
                img = self.transform1(img)
            self.imgs.append(img)
    def __getitem__(self, index):
        img = self.imgs[index]
        if self.transform2 is not None:
            img = self.transform2(img)
        return img
    def __len__(self):
        return len(self.imgs)

transform1 = transforms.Compose([transforms.Resize(64), transforms.CenterCrop(64)])
random_transforms = [transforms.RandomRotation(degrees=10)]
transform2 = transforms.Compose([transforms.RandomHorizontalFlip(p=0.5), transforms.RandomApply(random_transforms, p=0.3), transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
train_dataset = DogDataset(img_dir='../input/all-dogs/all-dogs/', transform1=transform1, transform2=transform2)
train_loader = DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True, num_workers=4)

x = next(iter(train_loader))
fig = plt.figure(figsize=(25, 16))
for ii, img in enumerate(x):
    ax = fig.add_subplot(4, 8, ii + 1, xticks=[], yticks=[])
    img = img.numpy().transpose(1, 2, 0)
    plt.imshow(img)