| from torch.utils.data import Dataset |
| import os |
| from PIL import Image |
| from torchvision import transforms |
|
|
|
|
| class ImageFolderDataset(Dataset): |
| def __init__(self, root, transform = None): |
| super(ImageFolderDataset, self).__init__() |
| self.root = root |
| self.transform = transform |
| self.files = list(os.listdir(root)) |
| self.files = [p for p in self.files if p.endswith(('.jpg', '.png', '.jpeg'))] |
|
|
| def __len__(self): |
| return len(self.files) |
|
|
| def __getitem__(self, idx): |
| image_path = os.path.join(self.root, self.files[idx]) |
| image = Image.open(image_path).convert('RGB') |
|
|
| if self.transform: |
| image = self.transform(image) |
|
|
| return image |
|
|
|
|
| def get_transform(size, crop, final_size): |
| transform_list = [] |
| if size > 0: |
| transform_list.append(transforms.Resize(size)) |
| if crop: |
| transform_list.append(transforms.RandomCrop(final_size)) |
| else: |
| transform_list.append(transforms.Resize(final_size)) |
|
|
| transform_list.append(transforms.ToTensor()) |
| return transforms.Compose(transform_list) |
| |
|
|
| def adaptive_instance_normalization(content_feat, style_feat): |
| |
| size = content_feat.size() |
| style_mean, style_std = calc_mean_std(style_feat) |
| content_mean, content_std = calc_mean_std(content_feat) |
| normalized_content_feat = (content_feat - content_mean.expand(size)) / content_std.expand(size) |
| return normalized_content_feat * style_std.expand(size) + style_mean.expand(size) |
|
|
| def calc_mean_std(feat, eps=1e-5): |
| |
| size = feat.size() |
| assert (len(size) == 4) |
| batch_size, channels = size[:2] |
| feat_mean = feat.view(batch_size, channels, -1).mean(dim=2).view(batch_size, channels, 1, 1) |
| feat_var = feat.view(batch_size, channels, -1).var(dim=2, unbiased=False) + eps |
| feat_std = feat_var.sqrt().view(batch_size, channels, 1, 1) |
| return feat_mean, feat_std |
