Spaces:
Paused
Paused
| 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") | |
| # default to score-sde preprocessing | |
| 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) | |