| import os |
| import numpy as np |
| import PIL |
| from PIL import Image |
| from torch.utils.data import Dataset |
| from torchvision import transforms |
|
|
|
|
| class LSUNBase(Dataset): |
| def __init__(self, |
| txt_file, |
| data_root, |
| size=None, |
| interpolation="bicubic", |
| flip_p=0.5 |
| ): |
| self.data_paths = txt_file |
| self.data_root = data_root |
| with open(self.data_paths, "r") as f: |
| self.image_paths = f.read().splitlines() |
| self._length = len(self.image_paths) |
| self.labels = { |
| "relative_file_path_": [l for l in self.image_paths], |
| "file_path_": [os.path.join(self.data_root, l) |
| for l in self.image_paths], |
| } |
|
|
| self.size = size |
| self.interpolation = {"linear": PIL.Image.LINEAR, |
| "bilinear": PIL.Image.BILINEAR, |
| "bicubic": PIL.Image.BICUBIC, |
| "lanczos": PIL.Image.LANCZOS, |
| }[interpolation] |
| self.flip = transforms.RandomHorizontalFlip(p=flip_p) |
|
|
| def __len__(self): |
| return self._length |
|
|
| def __getitem__(self, i): |
| example = dict((k, self.labels[k][i]) for k in self.labels) |
| image = Image.open(example["file_path_"]) |
| if not image.mode == "RGB": |
| image = image.convert("RGB") |
|
|
| |
| img = np.array(image).astype(np.uint8) |
| crop = min(img.shape[0], img.shape[1]) |
| h, w, = img.shape[0], img.shape[1] |
| img = img[(h - crop) // 2:(h + crop) // 2, |
| (w - crop) // 2:(w + crop) // 2] |
|
|
| image = Image.fromarray(img) |
| if self.size is not None: |
| image = image.resize((self.size, self.size), resample=self.interpolation) |
|
|
| image = self.flip(image) |
| image = np.array(image).astype(np.uint8) |
| example["image"] = (image / 127.5 - 1.0).astype(np.float32) |
| return example |
|
|
|
|
| class LSUNChurchesTrain(LSUNBase): |
| def __init__(self, **kwargs): |
| super().__init__(txt_file="data/lsun/church_outdoor_train.txt", data_root="data/lsun/churches", **kwargs) |
|
|
|
|
| class LSUNChurchesValidation(LSUNBase): |
| def __init__(self, flip_p=0., **kwargs): |
| super().__init__(txt_file="data/lsun/church_outdoor_val.txt", data_root="data/lsun/churches", |
| flip_p=flip_p, **kwargs) |
|
|
|
|
| class LSUNBedroomsTrain(LSUNBase): |
| def __init__(self, **kwargs): |
| super().__init__(txt_file="data/lsun/bedrooms_train.txt", data_root="data/lsun/bedrooms", **kwargs) |
|
|
|
|
| class LSUNBedroomsValidation(LSUNBase): |
| def __init__(self, flip_p=0.0, **kwargs): |
| super().__init__(txt_file="data/lsun/bedrooms_val.txt", data_root="data/lsun/bedrooms", |
| flip_p=flip_p, **kwargs) |
|
|
|
|
| class LSUNCatsTrain(LSUNBase): |
| def __init__(self, **kwargs): |
| super().__init__(txt_file="data/lsun/cat_train.txt", data_root="data/lsun/cats", **kwargs) |
|
|
|
|
| class LSUNCatsValidation(LSUNBase): |
| def __init__(self, flip_p=0., **kwargs): |
| super().__init__(txt_file="data/lsun/cat_val.txt", data_root="data/lsun/cats", |
| flip_p=flip_p, **kwargs) |
|
|
