stable-diffusion/ldm/data/lsun.py at main · ConAlgorithm/stable-diffusion · GitHub

History

92 lines (72 loc) · 3.2 KB

Raw

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)

stable-diffusion/ldm/data/lsun.py at main · ConAlgorithm/stable-diffusion · GitHub

Expand file tree

Breadcrumbs

lsun.py

Latest commit

History

Breadcrumbs

lsun.py

File metadata and controls