from typing import Dict, List, Optional, Type, Union from ..datasets.ab_dataset import ABDataset # from benchmark.data.visualize import visualize_classes_in_object_detection # from benchmark.scenario.val_domain_shift import get_val_domain_shift_transform from ..dataset import get_dataset import copy from torchvision.transforms import Compose from .merge_alias import merge_the_same_meaning_classes from ..datasets.registery import static_dataset_registery # some legacy aliases of variables: # ignore_classes == discarded classes # private_classes == unknown classes in partial / open-set / universal DA def _merge_the_same_meaning_classes(classes_info_of_all_datasets): final_classes_of_all_datasets, rename_map = merge_the_same_meaning_classes(classes_info_of_all_datasets) return final_classes_of_all_datasets, rename_map def _find_ignore_classes_when_sources_as_to_target_b(as_classes: List[List[str]], b_classes: List[str], da_mode): thres = {'da': 3, 'partial_da': 2, 'open_set_da': 1, 'universal_da': 0}[da_mode] from functools import reduce a_classes = reduce(lambda res, cur: res | set(cur), as_classes, set()) if set(a_classes) == set(b_classes): # a is equal to b, normal # 1. no ignore classes; 2. match class idx a_ignore_classes, b_ignore_classes = [], [] elif set(a_classes) > set(b_classes): # a contains b, partial a_ignore_classes, b_ignore_classes = [], [] if thres == 3 or thres == 1: # ignore extra classes in a a_ignore_classes = set(a_classes) - set(b_classes) elif set(a_classes) < set(b_classes): # a is contained by b, open set a_ignore_classes, b_ignore_classes = [], [] if thres == 3 or thres == 2: # ignore extra classes in b b_ignore_classes = set(b_classes) - set(a_classes) elif len(set(a_classes) & set(b_classes)) > 0: a_ignore_classes, b_ignore_classes = [], [] if thres == 3: a_ignore_classes = set(a_classes) - (set(a_classes) & set(b_classes)) b_ignore_classes = set(b_classes) - (set(a_classes) & set(b_classes)) elif thres == 2: b_ignore_classes = set(b_classes) - (set(a_classes) & set(b_classes)) elif thres == 1: a_ignore_classes = set(a_classes) - (set(a_classes) & set(b_classes)) else: return None # a has no intersection with b, none as_ignore_classes = [list(set(a_classes) & set(a_ignore_classes)) for a_classes in as_classes] return as_ignore_classes, list(b_ignore_classes) def _find_private_classes_when_sources_as_to_target_b(as_classes: List[List[str]], b_classes: List[str], da_mode): thres = {'da': 3, 'partial_da': 2, 'open_set_da': 1, 'universal_da': 0}[da_mode] from functools import reduce a_classes = reduce(lambda res, cur: res | set(cur), as_classes, set()) if set(a_classes) == set(b_classes): # a is equal to b, normal # 1. no ignore classes; 2. match class idx a_private_classes, b_private_classes = [], [] elif set(a_classes) > set(b_classes): # a contains b, partial a_private_classes, b_private_classes = [], [] # if thres == 2 or thres == 0: # ignore extra classes in a # a_private_classes = set(a_classes) - set(b_classes) # if thres == 0: # ignore extra classes in a # a_private_classes = set(a_classes) - set(b_classes) elif set(a_classes) < set(b_classes): # a is contained by b, open set a_private_classes, b_private_classes = [], [] if thres == 1 or thres == 0: # ignore extra classes in b b_private_classes = set(b_classes) - set(a_classes) elif len(set(a_classes) & set(b_classes)) > 0: a_private_classes, b_private_classes = [], [] if thres == 0: # a_private_classes = set(a_classes) - (set(a_classes) & set(b_classes)) b_private_classes = set(b_classes) - (set(a_classes) & set(b_classes)) elif thres == 1: b_private_classes = set(b_classes) - (set(a_classes) & set(b_classes)) elif thres == 2: # a_private_classes = set(a_classes) - (set(a_classes) & set(b_classes)) pass else: return None # a has no intersection with b, none return list(b_private_classes) class _ABDatasetMetaInfo: def __init__(self, name, classes, task_type, object_type, class_aliases, shift_type): self.name = name self.classes = classes self.class_aliases = class_aliases self.shift_type = shift_type self.task_type = task_type self.object_type = object_type def _get_dist_shift_type_when_source_a_to_target_b(a: _ABDatasetMetaInfo, b: _ABDatasetMetaInfo): if b.shift_type is None: return 'Dataset Shifts' if a.name in b.shift_type.keys(): return b.shift_type[a.name] mid_dataset_name = list(b.shift_type.keys())[0] mid_dataset_meta_info = _ABDatasetMetaInfo(mid_dataset_name, *static_dataset_registery[mid_dataset_name][1:]) return _get_dist_shift_type_when_source_a_to_target_b(a, mid_dataset_meta_info) + ' + ' + list(b.shift_type.values())[0] def _handle_all_datasets_v2(source_datasets: List[_ABDatasetMetaInfo], target_datasets: List[_ABDatasetMetaInfo], da_mode): # 1. merge the same meaning classes classes_info_of_all_datasets = { d.name: (d.classes, d.class_aliases) for d in source_datasets + target_datasets } final_classes_of_all_datasets, rename_map = _merge_the_same_meaning_classes(classes_info_of_all_datasets) all_datasets_classes = copy.deepcopy(final_classes_of_all_datasets) # print(all_datasets_known_classes) # 2. find ignored classes according to DA mode # source_datasets_ignore_classes, target_datasets_ignore_classes = {d.name: [] for d in source_datasets}, \ # {d.name: [] for d in target_datasets} # source_datasets_private_classes, target_datasets_private_classes = {d.name: [] for d in source_datasets}, \ # {d.name: [] for d in target_datasets} target_source_relationship_map = {td.name: {} for td in target_datasets} # source_target_relationship_map = {sd.name: [] for sd in source_datasets} # 1. construct target_source_relationship_map for sd in source_datasets:#sd和td使列表中每一个元素（类）的实例 for td in target_datasets: sc = all_datasets_classes[sd.name] tc = all_datasets_classes[td.name] if len(set(sc) & set(tc)) == 0:#只保留有相似类别的源域和目标域 continue target_source_relationship_map[td.name][sd.name] = _get_dist_shift_type_when_source_a_to_target_b(sd, td) # print(target_source_relationship_map) # exit() source_datasets_ignore_classes = {} for td_name, v1 in target_source_relationship_map.items(): for sd_name, v2 in v1.items(): source_datasets_ignore_classes[sd_name + '|' + td_name] = [] target_datasets_ignore_classes = {d.name: [] for d in target_datasets} target_datasets_private_classes = {d.name: [] for d in target_datasets} # 保证对于每个目标域上的DA都符合给定的label shift # 所以不同目标域就算对应同一个源域，该源域也可能不相同 for td_name, v1 in target_source_relationship_map.items(): sd_names = list(v1.keys()) sds_classes = [all_datasets_classes[sd_name] for sd_name in sd_names] td_classes = all_datasets_classes[td_name] ss_ignore_classes, t_ignore_classes = _find_ignore_classes_when_sources_as_to_target_b(sds_classes, td_classes, da_mode)#根据DA方式不同产生ignore_classes t_private_classes = _find_private_classes_when_sources_as_to_target_b(sds_classes, td_classes, da_mode) for sd_name, s_ignore_classes in zip(sd_names, ss_ignore_classes): source_datasets_ignore_classes[sd_name + '|' + td_name] = s_ignore_classes target_datasets_ignore_classes[td_name] = t_ignore_classes target_datasets_private_classes[td_name] = t_private_classes source_datasets_ignore_classes = {k: sorted(set(v), key=v.index) for k, v in source_datasets_ignore_classes.items()} target_datasets_ignore_classes = {k: sorted(set(v), key=v.index) for k, v in target_datasets_ignore_classes.items()} target_datasets_private_classes = {k: sorted(set(v), key=v.index) for k, v in target_datasets_private_classes.items()} # for k, v in source_datasets_ignore_classes.items(): # print(k, len(v)) # print() # for k, v in target_datasets_ignore_classes.items(): # print(k, len(v)) # print() # for k, v in target_datasets_private_classes.items(): # print(k, len(v)) # print() # print(source_datasets_private_classes, target_datasets_private_classes) # 3. reparse classes idx # 3.1. agg all used classes # all_used_classes = [] # all_datasets_private_class_idx_map = {} # source_datasets_classes_idx_map = {} # for td_name, v1 in target_source_relationship_map.items(): # for sd_name, v2 in v1.items(): # source_datasets_classes_idx_map[sd_name + '|' + td_name] = [] # target_datasets_classes_idx_map = {} global_idx = 0 all_used_classes_idx_map = {} # all_datasets_known_classes = {d: [] for d in final_classes_of_all_datasets.keys()} for dataset_name, classes in all_datasets_classes.items(): if dataset_name not in target_datasets_ignore_classes.keys(): ignore_classes = [0] * 100000 for sn, sic in source_datasets_ignore_classes.items(): if sn.startswith(dataset_name): if len(sic) < len(ignore_classes): ignore_classes = sic else: ignore_classes = target_datasets_ignore_classes[dataset_name] private_classes = [] \ if dataset_name not in target_datasets_ignore_classes.keys() else target_datasets_private_classes[dataset_name] for c in classes: if c not in ignore_classes and c not in all_used_classes_idx_map.keys() and c not in private_classes: all_used_classes_idx_map[c] = global_idx global_idx += 1 # print(all_used_classes_idx_map) # dataset_private_class_idx_offset = 0 target_private_class_idx = global_idx target_datasets_private_class_idx = {d: None for d in target_datasets_private_classes.keys()} for dataset_name, classes in final_classes_of_all_datasets.items(): if dataset_name not in target_datasets_private_classes.keys(): continue # ignore_classes = target_datasets_ignore_classes[dataset_name] private_classes = target_datasets_private_classes[dataset_name] # private_classes = [] \ # if dataset_name in source_datasets_private_classes.keys() else target_datasets_private_classes[dataset_name] # for c in classes: # if c not in ignore_classes and c not in all_used_classes_idx_map.keys() and c in private_classes: # all_used_classes_idx_map[c] = global_idx + dataset_private_class_idx_offset if len(private_classes) > 0: # all_datasets_private_class_idx[dataset_name] = global_idx + dataset_private_class_idx_offset # dataset_private_class_idx_offset += 1 # if dataset_name in source_datasets_private_classes.keys(): # if source_private_class_idx is None: # source_private_class_idx = global_idx if target_private_class_idx is None else target_private_class_idx + 1 # all_datasets_private_class_idx[dataset_name] = source_private_class_idx # else: # if target_private_class_idx is None: # target_private_class_idx = global_idx if source_private_class_idx is None else source_private_class_idx + 1 # all_datasets_private_class_idx[dataset_name] = target_private_class_idx target_datasets_private_class_idx[dataset_name] = target_private_class_idx target_private_class_idx += 1 # all_used_classes = sorted(set(all_used_classes), key=all_used_classes.index) # all_used_classes_idx_map = {c: i for i, c in enumerate(all_used_classes)} # print('rename_map', rename_map) # 3.2 raw_class -> rename_map[raw_classes] -> all_used_classes_idx_map all_datasets_e2e_idx_map = {} all_datasets_e2e_class_to_idx_map = {} for td_name, v1 in target_source_relationship_map.items(): sd_names = list(v1.keys()) sds_classes = [all_datasets_classes[sd_name] for sd_name in sd_names] td_classes = all_datasets_classes[td_name] for sd_name, sd_classes in zip(sd_names, sds_classes): cur_e2e_idx_map = {} cur_e2e_class_to_idx_map = {} for raw_ci, raw_c in enumerate(sd_classes): renamed_c = raw_c if raw_c not in rename_map[dataset_name] else rename_map[dataset_name][raw_c] ignore_classes = source_datasets_ignore_classes[sd_name + '|' + td_name] if renamed_c in ignore_classes: continue idx = all_used_classes_idx_map[renamed_c] cur_e2e_idx_map[raw_ci] = idx cur_e2e_class_to_idx_map[raw_c] = idx all_datasets_e2e_idx_map[sd_name + '|' + td_name] = cur_e2e_idx_map all_datasets_e2e_class_to_idx_map[sd_name + '|' + td_name] = cur_e2e_class_to_idx_map cur_e2e_idx_map = {} cur_e2e_class_to_idx_map = {} for raw_ci, raw_c in enumerate(td_classes): renamed_c = raw_c if raw_c not in rename_map[dataset_name] else rename_map[dataset_name][raw_c] ignore_classes = target_datasets_ignore_classes[td_name] if renamed_c in ignore_classes: continue if renamed_c in target_datasets_private_classes[td_name]: idx = target_datasets_private_class_idx[td_name] else: idx = all_used_classes_idx_map[renamed_c] cur_e2e_idx_map[raw_ci] = idx cur_e2e_class_to_idx_map[raw_c] = idx all_datasets_e2e_idx_map[td_name] = cur_e2e_idx_map all_datasets_e2e_class_to_idx_map[td_name] = cur_e2e_class_to_idx_map all_datasets_ignore_classes = {**source_datasets_ignore_classes, **target_datasets_ignore_classes} # all_datasets_private_classes = {**source_datasets_private_classes, **target_datasets_private_classes} classes_idx_set = [] for d, m in all_datasets_e2e_class_to_idx_map.items(): classes_idx_set += list(m.values()) classes_idx_set = set(classes_idx_set) num_classes = len(classes_idx_set) return all_datasets_ignore_classes, target_datasets_private_classes, \ all_datasets_e2e_idx_map, all_datasets_e2e_class_to_idx_map, target_datasets_private_class_idx, \ target_source_relationship_map, rename_map, num_classes def _build_scenario_info_v2( source_datasets_name: List[str], target_datasets_order: List[str], da_mode: str ): assert da_mode in ['close_set', 'partial', 'open_set', 'universal'] da_mode = {'close_set': 'da', 'partial': 'partial_da', 'open_set': 'open_set_da', 'universal': 'universal_da'}[da_mode] source_datasets_meta_info = [_ABDatasetMetaInfo(d, *static_dataset_registery[d][1:]) for d in source_datasets_name]#获知对应的名字和对应属性，要添加数据集时，直接register就行 target_datasets_meta_info = [_ABDatasetMetaInfo(d, *static_dataset_registery[d][1:]) for d in list(set(target_datasets_order))] all_datasets_ignore_classes, target_datasets_private_classes, \ all_datasets_e2e_idx_map, all_datasets_e2e_class_to_idx_map, target_datasets_private_class_idx, \ target_source_relationship_map, rename_map, num_classes \ = _handle_all_datasets_v2(source_datasets_meta_info, target_datasets_meta_info, da_mode) return all_datasets_ignore_classes, target_datasets_private_classes, \ all_datasets_e2e_idx_map, all_datasets_e2e_class_to_idx_map, target_datasets_private_class_idx, \ target_source_relationship_map, rename_map, num_classes def build_scenario_manually_v2( source_datasets_name: List[str], target_datasets_order: List[str], da_mode: str, data_dirs: Dict[str, str], # transforms: Optional[Dict[str, Compose]] = None ): configs = copy.deepcopy(locals())#返回当前局部变量 source_datasets_meta_info = [_ABDatasetMetaInfo(d, *static_dataset_registery[d][1:]) for d in source_datasets_name] target_datasets_meta_info = [_ABDatasetMetaInfo(d, *static_dataset_registery[d][1:]) for d in list(set(target_datasets_order))] all_datasets_ignore_classes, target_datasets_private_classes, \ all_datasets_e2e_idx_map, all_datasets_e2e_class_to_idx_map, target_datasets_private_class_idx, \ target_source_relationship_map, rename_map, num_classes \ = _build_scenario_info_v2(source_datasets_name, target_datasets_order, da_mode) # from rich.console import Console # console = Console(width=10000) # def print_obj(_o): # # import pprint # # s = pprint.pformat(_o, width=140, compact=True) # console.print(_o) # console.print('configs:', style='bold red') # print_obj(configs) # console.print('renamed classes:', style='bold red') # print_obj(rename_map) # console.print('discarded classes:', style='bold red') # print_obj(all_datasets_ignore_classes) # console.print('unknown classes:', style='bold red') # print_obj(target_datasets_private_classes) # console.print('class to index map:', style='bold red') # print_obj(all_datasets_e2e_class_to_idx_map) # console.print('index map:', style='bold red') # print_obj(all_datasets_e2e_idx_map) # console = Console() # # console.print('class distribution:', style='bold red') # # class_dist = { # # k: { # # '#known classes': len(all_datasets_known_classes[k]), # # '#unknown classes': len(all_datasets_private_classes[k]), # # '#discarded classes': len(all_datasets_ignore_classes[k]) # # } for k in all_datasets_ignore_classes.keys() # # } # # print_obj(class_dist) # console.print('corresponding sources of each target:', style='bold red') # print_obj(target_source_relationship_map) # return # res_source_datasets_map = {d: {split: get_dataset(d, data_dirs[d], split, getattr(transforms, d, None), # all_datasets_ignore_classes[d], all_datasets_e2e_idx_map[d]) # for split in ['train', 'val', 'test']} # for d in source_datasets_name} # res_target_datasets_map = {d: {'train': get_num_limited_dataset(get_dataset(d, data_dirs[d], 'test', getattr(transforms, d, None), # all_datasets_ignore_classes[d], all_datasets_e2e_idx_map[d]), # num_samples_in_each_target_domain), # 'test': get_dataset(d, data_dirs[d], 'test', getattr(transforms, d, None), # all_datasets_ignore_classes[d], all_datasets_e2e_idx_map[d]) # } # for d in list(set(target_datasets_order))} # res_source_datasets_map = {d: {split: get_dataset(d.split('|')[0], data_dirs[d.split('|')[0]], split, # getattr(transforms, d.split('|')[0], None), # all_datasets_ignore_classes[d], all_datasets_e2e_idx_map[d]) # for split in ['train', 'val', 'test']} # for d in all_datasets_ignore_classes.keys() if d.split('|')[0] in source_datasets_name} # from functools import reduce # res_offline_train_source_datasets_map = {} # res_offline_train_source_datasets_map_names = {} # 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 # res_target_datasets_map = {d: {split: get_dataset(d, data_dirs[d], split, getattr(transforms, d, None), # all_datasets_ignore_classes[d], all_datasets_e2e_idx_map[d]) # for split in ['train', 'val', 'test']} # for d in list(set(target_datasets_order))} from .scenario import Scenario, DatasetMetaInfo # test_scenario = Scenario( # config=configs, # offline_source_datasets_meta_info={ # d: DatasetMetaInfo(d, # {k: v for k, v in all_datasets_e2e_class_to_idx_map[res_offline_train_source_datasets_map_names[d]].items()}, # None) # for d in source_datasets_name # }, # offline_source_datasets={d: res_offline_train_source_datasets_map[d] for d in source_datasets_name}, # online_datasets_meta_info=[ # ( # {sd + '|' + d: DatasetMetaInfo(d, # {k: v for k, v in all_datasets_e2e_class_to_idx_map[sd + '|' + d].items()}, # None) # for sd in target_source_relationship_map[d].keys()}, # DatasetMetaInfo(d, # {k: v for k, v in all_datasets_e2e_class_to_idx_map[d].items() if k not in target_datasets_private_classes[d]}, # target_datasets_private_class_idx[d]) # ) # for d in target_datasets_order # ], # online_datasets={**res_source_datasets_map, **res_target_datasets_map}, # target_domains_order=target_datasets_order, # target_source_map=target_source_relationship_map, # num_classes=num_classes # ) import os os.environ['_ZQL_NUMC'] = str(num_classes) test_scenario = Scenario(config=configs, all_datasets_ignore_classes_map=all_datasets_ignore_classes, all_datasets_idx_map=all_datasets_e2e_idx_map, target_domains_order=target_datasets_order, target_source_map=target_source_relationship_map, all_datasets_e2e_class_to_idx_map=all_datasets_e2e_class_to_idx_map, num_classes=num_classes) return test_scenario if __name__ == '__main__': test_scenario = build_scenario_manually_v2(['CIFAR10', 'SVHN'], ['STL10', 'MNIST', 'STL10', 'USPS', 'MNIST', 'STL10'], 'close_set') print(test_scenario.num_classes)