Anonymous-123/ImageNet-Editing

1import math2import random3 4from PIL import Image5import blobfile as bf6from mpi4py import MPI7import numpy as np8from torch.utils.data import DataLoader, Dataset9 10 11def load_data(12    *,13    data_dir,14    batch_size,15    image_size,16    class_cond=False,17    deterministic=False,18    random_crop=False,19    random_flip=True,20):21    """22    For a dataset, create a generator over (images, kwargs) pairs.23 24    Each images is an NCHW float tensor, and the kwargs dict contains zero or25    more keys, each of which map to a batched Tensor of their own.26    The kwargs dict can be used for class labels, in which case the key is "y"27    and the values are integer tensors of class labels.28 29    :param data_dir: a dataset directory.30    :param batch_size: the batch size of each returned pair.31    :param image_size: the size to which images are resized.32    :param class_cond: if True, include a "y" key in returned dicts for class33                       label. If classes are not available and this is true, an34                       exception will be raised.35    :param deterministic: if True, yield results in a deterministic order.36    :param random_crop: if True, randomly crop the images for augmentation.37    :param random_flip: if True, randomly flip the images for augmentation.38    """39    if not data_dir:40        raise ValueError("unspecified data directory")41    all_files = _list_image_files_recursively(data_dir)42    classes = None43    if class_cond:44        # Assume classes are the first part of the filename,45        # before an underscore.46        class_names = [bf.basename(path).split("_")[0] for path in all_files]47        sorted_classes = {x: i for i, x in enumerate(sorted(set(class_names)))}48        classes = [sorted_classes[x] for x in class_names]49    dataset = ImageDataset(50        image_size,51        all_files,52        classes=classes,53        shard=MPI.COMM_WORLD.Get_rank(),54        num_shards=MPI.COMM_WORLD.Get_size(),55        random_crop=random_crop,56        random_flip=random_flip,57    )58    if deterministic:59        loader = DataLoader(60            dataset, batch_size=batch_size, shuffle=False, num_workers=1, drop_last=True61        )62    else:63        loader = DataLoader(64            dataset, batch_size=batch_size, shuffle=True, num_workers=1, drop_last=True65        )66    while True:67        yield from loader68 69 70def _list_image_files_recursively(data_dir):71    results = []72    for entry in sorted(bf.listdir(data_dir)):73        full_path = bf.join(data_dir, entry)74        ext = entry.split(".")[-1]75        if "." in entry and ext.lower() in ["jpg", "jpeg", "png", "gif"]:76            results.append(full_path)77        elif bf.isdir(full_path):78            results.extend(_list_image_files_recursively(full_path))79    return results80 81 82class ImageDataset(Dataset):83    def __init__(84        self,85        resolution,86        image_paths,87        classes=None,88        shard=0,89        num_shards=1,90        random_crop=False,91        random_flip=True,92    ):93        super().__init__()94        self.resolution = resolution95        self.local_images = image_paths[shard:][::num_shards]96        self.local_classes = None if classes is None else classes[shard:][::num_shards]97        self.random_crop = random_crop98        self.random_flip = random_flip99 100    def __len__(self):101        return len(self.local_images)102 103    def __getitem__(self, idx):104        path = self.local_images[idx]105        with bf.BlobFile(path, "rb") as f:106            pil_image = Image.open(f)107            pil_image.load()108        pil_image = pil_image.convert("RGB")109 110        if self.random_crop:111            arr = random_crop_arr(pil_image, self.resolution)112        else:113            arr = center_crop_arr(pil_image, self.resolution)114 115        if self.random_flip and random.random() < 0.5:116            arr = arr[:, ::-1]117 118        arr = arr.astype(np.float32) / 127.5 - 1119 120        out_dict = {}121        if self.local_classes is not None:122            out_dict["y"] = np.array(self.local_classes[idx], dtype=np.int64)123        return np.transpose(arr, [2, 0, 1]), out_dict124 125 126def center_crop_arr(pil_image, image_size):127    # We are not on a new enough PIL to support the `reducing_gap`128    # argument, which uses BOX downsampling at powers of two first.129    # Thus, we do it by hand to improve downsample quality.130    while min(*pil_image.size) >= 2 * image_size:131        pil_image = pil_image.resize(132            tuple(x // 2 for x in pil_image.size), resample=Image.BOX133        )134 135    scale = image_size / min(*pil_image.size)136    pil_image = pil_image.resize(137        tuple(round(x * scale) for x in pil_image.size), resample=Image.BICUBIC138    )139 140    arr = np.array(pil_image)141    crop_y = (arr.shape[0] - image_size) // 2142    crop_x = (arr.shape[1] - image_size) // 2143    return arr[crop_y : crop_y + image_size, crop_x : crop_x + image_size]144 145 146def random_crop_arr(pil_image, image_size, min_crop_frac=0.8, max_crop_frac=1.0):147    min_smaller_dim_size = math.ceil(image_size / max_crop_frac)148    max_smaller_dim_size = math.ceil(image_size / min_crop_frac)149    smaller_dim_size = random.randrange(min_smaller_dim_size, max_smaller_dim_size + 1)150 151    # We are not on a new enough PIL to support the `reducing_gap`152    # argument, which uses BOX downsampling at powers of two first.153    # Thus, we do it by hand to improve downsample quality.154    while min(*pil_image.size) >= 2 * smaller_dim_size:155        pil_image = pil_image.resize(156            tuple(x // 2 for x in pil_image.size), resample=Image.BOX157        )158 159    scale = smaller_dim_size / min(*pil_image.size)160    pil_image = pil_image.resize(161        tuple(round(x * scale) for x in pil_image.size), resample=Image.BICUBIC162    )163 164    arr = np.array(pil_image)165    crop_y = random.randrange(arr.shape[0] - image_size + 1)166    crop_x = random.randrange(arr.shape[1] - image_size + 1)167    return arr[crop_y : crop_y + image_size, crop_x : crop_x + image_size]168