data/build/scenario.py

import enum
from functools import reduce
from typing import Dict, List, Tuple
import numpy as np
import copy
from utils.common.log import logger
from ..datasets.ab_dataset import ABDataset
from ..dataloader import FastDataLoader, InfiniteDataLoader, build_dataloader
from data import get_dataset


class DatasetMetaInfo:
    def __init__(self, name, 
                 known_classes_name_idx_map, unknown_class_idx):
        
        assert unknown_class_idx not in known_classes_name_idx_map.keys()
        
        self.name = name
        self.unknown_class_idx = unknown_class_idx
        self.known_classes_name_idx_map = known_classes_name_idx_map
        
    @property
    def num_classes(self):
        return len(self.known_classes_idx) + 1
        
        
class MergedDataset:
    def __init__(self, datasets: List[ABDataset]):
        self.datasets = datasets
        self.datasets_len = [len(i) for i in self.datasets]
        logger.info(f'create MergedDataset: len of datasets {self.datasets_len}')
        self.datasets_cum_len = np.cumsum(self.datasets_len)

    def __getitem__(self, idx):
        for i, cum_len in enumerate(self.datasets_cum_len):
            if idx < cum_len:
                return self.datasets[i][idx - sum(self.datasets_len[0: i])]
            
    def __len__(self):
        return sum(self.datasets_len)
    
    
class IndexReturnedDataset:
    def __init__(self, dataset: ABDataset):
        self.dataset = dataset
        
    def __getitem__(self, idx):
        res = self.dataset[idx]

        if isinstance(res, (tuple, list)):
            return (*res, idx)
        else:
            return res, idx
            
    def __len__(self):
        return len(self.dataset)
    

# class Scenario:
#     def __init__(self, config,
#                  source_datasets_meta_info: Dict[str, DatasetMetaInfo], target_datasets_meta_info: Dict[str, DatasetMetaInfo], 
#                  target_source_map: Dict[str, Dict[str, str]], 
#                  target_domains_order: List[str],
#                  source_datasets: Dict[str, Dict[str, ABDataset]], target_datasets: Dict[str, Dict[str, ABDataset]]):
        
#         self.__config = config
#         self.__source_datasets_meta_info = source_datasets_meta_info
#         self.__target_datasets_meta_info = target_datasets_meta_info
#         self.__target_source_map = target_source_map
#         self.__target_domains_order = target_domains_order
#         self.__source_datasets = source_datasets
#         self.__target_datasets = target_datasets
    
#     # 1. basic
#     def get_config(self):
#         return copy.deepcopy(self.__config)
    
#     def get_task_type(self):
#         return list(self.__source_datasets.values())[0]['train'].task_type
    
#     def get_num_classes(self):
#         known_classes_idx = []
#         unknown_classes_idx = []
#         for v in self.__source_datasets_meta_info.values():
#             known_classes_idx += list(v.known_classes_name_idx_map.values())
#             unknown_classes_idx += [v.unknown_class_idx]
#         for v in self.__target_datasets_meta_info.values():
#             known_classes_idx += list(v.known_classes_name_idx_map.values())
#             unknown_classes_idx += [v.unknown_class_idx]
#         unknown_classes_idx = [i for i in unknown_classes_idx if i is not None]
#         # print(known_classes_idx, unknown_classes_idx)
#         res = len(set(known_classes_idx)), len(set(unknown_classes_idx)), len(set(known_classes_idx + unknown_classes_idx))
#         # print(res)
#         assert res[0] + res[1] == res[2]
#         return res
      
#     def build_dataloader(self, dataset: ABDataset, batch_size: int, num_workers: int, infinite: bool, shuffle_when_finite: bool):
#         if infinite:
#             dataloader = InfiniteDataLoader(
#                 dataset, None, batch_size, num_workers=num_workers)
#         else:
#             dataloader = FastDataLoader(
#                 dataset, batch_size, num_workers, shuffle=shuffle_when_finite)

#         return dataloader
    
#     def build_sub_dataset(self, dataset: ABDataset, indexes: List[int]):
#         from ..data.datasets.dataset_split import _SplitDataset
#         dataset.dataset = _SplitDataset(dataset.dataset, indexes)
#         return dataset
    
#     def build_index_returned_dataset(self, dataset: ABDataset):
#         return IndexReturnedDataset(dataset)
        
#     # 2. source
#     def get_source_datasets_meta_info(self):
#         return self.__source_datasets_meta_info
    
#     def get_source_datasets_name(self):
#         return list(self.__source_datasets.keys())
    
#     def get_merged_source_dataset(self, split):
#         source_train_datasets = {n: d[split] for n, d in self.__source_datasets.items()}
#         return MergedDataset(list(source_train_datasets.values()))
    
#     def get_source_datasets(self, split):
#         source_train_datasets = {n: d[split] for n, d in self.__source_datasets.items()}
#         return source_train_datasets
    
#     # 3. target **domain**
#     # (do we need such API `get_ith_target_domain()`?)
#     def get_target_domains_meta_info(self):
#         return self.__source_datasets_meta_info
    
#     def get_target_domains_order(self):
#         return self.__target_domains_order
    
#     def get_corr_source_datasets_name_of_target_domain(self, target_domain_name):
#         return self.__target_source_map[target_domain_name]
    
#     def get_limited_target_train_dataset(self):
#         if len(self.__target_domains_order) > 1:
#             raise RuntimeError('this API is only for pass-in scenario in user-defined online DA algorithm')
#         return list(self.__target_datasets.values())[0]['train']
    
#     def get_target_domains_iterator(self, split):
#         for target_domain_index, target_domain_name in enumerate(self.__target_domains_order):
#             target_dataset = self.__target_datasets[target_domain_name]
#             target_domain_meta_info = self.__target_datasets_meta_info[target_domain_name]
            
#             yield target_domain_index, target_domain_name, target_dataset[split], target_domain_meta_info
    
#     # 4. permission management
#     def get_sub_scenario(self, source_datasets_name, source_splits, target_domains_order, target_splits):
#         def get_split(dataset, splits):
#             res = {}
#             for s, d in dataset.items():
#                 if s in splits:
#                     res[s] = d
#             return res
        
#         return Scenario(
#             config=self.__config,
#             source_datasets_meta_info={k: v for k, v in self.__source_datasets_meta_info.items() if k in source_datasets_name},
#             target_datasets_meta_info={k: v for k, v in self.__target_datasets_meta_info.items() if k in target_domains_order},
#             target_source_map={k: v for k, v in self.__target_source_map.items() if k in target_domains_order},
#             target_domains_order=target_domains_order,
#             source_datasets={k: get_split(v, source_splits) for k, v in self.__source_datasets.items() if k in source_datasets_name},
#             target_datasets={k: get_split(v, target_splits) for k, v in self.__target_datasets.items() if k in target_domains_order}
#         )
    
#     def get_only_source_sub_scenario_for_exp_tracker(self):
#         return self.get_sub_scenario(self.get_source_datasets_name(), ['train', 'val', 'test'], [], [])
    
#     def get_only_source_sub_scenario_for_alg(self):
#         return self.get_sub_scenario(self.get_source_datasets_name(), ['train'], [], [])
    
#     def get_one_da_sub_scenario_for_alg(self, target_domain_name):
#         return self.get_sub_scenario(self.get_corr_source_datasets_name_of_target_domain(target_domain_name), 
#                                      ['train', 'val'], [target_domain_name], ['train'])


# class Scenario:
#     def __init__(self, config,
                 
#                  offline_source_datasets_meta_info: Dict[str, DatasetMetaInfo], 
#                  offline_source_datasets: Dict[str, ABDataset],
                 
#                  online_datasets_meta_info: List[Tuple[Dict[str, DatasetMetaInfo], DatasetMetaInfo]],
#                  online_datasets: Dict[str, ABDataset],
#                  target_domains_order: List[str],
#                  target_source_map: Dict[str, Dict[str, str]],
                 
#                  num_classes: int):
        
#         self.config = config
        
#         self.offline_source_datasets_meta_info = offline_source_datasets_meta_info
#         self.offline_source_datasets = offline_source_datasets
        
#         self.online_datasets_meta_info = online_datasets_meta_info
#         self.online_datasets = online_datasets
        
#         self.target_domains_order = target_domains_order
#         self.target_source_map = target_source_map
        
#         self.num_classes = num_classes
        
#     def get_offline_source_datasets(self, split):
#         return {n: d[split] for n, d in self.offline_source_datasets.items()}
    
#     def get_offline_source_merged_dataset(self, split):
#         return MergedDataset([d[split] for d in self.offline_source_datasets.values()])
    
#     def get_online_current_corresponding_source_datasets(self, domain_index, split):
#         cur_target_domain_name = self.target_domains_order[domain_index]
#         cur_source_datasets_name = list(self.target_source_map[cur_target_domain_name].keys())
#         cur_source_datasets = {n: self.online_datasets[n + '|' + cur_target_domain_name][split] for n in cur_source_datasets_name}
#         return cur_source_datasets
    
#     def get_online_current_corresponding_merged_source_dataset(self, domain_index, split):
#         cur_target_domain_name = self.target_domains_order[domain_index]
#         cur_source_datasets_name = list(self.target_source_map[cur_target_domain_name].keys())
#         cur_source_datasets = {n: self.online_datasets[n + '|' + cur_target_domain_name][split] for n in cur_source_datasets_name}
#         return MergedDataset([d for d in cur_source_datasets.values()])
    
#     def get_online_current_target_dataset(self, domain_index, split):
#         cur_target_domain_name = self.target_domains_order[domain_index]
#         return self.online_datasets[cur_target_domain_name][split]
    
#     def build_dataloader(self, dataset: ABDataset, batch_size: int, num_workers: int, 
#                          infinite: bool, shuffle_when_finite: bool, to_iterator: bool):
#         if infinite:
#             dataloader = InfiniteDataLoader(
#                 dataset, None, batch_size, num_workers=num_workers)
#         else:
#             dataloader = FastDataLoader(
#                 dataset, batch_size, num_workers, shuffle=shuffle_when_finite)
            
#         if to_iterator:
#             dataloader = iter(dataloader)

#         return dataloader
    
#     def build_sub_dataset(self, dataset: ABDataset, indexes: List[int]):
#         from data.datasets.dataset_split import _SplitDataset
#         dataset.dataset = _SplitDataset(dataset.dataset, indexes)
#         return dataset
    
#     def build_index_returned_dataset(self, dataset: ABDataset):
#         return IndexReturnedDataset(dataset)
    
#     def get_config(self):
#         return copy.deepcopy(self.config)
    
#     def get_task_type(self):
#         return list(self.online_datasets.values())[0]['train'].task_type
    
#     def get_num_classes(self):
#         return self.num_classes
    

class Scenario:
    def __init__(self, config, all_datasets_ignore_classes_map, all_datasets_idx_map, target_domains_order, target_source_map, 
                 all_datasets_e2e_class_to_idx_map,
                 num_classes):
        self.config = config 
        self.all_datasets_ignore_classes_map = all_datasets_ignore_classes_map
        self.all_datasets_idx_map = all_datasets_idx_map
        self.target_domains_order = target_domains_order
        self.target_source_map = target_source_map
        self.all_datasets_e2e_class_to_idx_map = all_datasets_e2e_class_to_idx_map
        self.num_classes = num_classes
        self.cur_domain_index = 0

        logger.info(f'[scenario build] # classes: {num_classes}')
        logger.debug(f'[scenario build] idx map: {all_datasets_idx_map}')
        
    def to_json(self):
        return dict(
            config=self.config, all_datasets_ignore_classes_map=self.all_datasets_ignore_classes_map,
            all_datasets_idx_map=self.all_datasets_idx_map, target_domains_order=self.target_domains_order,
            target_source_map=self.target_source_map, 
            all_datasets_e2e_class_to_idx_map=self.all_datasets_e2e_class_to_idx_map,
            num_classes=self.num_classes
        )
        
    def __str__(self):
        return f'Scenario({self.to_json()})'

    def get_offline_datasets(self, transform=None):
        # make source datasets which contains all unioned classes
        res_offline_train_source_datasets_map = {}

        from .. import get_dataset
        data_dirs = self.config['data_dirs']

        source_datasets_name = self.config['source_datasets_name']
        res_source_datasets_map = {d: {split: get_dataset(d.split('|')[0], data_dirs[d.split('|')[0]], split, 
                                                      transform,
                                                      self.all_datasets_ignore_classes_map[d], self.all_datasets_idx_map[d]) 
                                   for split in ['train', 'val', 'test']} 
                               for d in self.all_datasets_ignore_classes_map.keys() if d.split('|')[0] in source_datasets_name}
        
        for source_dataset_name in self.config['source_datasets_name']:
            source_datasets = [v for k, v in res_source_datasets_map.items() if source_dataset_name in k]

            # how to merge idx map?
            # 35 79 97
            idx_maps = [d['train'].idx_map for d in source_datasets]
            ignore_classes_list = [d['train'].ignore_classes for d in source_datasets]
            
            union_idx_map = {}
            for idx_map in idx_maps:
                for k, v in idx_map.items():
                    if k not in union_idx_map:
                        union_idx_map[k] = v
                    else:
                        assert union_idx_map[k] == v

            union_ignore_classes = reduce(lambda res, cur: res & set(cur), ignore_classes_list, set(ignore_classes_list[0]))
            assert len(union_ignore_classes) + len(union_idx_map) == len(source_datasets[0]['train'].raw_classes)

            logger.info(f'[scenario build] {source_dataset_name} has {len(union_idx_map)} classes in offline training')
            
            d = source_dataset_name
            res_offline_train_source_datasets_map[d] = {split: get_dataset(d, data_dirs[d], split, 
                                                      transform,
                                                      union_ignore_classes, union_idx_map) 
                                   for split in ['train', 'val', 'test']} 
        
        return res_offline_train_source_datasets_map
    
    def get_offline_datasets_args(self):
        # make source datasets which contains all unioned classes
        res_offline_train_source_datasets_map = {}

        from .. import get_dataset
        data_dirs = self.config['data_dirs']

        source_datasets_name = self.config['source_datasets_name']
        res_source_datasets_map = {d: {split: get_dataset(d.split('|')[0], data_dirs[d.split('|')[0]], split, 
                                                      None,
                                                      self.all_datasets_ignore_classes_map[d], self.all_datasets_idx_map[d]) 
                                   for split in ['train', 'val', 'test']} 
                               for d in self.all_datasets_ignore_classes_map.keys() if d.split('|')[0] in source_datasets_name}
        
        for source_dataset_name in self.config['source_datasets_name']:
            source_datasets = [v for k, v in res_source_datasets_map.items() if source_dataset_name in k]

            # how to merge idx map?
            # 35 79 97
            idx_maps = [d['train'].idx_map for d in source_datasets]
            ignore_classes_list = [d['train'].ignore_classes for d in source_datasets]
            
            union_idx_map = {}
            for idx_map in idx_maps:
                for k, v in idx_map.items():
                    if k not in union_idx_map:
                        union_idx_map[k] = v
                    else:
                        assert union_idx_map[k] == v

            union_ignore_classes = reduce(lambda res, cur: res & set(cur), ignore_classes_list, set(ignore_classes_list[0]))
            assert len(union_ignore_classes) + len(union_idx_map) == len(source_datasets[0]['train'].raw_classes)

            logger.info(f'[scenario build] {source_dataset_name} has {len(union_idx_map)} classes in offline training')
            
            d = source_dataset_name
            res_offline_train_source_datasets_map[d] = {split: dict(d, data_dirs[d], split, 
                                                      None,
                                                      union_ignore_classes, union_idx_map) 
                                   for split in ['train', 'val', 'test']} 

        return res_offline_train_source_datasets_map

        # for d in source_datasets_name:
        #     source_dataset_with_max_num_classes = None
            
        #     for ed_name, ed in res_source_datasets_map.items():
        #         if not ed_name.startswith(d):
        #             continue
                
        #         if source_dataset_with_max_num_classes is None:
        #             source_dataset_with_max_num_classes = ed
        #             res_offline_train_source_datasets_map_names[d] = ed_name
                    
        #         if len(ed['train'].ignore_classes) < len(source_dataset_with_max_num_classes['train'].ignore_classes):
        #             source_dataset_with_max_num_classes = ed
        #             res_offline_train_source_datasets_map_names[d] = ed_name
                    
        #     res_offline_train_source_datasets_map[d] = source_dataset_with_max_num_classes

        # return res_offline_train_source_datasets_map
        
    def get_online_ith_domain_datasets_args_for_inference(self, domain_index):
        target_dataset_name = self.target_domains_order[domain_index]
        # dataset_name: Any, root_dir: Any, split: Any, transform: Any | None = None, ignore_classes: Any = [], idx_map: Any | None = None
        
        if 'MM-CityscapesDet' in self.target_domains_order or 'CityscapesDet' in self.target_domains_order or 'BaiduPersonDet' in self.target_domains_order:
            logger.info(f'use val split for inference test (only Det workload)')
            split = 'test'
        else:
            split = 'train'
        
        return dict(dataset_name=target_dataset_name, 
                    root_dir=self.config['data_dirs'][target_dataset_name], 
                    split=split, 
                    transform=None, 
                    ignore_classes=self.all_datasets_ignore_classes_map[target_dataset_name], 
                    idx_map=self.all_datasets_idx_map[target_dataset_name])
    
    def get_online_ith_domain_datasets_args_for_training(self, domain_index):
        target_dataset_name = self.target_domains_order[domain_index]
        source_datasets_name = list(self.target_source_map[target_dataset_name].keys())

        res = {}
        # dataset_name: Any, root_dir: Any, split: Any, transform: Any | None = None, ignore_classes: Any = [], idx_map: Any | None = None
        res[target_dataset_name] = {split: dict(dataset_name=target_dataset_name, 
                    root_dir=self.config['data_dirs'][target_dataset_name], 
                    split=split, 
                    transform=None, 
                    ignore_classes=self.all_datasets_ignore_classes_map[target_dataset_name], 
                    idx_map=self.all_datasets_idx_map[target_dataset_name]) for split in ['train', 'val']}
        for d in source_datasets_name:
            res[d] = {split: dict(dataset_name=d, 
                    root_dir=self.config['data_dirs'][d], 
                    split=split, 
                    transform=None, 
                    ignore_classes=self.all_datasets_ignore_classes_map[d + '|' + target_dataset_name], 
                    idx_map=self.all_datasets_idx_map[d + '|' + target_dataset_name]) for split in ['train', 'val']}
        
        return res
    
    def get_online_cur_domain_datasets_args_for_inference(self):
        return self.get_online_ith_domain_datasets_args_for_inference(self.cur_domain_index)
    
    def get_online_cur_domain_datasets_args_for_training(self):
        return self.get_online_ith_domain_datasets_args_for_training(self.cur_domain_index)
    
    def get_online_cur_domain_datasets_for_training(self, transform=None):
        res = {}
        datasets_args = self.get_online_ith_domain_datasets_args_for_training(self.cur_domain_index)
        for dataset_name, dataset_args in datasets_args.items():
            res[dataset_name] = {}
            for split, args in dataset_args.items():
                if transform is not None:
                    args['transform'] = transform
                dataset = get_dataset(**args)
                res[dataset_name][split] = dataset
        return res
    
    def get_online_cur_domain_datasets_for_inference(self, transform=None):
        datasets_args = self.get_online_ith_domain_datasets_args_for_inference(self.cur_domain_index)
        if transform is not None:
            datasets_args['transform'] = transform
        return get_dataset(**datasets_args)
    
    def get_online_cur_domain_samples_for_training(self, num_samples, transform=None, collate_fn=None):
        dataset = self.get_online_cur_domain_datasets_for_training(transform=transform)
        dataset = dataset[self.target_domains_order[self.cur_domain_index]]['train']
        return next(iter(build_dataloader(dataset, num_samples, 0, True, None, collate_fn=collate_fn)))[0]

    def next_domain(self):
        self.cur_domain_index += 1