BitDance_B_4x/transformer/src/dataset.py

import contextlib
import io
import math
import os
import pickle
import tarfile
from functools import lru_cache

import numpy as np
import torch
from PIL import Image
from torch.utils.data import Dataset
from torchvision.datasets import ImageFolder
import torchvision.datasets as datasets


@contextlib.contextmanager
def numpy_seed(seed, *addl_seeds):
    """Context manager which seeds the NumPy PRNG with the specified seed and
    restores the state afterward"""
    if seed is None:
        yield
        return

    def check_seed(s):
        assert type(s) == int or type(s) == np.int32 or type(s) == np.int64

    check_seed(seed)
    if len(addl_seeds) > 0:
        for s in addl_seeds:
            check_seed(s)
        seed = int(hash((seed, *addl_seeds)) % 1e8)
    state = np.random.get_state()
    np.random.seed(seed)
    try:
        yield
    finally:
        np.random.set_state(state)


def build_flat_index(outer_path: str, idx_path: str):
    if os.path.exists(idx_path):
        print(f"Index file {idx_path} already exists. Skipping index building.")
        return pickle.load(open(idx_path, "rb"))
    entries = []  # (offset, size, label)
    cats = set()
    idx = 0
    with tarfile.open(outer_path, "r:") as outer:
        for sub in outer.getmembers():
            if not sub.isfile() or not sub.name.endswith(".tar"):
                continue
            outer_off = sub.offset_data
            sub_fobj = outer.extractfile(sub)
            with tarfile.open(fileobj=sub_fobj, mode="r:") as inner:
                for m in inner.getmembers():
                    if not m.isfile():
                        continue
                    cat = m.name.split("_", 1)[0]
                    cats.add(cat)
                    abs_off = outer_off + m.offset_data
                    entries.append((abs_off, m.size, cat))
                    if idx % 1000 == 1:
                        print(idx, m.name, abs_off, m.size, cat)
                    idx += 1
    sorted_cats = sorted(cats)
    cat2idx = {c: i for i, c in enumerate(sorted_cats)}

    flat = [(off, size, cat2idx[c]) for off, size, c in entries]

    os.makedirs(os.path.dirname(idx_path), exist_ok=True)
    with open(idx_path, "wb") as f:
        pickle.dump(
            flat,
            f,
        )
    print(f"Built flat index with {len(flat)} images.")
    return flat


class ImageNetTarDataset(Dataset):
    """
    ImageNet dataset stored in a tar file, avoid to decompress the whole dataset.
    You can direct use the original downloaded tar file (ILSVRC2012_img_train.tar) from official ImageNet website.
    The best practice is to copy the tar file to node's local disk or ramdisk (like /dev/shm/) first, to avoid remote I/O bottleneck.
    """

    def __init__(
        self,
        tar_file,
    ):
        self.tar_file = tar_file
        self.tar_handle = None
        self.files = build_flat_index(tar_file, tar_file + ".index")
        self.num_examples = len(self.files)

    def __len__(self):
        return self.num_examples

    def get_raw_image(self, index):
        if self.tar_handle is None:
            self.tar_handle = open(self.tar_file, "rb")

        offset, size, label = self.files[index]
        self.tar_handle.seek(offset)
        data = self.tar_handle.read(size)
        image = Image.open(io.BytesIO(data)).convert("RGB")
        return image, label

    @lru_cache(maxsize=16)
    def __getitem__(self, idx):
        return self.get_raw_image(idx)


def center_crop_arr(pil_image, image_size):
    """
    Center cropping implementation from ADM.
    https://github.com/openai/guided-diffusion/blob/8fb3ad9197f16bbc40620447b2742e13458d2831/guided_diffusion/image_datasets.py#L126
    """
    while min(*pil_image.size) >= 2 * image_size:
        pil_image = pil_image.resize(
            tuple(x // 2 for x in pil_image.size), resample=Image.BOX
        )

    scale = image_size / min(*pil_image.size)
    pil_image = pil_image.resize(
        tuple(round(x * scale) for x in pil_image.size), resample=Image.BICUBIC
    )

    arr = np.array(pil_image)
    crop_y = (arr.shape[0] - image_size) // 2
    crop_x = (arr.shape[1] - image_size) // 2
    return Image.fromarray(
        arr[crop_y : crop_y + image_size, crop_x : crop_x + image_size]
    )


def numpy_randrange(start, end):
    return int(np.random.randint(start, end))


def random_crop_arr(pil_image, image_size, min_crop_frac=0.8, max_crop_frac=1.0):
    min_smaller_dim_size = math.ceil(image_size / max_crop_frac)
    max_smaller_dim_size = math.ceil(image_size / min_crop_frac)
    smaller_dim_size = numpy_randrange(min_smaller_dim_size, max_smaller_dim_size + 1)

    # We are not on a new enough PIL to support the `reducing_gap`
    # argument, which uses BOX downsampling at powers of two first.
    # Thus, we do it by hand to improve downsample quality.
    while min(*pil_image.size) >= 2 * smaller_dim_size:
        pil_image = pil_image.resize(
            tuple(x // 2 for x in pil_image.size), resample=Image.BOX
        )

    scale = smaller_dim_size / min(*pil_image.size)
    pil_image = pil_image.resize(
        tuple(round(x * scale) for x in pil_image.size), resample=Image.BICUBIC
    )

    arr = np.array(pil_image)
    crop_y = numpy_randrange(0, arr.shape[0] - image_size + 1)
    crop_x = numpy_randrange(0, arr.shape[1] - image_size + 1)
    return Image.fromarray(
        arr[crop_y : crop_y + image_size, crop_x : crop_x + image_size]
    )


def crop(pil_image, left, top, right, bottom):
    """
    Crop the image to the specified box.
    """
    return pil_image.crop((left, top, right, bottom))


class ImageCropDataset(Dataset):

    def __init__(
        self,
        raw_dataset,
        resolution,
        patch_size,
        seed=42,
    ):
        self.raw_dataset = raw_dataset
        self.resolution = resolution
        self.patch_size = patch_size
        self.aug_ratio = 1.0
        self.seed = seed
        self.epoch = None

    def set_epoch(self, epoch):
        self.epoch = epoch

    def set_aug_ratio(self, aug_ratio):
        self.aug_ratio = aug_ratio

    def __len__(self):
        return len(self.raw_dataset)

    def crop_and_flip(self, image):
        is_aug = np.random.rand() < self.aug_ratio
        if not is_aug:
            image = center_crop_arr(image, self.resolution)
        else:
            image = random_crop_arr(image, self.resolution)

        arr = np.asarray(image)

        is_flip = int(np.random.randint(0, 2))
        if is_flip == 1:
            # horizontal flip
            arr = arr[:, ::-1, :]

        return arr.transpose(2, 0, 1)  # HWC to CHW

    def __getitem__(self, idx):
        with numpy_seed(self.seed, self.epoch, idx):
            image, label = self.raw_dataset[idx]
            samples = self.crop_and_flip(image)
            # to [-1, 1]
            samples = (samples.astype(np.float32) / 255.0 - 0.5) * 2.0
            samples = torch.from_numpy(samples).float()
            return (
                samples,
                torch.tensor(label).long(),
            )


def build_dataset(args):
    # use tarred imagenet dataset if data_path ends with .tar
    raw_dataset = (
        ImageNetTarDataset(args.data_path)
        if args.data_path.endswith(".tar")
        else ImageFolder(args.data_path)
    )
    return ImageCropDataset(
        raw_dataset,
        args.image_size,
        args.patch_size,
        seed=args.global_seed if hasattr(args, "global_seed") else 42,
    )