ContraCLIP/models/genforce/datasets/datasets.py

# python3.7
"""Contains the class of dataset."""

import os
import pickle
import string
import zipfile
import numpy as np
import cv2
import lmdb

import torch
from torch.utils.data import Dataset

from .transforms import progressive_resize_image
from .transforms import crop_resize_image
from .transforms import resize_image
from .transforms import normalize_image

try:
    import turbojpeg
    BASE_DIR = os.path.dirname(os.path.relpath(__file__))
    LIBRARY_NAME = 'libturbojpeg.so.0'
    LIBRARY_PATH = os.path.join(BASE_DIR, LIBRARY_NAME)
    jpeg = turbojpeg.TurboJPEG(LIBRARY_PATH)
except ImportError:
    jpeg = None

__all__ = ['BaseDataset']

_FORMATS_ALLOWED = ['dir', 'lmdb', 'list', 'zip']


class ZipLoader(object):
    """Defines a class to load zip file.

    This is a static class, which is used to solve the problem that different
    data workers can not share the same memory.
    """
    files = dict()

    @staticmethod
    def get_zipfile(file_path):
        """Fetches a zip file."""
        zip_files = ZipLoader.files
        if file_path not in zip_files:
            zip_files[file_path] = zipfile.ZipFile(file_path, 'r')
        return zip_files[file_path]

    @staticmethod
    def get_image(file_path, image_path):
        """Decodes an image from a particular zip file."""
        zip_file = ZipLoader.get_zipfile(file_path)
        image_str = zip_file.read(image_path)
        image_np = np.frombuffer(image_str, np.uint8)
        image = cv2.imdecode(image_np, cv2.IMREAD_COLOR)
        return image


class LmdbLoader(object):
    """Defines a class to load lmdb file.

    This is a static class, which is used to solve lmdb loading error
    when num_workers > 0
    """
    files = dict()

    @staticmethod
    def get_lmdbfile(file_path):
        """Fetches a lmdb file"""
        lmdb_files = LmdbLoader.files
        if 'env' not in lmdb_files:
            env = lmdb.open(file_path,
                            max_readers=1,
                            readonly=True,
                            lock=False,
                            readahead=False,
                            meminit=False)
            with env.begin(write=False) as txn:
                num_samples = txn.stat()['entries']
            cache_file = '_cache_' + ''.join(
                c for c in file_path if c in string.ascii_letters)
            if os.path.isfile(cache_file):
                keys = pickle.load(open(cache_file, "rb"))
            else:
                with env.begin(write=False) as txn:
                    keys = [key for key, _ in txn.cursor()]
                pickle.dump(keys, open(cache_file, "wb"))
            lmdb_files['env'] = env
            lmdb_files['num_samples'] = num_samples
            lmdb_files['keys'] = keys
        return lmdb_files

    @staticmethod
    def get_image(file_path, idx):
        """Decodes an image from a particular lmdb file"""
        lmdb_files = LmdbLoader.get_lmdbfile(file_path)
        env = lmdb_files['env']
        keys = lmdb_files['keys']
        with env.begin(write=False) as txn:
            imagebuf = txn.get(keys[idx])
        image_np = np.frombuffer(imagebuf, np.uint8)
        image = cv2.imdecode(image_np, cv2.IMREAD_COLOR)
        return image


class BaseDataset(Dataset):
    """Defines the base dataset class.

    This class supports loading data from a full-of-image folder, a lmdb
    database, or an image list. Images will be pre-processed based on the given
    `transform` function before fed into the data loader.

    NOTE: The loaded data will be returned as a directory, where there must be
    a key `image`.
    """
    def __init__(self,
                 root_dir,
                 resolution,
                 data_format='dir',
                 image_list_path=None,
                 mirror=0.0,
                 progressive_resize=True,
                 crop_resize_resolution=-1,
                 transform=normalize_image,
                 transform_kwargs=None,
                 **_unused_kwargs):
        """Initializes the dataset.

        Args:
            root_dir: Root directory containing the dataset.
            resolution: The resolution of the returned image.
            data_format: Format the dataset is stored. Supports `dir`, `lmdb`,
                and `list`. (default: `dir`)
            image_list_path: Path to the image list. This field is required if
                `data_format` is `list`. (default: None)
            mirror: The probability to do mirror augmentation. (default: 0.0)
            progressive_resize: Whether to resize images progressively.
                (default: True)
            crop_resize_resolution: The resolution of the output after crop
                and resize. (default: -1)
            transform: The transform function for pre-processing.
                (default: `datasets.transforms.normalize_image()`)
            transform_kwargs: The additional arguments for the `transform`
                function. (default: None)

        Raises:
            ValueError: If the input `data_format` is not supported.
            NotImplementedError: If the input `data_format` is not implemented.
        """
        if data_format.lower() not in _FORMATS_ALLOWED:
            raise ValueError(f'Invalid data format `{data_format}`!\n'
                             f'Supported formats: {_FORMATS_ALLOWED}.')

        self.root_dir = root_dir
        self.resolution = resolution
        self.data_format = data_format.lower()
        self.image_list_path = image_list_path
        self.mirror = np.clip(mirror, 0.0, 1.0)
        self.progressive_resize = progressive_resize
        self.crop_resize_resolution = crop_resize_resolution
        self.transform = transform
        self.transform_kwargs = transform_kwargs or dict()

        if self.data_format == 'dir':
            self.image_paths = sorted(os.listdir(self.root_dir))
            self.num_samples = len(self.image_paths)
        elif self.data_format == 'lmdb':
            lmdb_file = LmdbLoader.get_lmdbfile(self.root_dir)
            self.num_samples = lmdb_file['num_samples']
        elif self.data_format == 'list':
            self.metas = []
            assert os.path.isfile(self.image_list_path)
            with open(self.image_list_path) as f:
                for line in f:
                    fields = line.rstrip().split(' ')
                    if len(fields) == 1:
                        self.metas.append((fields[0], None))
                    else:
                        assert len(fields) == 2
                        self.metas.append((fields[0], int(fields[1])))
            self.num_samples = len(self.metas)
        elif self.data_format == 'zip':
            zip_file = ZipLoader.get_zipfile(self.root_dir)
            image_paths = [f for f in zip_file.namelist()
                           if ('.jpg' in f or '.jpeg' in f or '.png' in f)]
            self.image_paths = sorted(image_paths)
            self.num_samples = len(self.image_paths)
        else:
            raise NotImplementedError(f'Not implemented data format '
                                      f'`{self.data_format}`!')

    def __len__(self):
        return self.num_samples

    def __getitem__(self, idx):
        data = dict()

        # Load data.
        if self.data_format == 'dir':
            image_path = self.image_paths[idx]
            try:
                in_file = open(os.path.join(self.root_dir, image_path), 'rb')
                image = jpeg.decode(in_file.read())
            except:  # pylint: disable=bare-except
                image = cv2.imread(os.path.join(self.root_dir, image_path))
        elif self.data_format == 'lmdb':
            image = LmdbLoader.get_image(self.root_dir, idx)
        elif self.data_format == 'list':
            image_path, label = self.metas[idx]
            image = cv2.imread(os.path.join(self.root_dir, image_path))
            label = None if label is None else torch.LongTensor(label)
            # data.update({'label': label})
        elif self.data_format == 'zip':
            image_path = self.image_paths[idx]
            image = ZipLoader.get_image(self.root_dir, image_path)
        else:
            raise NotImplementedError(f'Not implemented data format '
                                      f'`{self.data_format}`!')

        image = image[:, :, ::-1]  # Converts BGR (cv2) to RGB.

        # Transform image.
        if self.crop_resize_resolution > 0:
            image = crop_resize_image(image, self.crop_resize_resolution)
        if self.progressive_resize:
            image = progressive_resize_image(image, self.resolution)
        image = image.transpose(2, 0, 1).astype(np.float32)
        if np.random.uniform() < self.mirror:
            image = image[:, :, ::-1]  # CHW
        image = torch.FloatTensor(image.copy())
        if not self.progressive_resize:
            image = resize_image(image, self.resolution)

        if self.transform is not None:
            image = self.transform(image, **self.transform_kwargs)
        data.update({'image': image})

        return data