| import os |
| import numpy as np |
| import albumentations |
| from torch.utils.data import Dataset |
|
|
| from taming.data.base import ImagePaths, NumpyPaths, ConcatDatasetWithIndex |
|
|
|
|
| class FacesBase(Dataset): |
| def __init__(self, *args, **kwargs): |
| super().__init__() |
| self.data = None |
| self.keys = None |
|
|
| def __len__(self): |
| return len(self.data) |
|
|
| def __getitem__(self, i): |
| example = self.data[i] |
| ex = {} |
| if self.keys is not None: |
| for k in self.keys: |
| ex[k] = example[k] |
| else: |
| ex = example |
| return ex |
|
|
|
|
| class CelebAHQTrain(FacesBase): |
| def __init__(self, size, keys=None): |
| super().__init__() |
| root = "data/celebahq" |
| with open("data/celebahqtrain.txt", "r") as f: |
| relpaths = f.read().splitlines() |
| paths = [os.path.join(root, relpath) for relpath in relpaths] |
| self.data = NumpyPaths(paths=paths, size=size, random_crop=False) |
| self.keys = keys |
|
|
|
|
| class CelebAHQValidation(FacesBase): |
| def __init__(self, size, keys=None): |
| super().__init__() |
| root = "data/celebahq" |
| with open("data/celebahqvalidation.txt", "r") as f: |
| relpaths = f.read().splitlines() |
| paths = [os.path.join(root, relpath) for relpath in relpaths] |
| self.data = NumpyPaths(paths=paths, size=size, random_crop=False) |
| self.keys = keys |
|
|
|
|
| class FFHQTrain(FacesBase): |
| def __init__(self, size, keys=None): |
| super().__init__() |
| root = "data/ffhq" |
| with open("data/ffhqtrain.txt", "r") as f: |
| relpaths = f.read().splitlines() |
| paths = [os.path.join(root, relpath) for relpath in relpaths] |
| self.data = ImagePaths(paths=paths, size=size, random_crop=False) |
| self.keys = keys |
|
|
|
|
| class FFHQValidation(FacesBase): |
| def __init__(self, size, keys=None): |
| super().__init__() |
| root = "data/ffhq" |
| with open("data/ffhqvalidation.txt", "r") as f: |
| relpaths = f.read().splitlines() |
| paths = [os.path.join(root, relpath) for relpath in relpaths] |
| self.data = ImagePaths(paths=paths, size=size, random_crop=False) |
| self.keys = keys |
|
|
|
|
| class FacesHQTrain(Dataset): |
| |
| def __init__(self, size, keys=None, crop_size=None, coord=False): |
| d1 = CelebAHQTrain(size=size, keys=keys) |
| d2 = FFHQTrain(size=size, keys=keys) |
| self.data = ConcatDatasetWithIndex([d1, d2]) |
| self.coord = coord |
| if crop_size is not None: |
| self.cropper = albumentations.RandomCrop(height=crop_size,width=crop_size) |
| if self.coord: |
| self.cropper = albumentations.Compose([self.cropper], |
| additional_targets={"coord": "image"}) |
|
|
| def __len__(self): |
| return len(self.data) |
|
|
| def __getitem__(self, i): |
| ex, y = self.data[i] |
| if hasattr(self, "cropper"): |
| if not self.coord: |
| out = self.cropper(image=ex["image"]) |
| ex["image"] = out["image"] |
| else: |
| h,w,_ = ex["image"].shape |
| coord = np.arange(h*w).reshape(h,w,1)/(h*w) |
| out = self.cropper(image=ex["image"], coord=coord) |
| ex["image"] = out["image"] |
| ex["coord"] = out["coord"] |
| ex["class"] = y |
| return ex |
|
|
|
|
| class FacesHQValidation(Dataset): |
| |
| def __init__(self, size, keys=None, crop_size=None, coord=False): |
| d1 = CelebAHQValidation(size=size, keys=keys) |
| d2 = FFHQValidation(size=size, keys=keys) |
| self.data = ConcatDatasetWithIndex([d1, d2]) |
| self.coord = coord |
| if crop_size is not None: |
| self.cropper = albumentations.CenterCrop(height=crop_size,width=crop_size) |
| if self.coord: |
| self.cropper = albumentations.Compose([self.cropper], |
| additional_targets={"coord": "image"}) |
|
|
| def __len__(self): |
| return len(self.data) |
|
|
| def __getitem__(self, i): |
| ex, y = self.data[i] |
| if hasattr(self, "cropper"): |
| if not self.coord: |
| out = self.cropper(image=ex["image"]) |
| ex["image"] = out["image"] |
| else: |
| h,w,_ = ex["image"].shape |
| coord = np.arange(h*w).reshape(h,w,1)/(h*w) |
| out = self.cropper(image=ex["image"], coord=coord) |
| ex["image"] = out["image"] |
| ex["coord"] = out["coord"] |
| ex["class"] = y |
| return ex |
|
|
