Upload folder using huggingface_hub

866ee56 verified over 1 year ago

18.3 kB

# --------------------------------------------------------
# InternVL
# Copyright (c) 2023 OpenGVLab
# Licensed under The MIT License [see LICENSE for details]
# --------------------------------------------------------

import io
import json
import logging
import math
import os
import os.path as osp
import re
import time
from abc import abstractmethod

import mmcv
import torch
import torch.distributed as dist
import torch.utils.data as data
from mmcv.fileio import FileClient
from PIL import Image
from tqdm import tqdm, trange

from .zipreader import ZipReader, is_zip_path

_logger = logging.getLogger(__name__)

_ERROR_RETRY = 50


def has_file_allowed_extension(filename, extensions):
    """Checks if a file is an allowed extension.
    Args:
        filename (string): path to a file
    Returns:
        bool: True if the filename ends with a known image extension
    """
    filename_lower = filename.lower()
    return any(filename_lower.endswith(ext) for ext in extensions)


def find_classes(dir):
    classes = [
        d for d in os.listdir(dir) if os.path.isdir(os.path.join(dir, d))
    ]
    classes.sort()
    class_to_idx = {classes[i]: i for i in range(len(classes))}
    return classes, class_to_idx


def make_dataset(dir, class_to_idx, extensions):
    images = []
    dir = os.path.expanduser(dir)
    for target in sorted(os.listdir(dir)):
        d = os.path.join(dir, target)
        if not os.path.isdir(d):
            continue
        for root, _, fnames in sorted(os.walk(d)):
            for fname in sorted(fnames):
                if has_file_allowed_extension(fname, extensions):
                    path = os.path.join(root, fname)
                    item = (path, class_to_idx[target])
                    images.append(item)

    return images


def make_dataset_with_ann(ann_file, img_prefix, extensions):
    images = []
    with open(ann_file, 'r') as f:
        contents = f.readlines()
        for line_str in contents:
            path_contents = [c for c in line_str.split('\t')]
            im_file_name = path_contents[0]
            class_index = int(path_contents[1])
            assert str.lower(os.path.splitext(im_file_name)[-1]) in extensions
            item = (os.path.join(img_prefix, im_file_name), class_index)
            images.append(item)

    return images


class DatasetFolder(data.Dataset):
    """A generic data loader where the samples are arranged in this way: ::
    root/class_x/xxx.ext
        root/class_x/xxy.ext
        root/class_x/xxz.ext
        root/class_y/123.ext
        root/class_y/nsdf3.ext
        root/class_y/asd932_.ext
    Args:
        root (string): Root directory path.
        loader (callable): A function to load a sample given its path.
        extensions (list[string]): A list of allowed extensions.
        transform (callable, optional): A function/transform that takes in
            a sample and returns a transformed version.
            E.g, ``transforms.RandomCrop`` for images.
        target_transform (callable, optional): A function/transform that takes
            in the target and transforms it.
     Attributes:
        samples (list): List of (sample path, class_index) tuples
    """

    def __init__(self,
                 root,
                 loader,
                 extensions,
                 ann_file='',
                 img_prefix='',
                 transform=None,
                 target_transform=None,
                 cache_mode='no'):
        # image folder mode
        if ann_file == '':
            _, class_to_idx = find_classes(root)
            samples = make_dataset(root, class_to_idx, extensions)
        # zip mode
        else:
            samples = make_dataset_with_ann(os.path.join(root, ann_file),
                                            os.path.join(root, img_prefix),
                                            extensions)

        if len(samples) == 0:
            raise (RuntimeError('Found 0 files in subfolders of: ' + root +
                                '\n' + 'Supported extensions are: ' +
                                ','.join(extensions)))

        self.root = root
        self.loader = loader
        self.extensions = extensions

        self.samples = samples
        self.labels = [y_1k for _, y_1k in samples]
        self.classes = list(set(self.labels))

        self.transform = transform
        self.target_transform = target_transform

        self.cache_mode = cache_mode
        if self.cache_mode != 'no':
            self.init_cache()

    def init_cache(self):
        assert self.cache_mode in ['part', 'full']
        n_sample = len(self.samples)
        global_rank = dist.get_rank()
        world_size = dist.get_world_size()

        samples_bytes = [None for _ in range(n_sample)]
        start_time = time.time()
        for index in range(n_sample):
            if index % (n_sample // 10) == 0:
                t = time.time() - start_time
                print(
                    f'global_rank {dist.get_rank()} cached {index}/{n_sample} takes {t:.2f}s per block'
                )
                start_time = time.time()
            path, target = self.samples[index]
            if self.cache_mode == 'full':
                samples_bytes[index] = (ZipReader.read(path), target)
            elif self.cache_mode == 'part' and index % world_size == global_rank:
                samples_bytes[index] = (ZipReader.read(path), target)
            else:
                samples_bytes[index] = (path, target)
        self.samples = samples_bytes

    def __getitem__(self, index):
        """
        Args:
            index (int): Index
        Returns:
            tuple: (sample, target) where target is class_index of the target class.
        """
        path, target = self.samples[index]
        sample = self.loader(path)
        if self.transform is not None:
            sample = self.transform(sample)
        if self.target_transform is not None:
            target = self.target_transform(target)

        return sample, target

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

    def __repr__(self):
        fmt_str = 'Dataset ' + self.__class__.__name__ + '\n'
        fmt_str += '    Number of datapoints: {}\n'.format(self.__len__())
        fmt_str += '    Root Location: {}\n'.format(self.root)
        tmp = '    Transforms (if any): '
        fmt_str += '{0}{1}\n'.format(
            tmp,
            self.transform.__repr__().replace('\n', '\n' + ' ' * len(tmp)))
        tmp = '    Target Transforms (if any): '
        fmt_str += '{0}{1}'.format(
            tmp,
            self.target_transform.__repr__().replace('\n',
                                                     '\n' + ' ' * len(tmp)))

        return fmt_str


IMG_EXTENSIONS = ['.jpg', '.jpeg', '.png', '.ppm', '.bmp', '.pgm', '.tif']


def pil_loader(path):
    # open path as file to avoid ResourceWarning (https://github.com/python-pillow/Pillow/issues/835)
    if isinstance(path, bytes):
        img = Image.open(io.BytesIO(path))
    elif is_zip_path(path):
        data = ZipReader.read(path)
        img = Image.open(io.BytesIO(data))
    else:
        with open(path, 'rb') as f:
            img = Image.open(f)
            return img.convert('RGB')

    return img.convert('RGB')


def accimage_loader(path):
    import accimage
    try:
        return accimage.Image(path)
    except IOError:
        # Potentially a decoding problem, fall back to PIL.Image
        return pil_loader(path)


def default_img_loader(path):
    from torchvision import get_image_backend
    if get_image_backend() == 'accimage':
        return accimage_loader(path)
    else:
        return pil_loader(path)


class CachedImageFolder(DatasetFolder):
    """A generic data loader where the images are arranged in this way: ::
    root/dog/xxx.png
        root/dog/xxy.png
        root/dog/xxz.png
        root/cat/123.png
        root/cat/nsdf3.png
        root/cat/asd932_.png
    Args:
        root (string): Root directory path.
        transform (callable, optional): A function/transform that  takes in an PIL image
            and returns a transformed version. E.g, ``transforms.RandomCrop``
        target_transform (callable, optional): A function/transform that takes in the
            target and transforms it.
        loader (callable, optional): A function to load an image given its path.
     Attributes:
        imgs (list): List of (image path, class_index) tuples
    """

    def __init__(self,
                 root,
                 ann_file='',
                 img_prefix='',
                 transform=None,
                 target_transform=None,
                 loader=default_img_loader,
                 cache_mode='no'):
        super(CachedImageFolder,
              self).__init__(root,
                             loader,
                             IMG_EXTENSIONS,
                             ann_file=ann_file,
                             img_prefix=img_prefix,
                             transform=transform,
                             target_transform=target_transform,
                             cache_mode=cache_mode)
        self.imgs = self.samples

    def __getitem__(self, index):
        """
        Args:
            index (int): Index
        Returns:
            tuple: (image, target) where target is class_index of the target class.
        """
        path, target = self.samples[index]
        image = self.loader(path)
        if self.transform is not None:
            img = self.transform(image)
        else:
            img = image
        if self.target_transform is not None:
            target = self.target_transform(target)

        return img, target


class ImageCephDataset(data.Dataset):

    def __init__(self,
                 root,
                 split,
                 parser=None,
                 transform=None,
                 target_transform=None,
                 on_memory=False):
        if '22k' in root:
            # Imagenet 22k
            annotation_root = 'meta_data/'
        else:
            # Imagenet
            annotation_root = 'meta_data/'
        if parser is None or isinstance(parser, str):
            parser = ParserCephImage(root=root,
                                     split=split,
                                     annotation_root=annotation_root,
                                     on_memory=on_memory)
        self.parser = parser
        self.transform = transform
        self.target_transform = target_transform
        self._consecutive_errors = 0

    def __getitem__(self, index):
        img, target = self.parser[index]
        self._consecutive_errors = 0
        if self.transform is not None:
            img = self.transform(img)
        if target is None:
            target = -1
        elif self.target_transform is not None:
            target = self.target_transform(target)
        return img, target

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

    def filename(self, index, basename=False, absolute=False):
        return self.parser.filename(index, basename, absolute)

    def filenames(self, basename=False, absolute=False):
        return self.parser.filenames(basename, absolute)


class Parser:

    def __init__(self):
        pass

    @abstractmethod
    def _filename(self, index, basename=False, absolute=False):
        pass

    def filename(self, index, basename=False, absolute=False):
        return self._filename(index, basename=basename, absolute=absolute)

    def filenames(self, basename=False, absolute=False):
        return [
            self._filename(index, basename=basename, absolute=absolute)
            for index in range(len(self))
        ]


class ParserCephImage(Parser):

    def __init__(self,
                 root,
                 split,
                 annotation_root,
                 on_memory=False,
                 **kwargs):
        super().__init__()

        self.file_client = None
        self.kwargs = kwargs

        self.root = root  # dataset:s3://imagenet22k
        if '22k' in root:
            self.io_backend = 'petrel'
            with open(osp.join(annotation_root, '22k_class_to_idx.json'),
                      'r') as f:
                self.class_to_idx = json.loads(f.read())
            with open(osp.join(annotation_root, '22k_label.txt'), 'r') as f:
                self.samples = f.read().splitlines()
        else:
            self.io_backend = 'disk'
            self.class_to_idx = None
            with open(osp.join(annotation_root, f'{split}.txt'), 'r') as f:
                self.samples = f.read().splitlines()
        local_rank = None
        local_size = None
        self._consecutive_errors = 0
        self.on_memory = on_memory
        if on_memory:
            self.holder = {}
            if local_rank is None:
                local_rank = int(os.environ.get('LOCAL_RANK', 0))
            if local_size is None:
                local_size = int(os.environ.get('LOCAL_SIZE', 1))
            self.local_rank = local_rank
            self.local_size = local_size
            self.rank = int(os.environ['RANK'])
            self.world_size = int(os.environ['WORLD_SIZE'])
            self.num_replicas = int(os.environ['WORLD_SIZE'])
            self.num_parts = local_size
            self.num_samples = int(
                math.ceil(len(self.samples) * 1.0 / self.num_replicas))
            self.total_size = self.num_samples * self.num_replicas
            self.total_size_parts = self.num_samples * self.num_replicas // self.num_parts
            self.load_onto_memory_v2()

    def load_onto_memory(self):
        print('Loading images onto memory...', self.local_rank,
              self.local_size)
        if self.file_client is None:
            self.file_client = FileClient(self.io_backend, **self.kwargs)
        for index in trange(len(self.samples)):
            if index % self.local_size != self.local_rank:
                continue
            path, _ = self.samples[index].split(' ')
            path = osp.join(self.root, path)
            img_bytes = self.file_client.get(path)
            self.holder[path] = img_bytes

        print('Loading complete!')

    def load_onto_memory_v2(self):
        # print("Loading images onto memory...", self.local_rank, self.local_size)
        t = torch.Generator()
        t.manual_seed(0)
        indices = torch.randperm(len(self.samples), generator=t).tolist()
        # indices = range(len(self.samples))
        indices = [i for i in indices if i % self.num_parts == self.local_rank]
        # add extra samples to make it evenly divisible
        indices += indices[:(self.total_size_parts - len(indices))]
        assert len(indices) == self.total_size_parts

        # subsample
        indices = indices[self.rank // self.num_parts:self.
                          total_size_parts:self.num_replicas // self.num_parts]
        assert len(indices) == self.num_samples

        if self.file_client is None:
            self.file_client = FileClient(self.io_backend, **self.kwargs)
        for index in tqdm(indices):
            if index % self.local_size != self.local_rank:
                continue
            path, _ = self.samples[index].split(' ')
            path = osp.join(self.root, path)
            img_bytes = self.file_client.get(path)

            self.holder[path] = img_bytes

        print('Loading complete!')

    def __getitem__(self, index):
        if self.file_client is None:
            self.file_client = FileClient(self.io_backend, **self.kwargs)

        filepath, target = self.samples[index].split(' ')
        filepath = osp.join(self.root, filepath)

        try:
            if self.on_memory:
                img_bytes = self.holder[filepath]
            else:
                # pass
                img_bytes = self.file_client.get(filepath)
            img = mmcv.imfrombytes(img_bytes)[:, :, ::-1]
        except Exception as e:
            _logger.warning(
                f'Skipped sample (index {index}, file {filepath}). {str(e)}')
            self._consecutive_errors += 1
            if self._consecutive_errors < _ERROR_RETRY:
                return self.__getitem__((index + 1) % len(self))
            else:
                raise e
        self._consecutive_errors = 0

        img = Image.fromarray(img)
        try:
            if self.class_to_idx is not None:
                target = self.class_to_idx[target]
            else:
                target = int(target)
        except:
            print(filepath, target)
            exit()

        return img, target

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

    def _filename(self, index, basename=False, absolute=False):
        filename, _ = self.samples[index].split(' ')
        filename = osp.join(self.root, filename)

        return filename


def get_temporal_info(date, miss_hour=False):
    try:
        if date:
            if miss_hour:
                pattern = re.compile(r'(\d*)-(\d*)-(\d*)', re.I)
            else:
                pattern = re.compile(r'(\d*)-(\d*)-(\d*) (\d*):(\d*):(\d*)',
                                     re.I)
            m = pattern.match(date.strip())

            if m:
                year = int(m.group(1))
                month = int(m.group(2))
                day = int(m.group(3))
                x_month = math.sin(2 * math.pi * month / 12)
                y_month = math.cos(2 * math.pi * month / 12)
                if miss_hour:
                    x_hour = 0
                    y_hour = 0
                else:
                    hour = int(m.group(4))
                    x_hour = math.sin(2 * math.pi * hour / 24)
                    y_hour = math.cos(2 * math.pi * hour / 24)
                return [x_month, y_month, x_hour, y_hour]
            else:
                return [0, 0, 0, 0]
        else:
            return [0, 0, 0, 0]
    except:
        return [0, 0, 0, 0]


def get_spatial_info(latitude, longitude):
    if latitude and longitude:
        latitude = math.radians(latitude)
        longitude = math.radians(longitude)
        x = math.cos(latitude) * math.cos(longitude)
        y = math.cos(latitude) * math.sin(longitude)
        z = math.sin(latitude)
        return [x, y, z]
    else:
        return [0, 0, 0]

	# --------------------------------------------------------
	# InternVL
	# Copyright (c) 2023 OpenGVLab
	# Licensed under The MIT License [see LICENSE for details]
	# --------------------------------------------------------

	import io
	import json
	import logging
	import math
	import os
	import os.path as osp
	import re
	import time
	from abc import abstractmethod

	import mmcv
	import torch
	import torch.distributed as dist
	import torch.utils.data as data
	from mmcv.fileio import FileClient
	from PIL import Image
	from tqdm import tqdm, trange

	from .zipreader import ZipReader, is_zip_path

	_logger = logging.getLogger(__name__)

	_ERROR_RETRY = 50


	def has_file_allowed_extension(filename, extensions):
	"""Checks if a file is an allowed extension.

	Args:
	filename (string): path to a file
	Returns:
	bool: True if the filename ends with a known image extension
	"""
	filename_lower = filename.lower()
	return any(filename_lower.endswith(ext) for ext in extensions)


	def find_classes(dir):
	classes = [
	d for d in os.listdir(dir) if os.path.isdir(os.path.join(dir, d))
	]
	classes.sort()
	class_to_idx = {classes[i]: i for i in range(len(classes))}
	return classes, class_to_idx


	def make_dataset(dir, class_to_idx, extensions):
	images = []
	dir = os.path.expanduser(dir)
	for target in sorted(os.listdir(dir)):
	d = os.path.join(dir, target)
	if not os.path.isdir(d):
	continue
	for root, _, fnames in sorted(os.walk(d)):
	for fname in sorted(fnames):
	if has_file_allowed_extension(fname, extensions):
	path = os.path.join(root, fname)
	item = (path, class_to_idx[target])
	images.append(item)

	return images


	def make_dataset_with_ann(ann_file, img_prefix, extensions):
	images = []
	with open(ann_file, 'r') as f:
	contents = f.readlines()
	for line_str in contents:
	path_contents = [c for c in line_str.split('\t')]
	im_file_name = path_contents[0]
	class_index = int(path_contents[1])
	assert str.lower(os.path.splitext(im_file_name)[-1]) in extensions
	item = (os.path.join(img_prefix, im_file_name), class_index)
	images.append(item)

	return images


	class DatasetFolder(data.Dataset):
	"""A generic data loader where the samples are arranged in this way: ::

	root/class_x/xxx.ext
	root/class_x/xxy.ext
	root/class_x/xxz.ext
	root/class_y/123.ext
	root/class_y/nsdf3.ext
	root/class_y/asd932_.ext
	Args:
	root (string): Root directory path.
	loader (callable): A function to load a sample given its path.
	extensions (list[string]): A list of allowed extensions.
	transform (callable, optional): A function/transform that takes in
	a sample and returns a transformed version.
	E.g, ``transforms.RandomCrop`` for images.
	target_transform (callable, optional): A function/transform that takes
	in the target and transforms it.
	Attributes:
	samples (list): List of (sample path, class_index) tuples
	"""

	def __init__(self,
	root,
	loader,
	extensions,
	ann_file='',
	img_prefix='',
	transform=None,
	target_transform=None,
	cache_mode='no'):
	# image folder mode
	if ann_file == '':
	_, class_to_idx = find_classes(root)
	samples = make_dataset(root, class_to_idx, extensions)
	# zip mode
	else:
	samples = make_dataset_with_ann(os.path.join(root, ann_file),
	os.path.join(root, img_prefix),
	extensions)

	if len(samples) == 0:
	raise (RuntimeError('Found 0 files in subfolders of: ' + root +
	'\n' + 'Supported extensions are: ' +
	','.join(extensions)))

	self.root = root
	self.loader = loader
	self.extensions = extensions

	self.samples = samples
	self.labels = [y_1k for _, y_1k in samples]
	self.classes = list(set(self.labels))

	self.transform = transform
	self.target_transform = target_transform

	self.cache_mode = cache_mode
	if self.cache_mode != 'no':
	self.init_cache()

	def init_cache(self):
	assert self.cache_mode in ['part', 'full']
	n_sample = len(self.samples)
	global_rank = dist.get_rank()
	world_size = dist.get_world_size()

	samples_bytes = [None for _ in range(n_sample)]
	start_time = time.time()
	for index in range(n_sample):
	if index % (n_sample // 10) == 0:
	t = time.time() - start_time
	print(
	f'global_rank {dist.get_rank()} cached {index}/{n_sample} takes {t:.2f}s per block'
	)
	start_time = time.time()
	path, target = self.samples[index]
	if self.cache_mode == 'full':
	samples_bytes[index] = (ZipReader.read(path), target)
	elif self.cache_mode == 'part' and index % world_size == global_rank:
	samples_bytes[index] = (ZipReader.read(path), target)
	else:
	samples_bytes[index] = (path, target)
	self.samples = samples_bytes

	def __getitem__(self, index):
	"""
	Args:
	index (int): Index
	Returns:
	tuple: (sample, target) where target is class_index of the target class.
	"""
	path, target = self.samples[index]
	sample = self.loader(path)
	if self.transform is not None:
	sample = self.transform(sample)
	if self.target_transform is not None:
	target = self.target_transform(target)

	return sample, target

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

	def __repr__(self):
	fmt_str = 'Dataset ' + self.__class__.__name__ + '\n'
	fmt_str += ' Number of datapoints: {}\n'.format(self.__len__())
	fmt_str += ' Root Location: {}\n'.format(self.root)
	tmp = ' Transforms (if any): '
	fmt_str += '{0}{1}\n'.format(
	tmp,
	self.transform.__repr__().replace('\n', '\n' + ' ' * len(tmp)))
	tmp = ' Target Transforms (if any): '
	fmt_str += '{0}{1}'.format(
	tmp,
	self.target_transform.__repr__().replace('\n',
	'\n' + ' ' * len(tmp)))

	return fmt_str


	IMG_EXTENSIONS = ['.jpg', '.jpeg', '.png', '.ppm', '.bmp', '.pgm', '.tif']


	def pil_loader(path):
	# open path as file to avoid ResourceWarning (https://github.com/python-pillow/Pillow/issues/835)
	if isinstance(path, bytes):
	img = Image.open(io.BytesIO(path))
	elif is_zip_path(path):
	data = ZipReader.read(path)
	img = Image.open(io.BytesIO(data))
	else:
	with open(path, 'rb') as f:
	img = Image.open(f)
	return img.convert('RGB')

	return img.convert('RGB')


	def accimage_loader(path):
	import accimage
	try:
	return accimage.Image(path)
	except IOError:
	# Potentially a decoding problem, fall back to PIL.Image
	return pil_loader(path)


	def default_img_loader(path):
	from torchvision import get_image_backend
	if get_image_backend() == 'accimage':
	return accimage_loader(path)
	else:
	return pil_loader(path)


	class CachedImageFolder(DatasetFolder):
	"""A generic data loader where the images are arranged in this way: ::

	root/dog/xxx.png
	root/dog/xxy.png
	root/dog/xxz.png
	root/cat/123.png
	root/cat/nsdf3.png
	root/cat/asd932_.png
	Args:
	root (string): Root directory path.
	transform (callable, optional): A function/transform that takes in an PIL image
	and returns a transformed version. E.g, ``transforms.RandomCrop``
	target_transform (callable, optional): A function/transform that takes in the
	target and transforms it.
	loader (callable, optional): A function to load an image given its path.
	Attributes:
	imgs (list): List of (image path, class_index) tuples
	"""

	def __init__(self,
	root,
	ann_file='',
	img_prefix='',
	transform=None,
	target_transform=None,
	loader=default_img_loader,
	cache_mode='no'):
	super(CachedImageFolder,
	self).__init__(root,
	loader,
	IMG_EXTENSIONS,
	ann_file=ann_file,
	img_prefix=img_prefix,
	transform=transform,
	target_transform=target_transform,
	cache_mode=cache_mode)
	self.imgs = self.samples

	def __getitem__(self, index):
	"""
	Args:
	index (int): Index
	Returns:
	tuple: (image, target) where target is class_index of the target class.
	"""
	path, target = self.samples[index]
	image = self.loader(path)
	if self.transform is not None:
	img = self.transform(image)
	else:
	img = image
	if self.target_transform is not None:
	target = self.target_transform(target)

	return img, target


	class ImageCephDataset(data.Dataset):

	def __init__(self,
	root,
	split,
	parser=None,
	transform=None,
	target_transform=None,
	on_memory=False):
	if '22k' in root:
	# Imagenet 22k
	annotation_root = 'meta_data/'
	else:
	# Imagenet
	annotation_root = 'meta_data/'
	if parser is None or isinstance(parser, str):
	parser = ParserCephImage(root=root,
	split=split,
	annotation_root=annotation_root,
	on_memory=on_memory)
	self.parser = parser
	self.transform = transform
	self.target_transform = target_transform
	self._consecutive_errors = 0

	def __getitem__(self, index):
	img, target = self.parser[index]
	self._consecutive_errors = 0
	if self.transform is not None:
	img = self.transform(img)
	if target is None:
	target = -1
	elif self.target_transform is not None:
	target = self.target_transform(target)
	return img, target

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

	def filename(self, index, basename=False, absolute=False):
	return self.parser.filename(index, basename, absolute)

	def filenames(self, basename=False, absolute=False):
	return self.parser.filenames(basename, absolute)


	class Parser:

	def __init__(self):
	pass

	@abstractmethod
	def _filename(self, index, basename=False, absolute=False):
	pass

	def filename(self, index, basename=False, absolute=False):
	return self._filename(index, basename=basename, absolute=absolute)

	def filenames(self, basename=False, absolute=False):
	return [
	self._filename(index, basename=basename, absolute=absolute)
	for index in range(len(self))
	]


	class ParserCephImage(Parser):

	def __init__(self,
	root,
	split,
	annotation_root,
	on_memory=False,
	**kwargs):
	super().__init__()

	self.file_client = None
	self.kwargs = kwargs

	self.root = root # dataset:s3://imagenet22k
	if '22k' in root:
	self.io_backend = 'petrel'
	with open(osp.join(annotation_root, '22k_class_to_idx.json'),
	'r') as f:
	self.class_to_idx = json.loads(f.read())
	with open(osp.join(annotation_root, '22k_label.txt'), 'r') as f:
	self.samples = f.read().splitlines()
	else:
	self.io_backend = 'disk'
	self.class_to_idx = None
	with open(osp.join(annotation_root, f'{split}.txt'), 'r') as f:
	self.samples = f.read().splitlines()
	local_rank = None
	local_size = None
	self._consecutive_errors = 0
	self.on_memory = on_memory
	if on_memory:
	self.holder = {}
	if local_rank is None:
	local_rank = int(os.environ.get('LOCAL_RANK', 0))
	if local_size is None:
	local_size = int(os.environ.get('LOCAL_SIZE', 1))
	self.local_rank = local_rank
	self.local_size = local_size
	self.rank = int(os.environ['RANK'])
	self.world_size = int(os.environ['WORLD_SIZE'])
	self.num_replicas = int(os.environ['WORLD_SIZE'])
	self.num_parts = local_size
	self.num_samples = int(
	math.ceil(len(self.samples) * 1.0 / self.num_replicas))
	self.total_size = self.num_samples * self.num_replicas
	self.total_size_parts = self.num_samples * self.num_replicas // self.num_parts
	self.load_onto_memory_v2()

	def load_onto_memory(self):
	print('Loading images onto memory...', self.local_rank,
	self.local_size)
	if self.file_client is None:
	self.file_client = FileClient(self.io_backend, **self.kwargs)
	for index in trange(len(self.samples)):
	if index % self.local_size != self.local_rank:
	continue
	path, _ = self.samples[index].split(' ')
	path = osp.join(self.root, path)
	img_bytes = self.file_client.get(path)
	self.holder[path] = img_bytes

	print('Loading complete!')

	def load_onto_memory_v2(self):
	# print("Loading images onto memory...", self.local_rank, self.local_size)
	t = torch.Generator()
	t.manual_seed(0)
	indices = torch.randperm(len(self.samples), generator=t).tolist()
	# indices = range(len(self.samples))
	indices = [i for i in indices if i % self.num_parts == self.local_rank]
	# add extra samples to make it evenly divisible
	indices += indices[:(self.total_size_parts - len(indices))]
	assert len(indices) == self.total_size_parts

	# subsample
	indices = indices[self.rank // self.num_parts:self.
	total_size_parts:self.num_replicas // self.num_parts]
	assert len(indices) == self.num_samples

	if self.file_client is None:
	self.file_client = FileClient(self.io_backend, **self.kwargs)
	for index in tqdm(indices):
	if index % self.local_size != self.local_rank:
	continue
	path, _ = self.samples[index].split(' ')
	path = osp.join(self.root, path)
	img_bytes = self.file_client.get(path)

	self.holder[path] = img_bytes

	print('Loading complete!')

	def __getitem__(self, index):
	if self.file_client is None:
	self.file_client = FileClient(self.io_backend, **self.kwargs)

	filepath, target = self.samples[index].split(' ')
	filepath = osp.join(self.root, filepath)

	try:
	if self.on_memory:
	img_bytes = self.holder[filepath]
	else:
	# pass
	img_bytes = self.file_client.get(filepath)
	img = mmcv.imfrombytes(img_bytes)[:, :, ::-1]
	except Exception as e:
	_logger.warning(
	f'Skipped sample (index {index}, file {filepath}). {str(e)}')
	self._consecutive_errors += 1
	if self._consecutive_errors < _ERROR_RETRY:
	return self.__getitem__((index + 1) % len(self))
	else:
	raise e
	self._consecutive_errors = 0

	img = Image.fromarray(img)
	try:
	if self.class_to_idx is not None:
	target = self.class_to_idx[target]
	else:
	target = int(target)
	except:
	print(filepath, target)
	exit()

	return img, target

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

	def _filename(self, index, basename=False, absolute=False):
	filename, _ = self.samples[index].split(' ')
	filename = osp.join(self.root, filename)

	return filename


	def get_temporal_info(date, miss_hour=False):
	try:
	if date:
	if miss_hour:
	pattern = re.compile(r'(\d)-(\d)-(\d*)', re.I)
	else:
	pattern = re.compile(r'(\d)-(\d)-(\d) (\d):(\d):(\d)',
	re.I)
	m = pattern.match(date.strip())

	if m:
	year = int(m.group(1))
	month = int(m.group(2))
	day = int(m.group(3))
	x_month = math.sin(2 * math.pi * month / 12)
	y_month = math.cos(2 * math.pi * month / 12)
	if miss_hour:
	x_hour = 0
	y_hour = 0
	else:
	hour = int(m.group(4))
	x_hour = math.sin(2 * math.pi * hour / 24)
	y_hour = math.cos(2 * math.pi * hour / 24)
	return [x_month, y_month, x_hour, y_hour]
	else:
	return [0, 0, 0, 0]
	else:
	return [0, 0, 0, 0]
	except:
	return [0, 0, 0, 0]


	def get_spatial_info(latitude, longitude):
	if latitude and longitude:
	latitude = math.radians(latitude)
	longitude = math.radians(longitude)
	x = math.cos(latitude) * math.cos(longitude)
	y = math.cos(latitude) * math.sin(longitude)
	z = math.sin(latitude)
	return [x, y, z]
	else:
	return [0, 0, 0]