AdaIN_dataset.py

import os
import glob
import numpy as np
from tqdm import tqdm
import torch
from torch.utils.data import Dataset
from torchvision import transforms
from PIL import Image

normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                 std=[0.229, 0.224, 0.225])

def denorm(tensor, device):
    std = torch.Tensor([0.229, 0.224, 0.225]).reshape(-1, 1, 1).to(device)
    mean = torch.Tensor([0.485, 0.456, 0.406]).reshape(-1, 1, 1).to(device)
    res = torch.clamp(tensor * std + mean, 0, 1)
    return res

class PreprocessDataset(Dataset):
    def __init__(self, content_dataframe, style_dataframe, transform):
        self.content_dataframe = content_dataframe
        self.style_dataframe = style_dataframe

        self.transform = transform
    
    @staticmethod
    def _resize(image):
        H, W = image.size
        if H < W:
            ratio = W / H
            H = 512
            W = int(ratio * H)
        else:
            ratio = H / W
            W = 512
            H = int(ratio * W)
        img = image.resize((H, W), Image.ANTIALIAS)
        return image
    
    def __len__(self):
        return len(self.style_dataframe)
        #return len(self.images_pairs)

    def __getitem__(self, index):
        content_row = self.content_dataframe.iloc[index]
        style_row = self.style_dataframe.iloc[index]
        
        content_image = self._resize(Image.open((content_row["file_path"])).convert('RGB'))
        style_image = self._resize(Image.open((style_row["file_path"])).convert('RGB'))        
        content_image = self.transform(content_image)
        style_image = self.transform(style_image)
        
        return content_image, style_image