datasets.py

# Copyright (c) 2021 PPViT Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
Dataset related classes and methods for ViT training and validation
Cifar10, Cifar100 and ImageNet2012 are supported
"""

import os
import math
import random
import PIL
import paddle
import paddle.nn as nn
from paddle.io import Dataset, DataLoader, DistributedBatchSampler
from paddle.vision import transforms, datasets, image_load

class ImageNet2012Dataset(Dataset):
    """Build ImageNet2012 dataset

    This class gets train/val imagenet datasets, which loads transfomed data and labels.

    Attributes:
        file_folder: path where imagenet images are stored
        transform: preprocessing ops to apply on image
        img_path_list: list of full path of images in whole dataset
        label_list: list of labels of whole dataset
    """

    def __init__(self, file_folder, mode="train", transform=None):
        """Init ImageNet2012 Dataset with dataset file path, mode(train/val), and transform"""
        super(ImageNet2012Dataset, self).__init__()
        assert mode in ["train", "val"]
        self.file_folder = file_folder
        self.transform = transform
        self.img_path_list = []
        self.label_list = []

        if mode == "train":
            self.list_file = os.path.join(self.file_folder, "train_list.txt")
        else:
            self.list_file = os.path.join(self.file_folder, "val_list.txt")

        with open(self.list_file, 'r') as infile:
            for line in infile:
                img_path = line.strip().split()[0]
                img_label = int(line.strip().split()[1])
                self.img_path_list.append(os.path.join(self.file_folder, img_path))
                self.label_list.append(img_label)
        print(f'----- Imagenet2012 image {mode} list len = {len(self.label_list)}')

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

    def __getitem__(self, index):
        data = image_load(self.img_path_list[index]).convert('RGB')
        if isinstance(self.transform, (list, tuple)):
            data = [trans(data) for trans in self.transform]
        else:
            data = self.transform(data)
        label = self.label_list[index]

        return data, label


def get_train_dino_transforms(config):
    flip_and_color_jitter = transforms.Compose([
        transforms.RandomHorizontalFlip(prob=0.5),
        RandomApply([transforms.ColorJitter(brightness=0.4,
                                            contrast=0.4,
                                            saturation=0.2,
                                            hue=0.1)], prob=0.8),
        RandomGrayscale(prob=0.2),
    ])

    normalize = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize(mean=config.DATA.IMAGENET_MEAN, std=config.DATA.IMAGENET_STD),
    ])
    # 1st global crop
    global_transforms1 = transforms.Compose([
        transforms.RandomResizedCrop((config.DATA.IMAGE_SIZE, config.DATA.IMAGE_SIZE),
                                      scale=(config.DATA.GLOBAL_CROPS_SCALE),
                                      interpolation='bicubic'),
        flip_and_color_jitter,
        GaussianBlur(prob=1.0),
        normalize,
    ])
    # 2nd global crop
    global_transforms2 = transforms.Compose([
        transforms.RandomResizedCrop((config.DATA.IMAGE_SIZE, config.DATA.IMAGE_SIZE),
                                      scale=(config.DATA.GLOBAL_CROPS_SCALE),
                                      interpolation='bicubic'),
        flip_and_color_jitter,
        GaussianBlur(prob=0.1),
        Solarization(prob=0.2),
        normalize,
    ])
    # local small crops
    local_transforms = transforms.Compose([
        transforms.RandomResizedCrop(
            (config.DATA.SMALL_CROP_IMAGE_SIZE, config.DATA.SMALL_CROP_IMAGE_SIZE),
            scale=(config.DATA.LOCAL_CROPS_SCALE),
            interpolation='bicubic'),
        flip_and_color_jitter,
        GaussianBlur(prob=0.5),
        normalize,
    ])

    transforms_list = [global_transforms1, global_transforms2]
    for _ in range(config.DATA.LOCAL_CROPS_NUMBER):
        transforms_list.append(local_transforms)
    return transforms_list


class GaussianBlur():
    def __init__(self, prob=0.5, radius_min=0.1, radius_max=2.):
        self.prob = prob
        self.radius_min = radius_min
        self.radius_max = radius_max

    def __call__(self, img):
        if random.random() < self.prob:
            return img
        return img.filter(
            PIL.ImageFilter.GaussianBlur(
                radius=random.uniform(self.radius_min, self.radius_max)))


class Solarization():
    def __init__(self, prob):
        self.prob = prob

    def __call__(self, img):
        if random.random() < self.prob:
            return PIL.ImageOps.solarize(img)
        return img


class RandomGrayscale(nn.Layer):
    def __init__(self, prob=0.1):
        super().__init__()
        self.prob = prob

    def get_image_channels(self, img):
        if isinstance(img, paddle.Tensor):
            return img.shape[1] 
        return len(PIL.Image.Image.getbands(img))

    def forward(self, img):
        if self.prob > paddle.rand([1]):
            ch = self.get_image_channels(img)
            return transforms.to_grayscale(img, num_output_channels=self.get_image_channels(img))
        return img


class RandomApply(nn.Layer):
    def __init__(self, transforms, prob=0.5):
        super().__init__()
        self.prob = prob
        self.transforms = transforms

    def forward(self, img):
        if self.prob > paddle.rand([1]):
            return img
        for t in self.transforms:
            img = t(img)
        return img


def get_train_transforms(config):
    """ Get training transforms

    For training, a RandomResizedCrop is applied, then normalization is applied with
    [0.5, 0.5, 0.5] mean and std. The input pixel values must be rescaled to [0, 1.]
    Outputs is converted to tensor

    Args:
        config: configs contains IMAGE_SIZE, see config.py for details
    Returns:
        transforms_train: training transforms
    """

    transforms_train = transforms.Compose([
        transforms.RandomResizedCrop((config.DATA.IMAGE_SIZE, config.DATA.IMAGE_SIZE),
                                     scale=(0.05, 1.0)),
        transforms.ToTensor(),
        transforms.Normalize(mean=config.DATA.IMAGENET_MEAN, std=config.DATA.IMAGENET_STD),
    ])
    return transforms_train


def get_val_transforms(config):
    """ Get training transforms

    For validation, image is first Resize then CenterCrop to image_size.
    Then normalization is applied with [0.5, 0.5, 0.5] mean and std.
    The input pixel values must be rescaled to [0, 1.]
    Outputs is converted to tensor

    Args:
        config: configs contains IMAGE_SIZE, see config.py for details
    Returns:
        transforms_train: training transforms
    """

    scale_size = int(math.floor(config.DATA.IMAGE_SIZE / config.DATA.CROP_PCT))
    transforms_val = transforms.Compose([
        transforms.Resize(scale_size, 'bicubic'), # single int for resize shorter side of image
        transforms.CenterCrop((config.DATA.IMAGE_SIZE, config.DATA.IMAGE_SIZE)),
        transforms.ToTensor(),
        transforms.Normalize(mean=config.DATA.IMAGENET_MEAN, std=config.DATA.IMAGENET_STD),
    ])
    return transforms_val


def get_dataset(config, mode='train'):
    """ Get dataset from config and mode (train/val)

    Returns the related dataset object according to configs and mode(train/val)

    Args:
        config: configs contains dataset related settings. see config.py for details
    Returns:
        dataset: dataset object
    """
    assert mode in ['train', 'val']
    if config.DATA.DATASET == "cifar10":
        if mode == 'train':
            dataset = datasets.Cifar10(mode=mode, transform=get_train_transforms(config))
        else:
            mode = 'test'
            dataset = datasets.Cifar10(mode=mode, transform=get_val_transforms(config))
    elif config.DATA.DATASET == "cifar100":
        if mode == 'train':
            dataset = datasets.Cifar100(mode=mode, transform=get_train_transforms(config))
        else:
            mode = 'test'
            dataset = datasets.Cifar100(mode=mode, transform=get_val_transforms(config))
    elif config.DATA.DATASET == "imagenet2012":
        if mode == 'train':
            dataset = ImageNet2012Dataset(config.DATA.DATA_PATH,
                                          mode=mode,
                                          transform=get_train_dino_transforms(config))
        else:
            dataset = ImageNet2012Dataset(config.DATA.DATA_PATH,
                                          mode=mode,
                                          transform=get_val_transforms(config))
    else:
        raise NotImplementedError(
            "[{config.DATA.DATASET}] Only cifar10, cifar100, imagenet2012 are supported now")
    return dataset


def get_dataloader(config, dataset, mode='train', multi_process=False):
    """Get dataloader with config, dataset, mode as input, allows multiGPU settings.

        Multi-GPU loader is implements as distributedBatchSampler.

    Args:
        config: see config.py for details
        dataset: paddle.io.dataset object
        mode: train/val
        multi_process: if True, use DistributedBatchSampler to support multi-processing
    Returns:
        dataloader: paddle.io.DataLoader object.
    """

    if mode == 'train':
        batch_size = config.DATA.BATCH_SIZE
    else:
        batch_size = config.DATA.BATCH_SIZE_EVAL

    if multi_process is True:
        sampler = DistributedBatchSampler(dataset,
                                          batch_size=batch_size,
                                          shuffle=(mode == 'train'))
        dataloader = DataLoader(dataset,
                                batch_sampler=sampler,
                                num_workers=config.DATA.NUM_WORKERS)
    else:
        dataloader = DataLoader(dataset,
                                batch_size=batch_size,
                                num_workers=config.DATA.NUM_WORKERS,
                                shuffle=(mode == 'train'))
    return dataloader