text: import copy
  import time
  
  import torch
  import torchvision
  from torch.utils.data import DataLoader, RandomSampler, SequentialSampler, random_split
  from torch.utils.data.distributed import DistributedSampler
  from torchvision.datasets import PhotoTour, Kinetics400, HMDB51, UCF101, Cityscapes, CocoCaptions, CocoDetection, \
      SBDataset, VOCSegmentation, VOCDetection
  
  from torchdistill.common.constant import def_logger
  from torchdistill.datasets.coco import ImageToTensor, Compose, CocoRandomHorizontalFlip, get_coco
  from torchdistill.datasets.collator import get_collate_func
  from torchdistill.datasets.registry import DATASET_DICT
  from torchdistill.datasets.sample_loader import get_sample_loader
  from torchdistill.datasets.sampler import get_batch_sampler
  from torchdistill.datasets.transform import TRANSFORM_CLASS_DICT, CustomCompose
  from torchdistill.datasets.wrapper import default_idx2subpath, BaseDatasetWrapper, CacheableDataset, get_dataset_wrapper
  
  logger = def_logger.getChild(__name__)
  
  TRANSFORM_CLASS_DICT.update(torchvision.transforms.__dict__)
  
  
  def load_coco_dataset(img_dir_path, ann_file_path, annotated_only, random_horizontal_flip=None, is_segment=False,
                        transforms=None, jpeg_quality=None):
      if transforms is None:
          transform_list = [ImageToTensor(jpeg_quality)]
          if random_horizontal_flip is not None and not is_segment:
              transform_list.append(CocoRandomHorizontalFlip(random_horizontal_flip))
          transforms = Compose(transform_list)
      return get_coco(img_dir_path=img_dir_path, ann_file_path=ann_file_path,
                      transforms=transforms, annotated_only=annotated_only, is_segment=is_segment)
  
  
  def build_transform(transform_params_config, compose_cls=None):
      if not isinstance(transform_params_config, (dict, list)) or len(transform_params_config) == 0:
          return None
  
      if isinstance(compose_cls, str):
          compose_cls = TRANSFORM_CLASS_DICT[compose_cls]
  
      component_list = list()
      if isinstance(transform_params_config, dict):
          for component_key in sorted(transform_params_config.keys()):
              component_config = transform_params_config[component_key]
              params_config = component_config.get('params', dict())
              if params_config is None:
                  params_config = dict()
  
              component = TRANSFORM_CLASS_DICT[component_config['type']](**params_config)
              component_list.append(component)
      else:
          for component_config in transform_params_config:
              params_config = component_config.get('params', dict())
              if params_config is None:
                  params_config = dict()
  
              component = TRANSFORM_CLASS_DICT[component_config['type']](**params_config)
              component_list.append(component)
      return torchvision.transforms.Compose(component_list) if compose_cls is None else compose_cls(component_list)
  
  
  def get_torchvision_dataset(dataset_cls, dataset_params_config):
      params_config = dataset_params_config.copy()
      transform_compose_cls_name = params_config.pop('transform_compose_cls', None)
      transform = build_transform(params_config.pop('transform_params', None), compose_cls=transform_compose_cls_name)
      target_transform_compose_cls_name = params_config.pop('target_transform_compose_cls', None)
      target_transform = build_transform(params_config.pop('target_transform_params', None),
                                         compose_cls=target_transform_compose_cls_name)
      transforms_compose_cls_name = params_config.pop('transforms_compose_cls', None)
      transforms = build_transform(params_config.pop('transforms_params', None), compose_cls=transforms_compose_cls_name)
      if 'loader' in params_config:
          loader_config = params_config.pop('loader')
          loader_type = loader_config['type']
          loader_params_config = loader_config.get('params', None)
          loader = get_sample_loader(loader_type) if loader_params_config is None \
              else get_sample_loader(loader_type, **loader_params_config)
          params_config['loader'] = loader
  
      # For datasets without target_transform
      if dataset_cls in (PhotoTour, Kinetics400, HMDB51, UCF101):
          return dataset_cls(transform=transform, **params_config)
      # For datasets with transforms
      if dataset_cls in (Cityscapes, CocoCaptions, CocoDetection, SBDataset, VOCSegmentation, VOCDetection):
          return dataset_cls(transform=transform, target_transform=target_transform,
                             transforms=transforms, **params_config)
      return dataset_cls(transform=transform, target_transform=target_transform, **params_config)
  
  
  def split_dataset(org_dataset, random_split_config, dataset_id, dataset_dict):
      org_dataset_length = len(org_dataset)
      logger.info('Splitting `{}` dataset ({} samples in total)'.format(dataset_id, org_dataset_length))
      lengths = random_split_config['lengths']
      total_length = sum(lengths)
      if total_length != org_dataset_length:
          lengths = [int((l / total_length) * org_dataset_length) for l in lengths]
          if len(lengths) > 1 and sum(lengths) != org_dataset_length:
              lengths[-1] = org_dataset_length - sum(lengths[:-1])
  
      manual_seed = random_split_config.get('generator_seed', None)
      sub_datasets = random_split(org_dataset, lengths) if manual_seed is None \
          else random_split(org_dataset, lengths, generator=torch.Generator().manual_seed(manual_seed))
      # Deep-copy dataset to configure transforms independently as dataset in Subset class is shallow-copied
      for sub_dataset in sub_datasets:
          sub_dataset.dataset = copy.deepcopy(sub_dataset.dataset)
  
      sub_splits_config = random_split_config['sub_splits']
      assert len(sub_datasets) == len(sub_splits_config), \
          'len(lengths) `{}` should be equal to len(sub_splits) `{}`'.format(len(sub_datasets), len(sub_splits_config))
      for sub_dataset, sub_split_params in zip(sub_datasets, sub_splits_config):
          sub_dataset_id = sub_split_params['dataset_id']
          logger.info('new dataset_id: `{}` ({} samples)'.format(sub_dataset_id, len(sub_dataset)))
          params_config = sub_split_params.copy()
          transform = build_transform(params_config.pop('transform_params', None))
          target_transform = build_transform(params_config.pop('transform_params', None))
          transforms = build_transform(params_config.pop('transforms_params', None))
          if hasattr(sub_dataset.dataset, 'transform') and transform is not None:
              sub_dataset.dataset.transform = transform
          if hasattr(sub_dataset.dataset, 'target_transform') and target_transform is not None:
              sub_dataset.dataset.target_transform = target_transform
          if hasattr(sub_dataset.dataset, 'transforms') and transforms is not None:
              sub_dataset.dataset.transforms = transforms
          dataset_dict[sub_dataset_id] = sub_dataset
  
  
  def get_dataset_dict(dataset_config):
      dataset_type = dataset_config['type']
      dataset_dict = dict()
      if dataset_type == 'cocodetect':
          dataset_splits_config = dataset_config['splits']
          for split_name in dataset_splits_config.keys():
              split_config = dataset_splits_config[split_name]
              is_segment = split_config.get('is_segment', False)
              compose_cls = CustomCompose if is_segment else None
              transforms = build_transform(split_config.get('transforms_params', None), compose_cls=compose_cls)
              dataset_dict[split_config['dataset_id']] =\
                  load_coco_dataset(split_config['images'], split_config['annotations'],
                                    split_config['annotated_only'], split_config.get('random_horizontal_flip', None),
                                    is_segment, transforms, split_config.get('jpeg_quality', None))
      elif dataset_type in DATASET_DICT:
          dataset_cls_or_func = DATASET_DICT[dataset_type]
          is_torchvision = dataset_type in torchvision.datasets.__dict__
          dataset_splits_config = dataset_config['splits']
          for split_name in dataset_splits_config.keys():
              st = time.time()
              logger.info('Loading {} data'.format(split_name))
              split_config = dataset_splits_config[split_name]
              org_dataset = get_torchvision_dataset(dataset_cls_or_func, split_config['params']) if is_torchvision \
                  else dataset_cls_or_func(**split_config['params'])
              print(dataset_cls_or_func,split_config['params'])
              dataset_id = split_config['dataset_id']
              random_split_config = split_config.get('random_split', None)
              if random_split_config is None:
                  dataset_dict[dataset_id] = org_dataset
              else:
                  split_dataset(org_dataset, random_split_config, dataset_id, dataset_dict)
              logger.info('dataset_id `{}`: {} sec'.format(dataset_id, time.time() - st))
      else:
          raise ValueError('dataset_type `{}` is not expected'.format(dataset_type))
      return dataset_dict
  
  
  def get_all_datasets(datasets_config):
      dataset_dict = dict()
      for dataset_name in datasets_config.keys():
          sub_dataset_dict = get_dataset_dict(datasets_config[dataset_name])
          dataset_dict.update(sub_dataset_dict)
      return dataset_dict
  
  
  def build_data_loader(dataset, data_loader_config, distributed, accelerator=None):
      num_workers = data_loader_config['num_workers']
      cache_dir_path = data_loader_config.get('cache_output', None)
      dataset_wrapper_config = data_loader_config.get('dataset_wrapper', None)
      if isinstance(dataset_wrapper_config, dict) and len(dataset_wrapper_config) > 0:
          dataset = get_dataset_wrapper(dataset_wrapper_config['name'], dataset, **dataset_wrapper_config['params'])
      elif cache_dir_path is not None:
          dataset = CacheableDataset(dataset, cache_dir_path, idx2subpath_func=default_idx2subpath)
      elif data_loader_config.get('requires_supp', False):
          dataset = BaseDatasetWrapper(dataset)
  
      sampler = DistributedSampler(dataset) if distributed and accelerator is None \
          else RandomSampler(dataset) if data_loader_config.get('random_sample', False) else SequentialSampler(dataset)
      batch_sampler_config = data_loader_config.get('batch_sampler', None)
      batch_sampler = None if batch_sampler_config is None \
          else get_batch_sampler(dataset, batch_sampler_config['type'], sampler, **batch_sampler_config['params'])
      collate_fn = get_collate_func(data_loader_config.get('collate_fn', None))
      drop_last = data_loader_config.get('drop_last', False)
      if batch_sampler is not None:
          return DataLoader(dataset, batch_sampler=batch_sampler, num_workers=num_workers,
                            collate_fn=collate_fn, drop_last=drop_last)
  
      batch_size = data_loader_config['batch_size']
      pin_memory = data_loader_config.get('pin_memory', True)
      return DataLoader(dataset, batch_size=batch_size, sampler=sampler,
                        num_workers=num_workers, collate_fn=collate_fn, pin_memory=pin_memory, drop_last=drop_last)
  
  
  def build_data_loaders(dataset_dict, data_loader_configs, distributed, accelerator=None):
      data_loader_list = list()
      for data_loader_config in data_loader_configs:
          dataset_id = data_loader_config.get('dataset_id', None)
          data_loader = None if dataset_id is None or dataset_id not in dataset_dict \
              else build_data_loader(dataset_dict[dataset_id], data_loader_config, distributed, accelerator)
          data_loader_list.append(data_loader)
      return data_loader_list
  