Upload 1912 files

.gitattributes

@@ -33,3 +33,10 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+data/datasets/visual_question_answering/generate_c_image/imagenet_c/frost/frost1.png filter=lfs diff=lfs merge=lfs -text
+data/datasets/visual_question_answering/generate_c_image/robustness-master/assets/spatter.gif filter=lfs diff=lfs merge=lfs -text
+data/datasets/visual_question_answering/generate_c_image/robustness-master/assets/tilt.gif filter=lfs diff=lfs merge=lfs -text
+data/datasets/visual_question_answering/generate_c_image/robustness-master/assets/translate.gif filter=lfs diff=lfs merge=lfs -text
+data/datasets/visual_question_answering/generate_c_image/robustness-master/ImageNet-C/create_c/frost1.png filter=lfs diff=lfs merge=lfs -text
+data/datasets/visual_question_answering/generate_c_image/robustness-master/ImageNet-C/imagenet_c/imagenet_c/frost/frost1.png filter=lfs diff=lfs merge=lfs -text
+new_results.png filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,5 @@
+results
+logs
+entry_model
+__pycache__
+backup_codes

README.md

@@ -1,10 +1,5 @@
 ---
-title: EdgeTA
-emoji: 📊
-colorFrom: green
-colorTo: red
+title: EdgeFM
 sdk: static
-pinned: false
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+app_file: index.html
+---

data/README.md

@@ -0,0 +1,94 @@
+## How to implement a dataset?
+
+For example, we want to implement a image classification dataset.
+
+1. create a file in corresponding directory, i.e. `benchmark/data/datasets/image_classification`
+
+2. create a class (inherited from `benchmark.data.datasets.ab_dataset.ABDataset`), e.g. `class YourDataset(ABDataset)`
+
+3. register your dataset with `benchmark.data.datasets.registry.dataset_register(name, classes, classes_aliases)`, which represents the name of your dataset, the classes of your dataset, and the possible aliases of the classes. Examples refer to `benchmark/data/datasets/image_classification/cifar10.py` or other files. 
+
+   Note that the order of `classes` must match the indexes. For example, `classes` of MNIST must be `['0', '1', '2', ..., '9']`, which means 0-th class is '0', 1-st class is '1', 2-nd class is '2', ...; `['1', '2', '0', ...]` is not correct because 0-th class is not '1' and 1-st class is not '2'.
+
+   How to get `classes` of a dataset? For PyTorch built-in dataset (CIFAR10, MNIST, ...) and general dataset build by `ImageFolder`, you can initialize it (e.g. `dataset = CIFAR10(...)`) and get its classes by `dataset.classes`. 
+
+   ```python
+   # How to get classes in CIFAR10?
+   from torchvision.datasets import CIFAR10
+   dataset = CIFAR10(...)
+   print(dataset.classes)
+   # copy this output to @dataset_register(classes=<what you copied>)
+   
+   # it's not recommended to dynamically get classes, e.g.:
+   # this works but runs slowly!
+   from torchvision.datasets import CIFAR10 as RawCIFAR10
+   dataset = RawCIFAR10(...)
+   
+   @dataset_register(
+   	name='CIFAR10',
+       classes=dataset.classes
+   )
+   class CIFAR10(ABDataset):
+       # ...
+   ```
+
+   For object detection dataset, you can read the annotation JSON file and find `categories` information in it.
+
+4. implement abstract function `create_dataset(self, root_dir: str, split: str, transform: Optional[Compose], classes: List[str], ignore_classes: List[str], idx_map: Optional[Dict[int, int]])`.
+
+   Arguments:
+
+   - `root_dir`: the location of data
+   - `split`: `train / val / test`
+   - `transform`: preprocess function in `torchvision.transforms`
+   - `classes`: the same value with `dataset_register.classes`
+   - `ignore_classes`: **classes should be discarded. You should remove images which belong to these ignore classes.**
+   - `idx_map`: **map the original class index to new class index. For example, `{0: 2}` means the index of 0-th class will be 2 instead of 0. You should implement this by modifying the stored labels in the original dataset. **
+
+   You should do five things in this function:
+
+   1. if no user-defined transform is passed, you should implemented the default transform
+   2. create the original dataset
+   3. remove ignored classes in the original dataset if there are ignored classes
+   4. map the original class index to new class index if there is index map
+   5. split the original dataset to train / val / test dataset. If there's no val dataset in original dataset (e.g. DomainNetReal), you should split the original dataset to train / val / test dataset. If there's already val dataset in original dataset (e.g. CIFAR10 and ImageNet), regard the original val dataset as test dataset, and split the original train dataset into train / val dataset. Details just refer to existed files.
+
+Example (`benchmark/data/datasets/image_classification/cifar10.py`):
+
+```python
+@dataset_register(
+    name='CIFAR10', 
+    # means in the original CIFAR10, 0-th class is airplane, 1-st class is automobile, ...
+    classes=['airplane', 'automobile', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck'], 
+    # means 'automobile' and 'car' are the same thing actually
+    class_aliases=[['automobile', 'car']] 
+)
+class CIFAR10(ABDataset):    
+    def create_dataset(self, root_dir: str, split: str, transform: Optional[Compose], 
+                       classes: List[str], ignore_classes: List[str], idx_map: Optional[Dict[int, int]]):
+        # 1. if no user-defined transform is passed, you should implemented the default transform
+        if transform is None:
+            transform = cifar_like_image_train_aug() if split == 'train' else cifar_like_image_test_aug()
+        # 2. create the original dataset
+        dataset = RawCIFAR10(root_dir, split != 'test', transform=transform, download=True)
+        
+        # 3. remove ignored classes in the original dataset if there are ignored classes
+        dataset.targets = np.asarray(dataset.targets)
+        if len(ignore_classes) > 0: 
+            for ignore_class in ignore_classes:
+                dataset.data = dataset.data[dataset.targets != classes.index(ignore_class)]
+                dataset.targets = dataset.targets[dataset.targets != classes.index(ignore_class)]
+        
+        # 4. map the original class index to new class index if there is index map
+        if idx_map is not None:
+            for ti, t in enumerate(dataset.targets):
+                dataset.targets[ti] = idx_map[t]
+        
+        # 5. split the original dataset to train / val / test dataset.
+        # there is not val dataset in CIFAR10 dataset, so we split the val dataset from the train dataset.
+        if split != 'test':
+            dataset = train_val_split(dataset, split)
+        return dataset
+```
+
+After implementing a new dataset, you can create a test file in `example` and load the dataset by `benchmark.data.dataset.get_dataset()`. Try using this dataset to ensure it works. (Example: `example/1.py`)

data/__init__.py

@@ -0,0 +1,14 @@
+from .dataset import get_dataset
+from .build.build import build_scenario_manually_v2 as build_scenario
+from .dataloader import build_dataloader
+from .build.scenario import IndexReturnedDataset, MergedDataset
+from .datasets.ab_dataset import ABDataset
+from .build.scenario import Scenario
+
+from .build_cl.build import build_cl_scenario
+from .build_cl.scenario import Scenario as CLScenario
+
+from .build_gen.build import build_scenario_manually_v2 as build_gen_scenario
+from .build_gen.scenario import Scenario as GenScenario
+
+from .datasets.dataset_split import split_dataset

data/build/build.py

@@ -0,0 +1,495 @@
+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)
+    

data/build/merge_alias.py

@@ -0,0 +1,106 @@
+from re import L
+from typing import Dict, List
+from collections import Counter
+
+    
+def grouping(bondlist):
+    # reference: https://blog.csdn.net/YnagShanwen/article/details/111344386
+    groups = [] 
+    break1 = False    
+    while bondlist:
+        pair1 = bondlist.pop(0)        
+        a = 11111
+        b = 10000
+        while b != a:
+            a = b
+            for atomid in pair1:
+                for i,pair2 in enumerate(bondlist):            
+                    if atomid in pair2:
+                        pair1 = pair1 + pair2
+                        bondlist.pop(i)
+                        if not bondlist:
+                            break1 = True
+                        break
+                if break1:
+                    break
+            b = len(pair1)
+        groups.append(pair1)
+    return groups
+
+
+def build_semantic_class_info(classes: List[str], aliases: List[List[str]]):
+    res = []
+    for c in classes:
+        # print(res)
+        if len(aliases) == 0:
+            res += [[c]]
+        else:
+            find_alias = False
+            for alias in aliases:
+                if c in alias:
+                    res += [alias]
+                    find_alias = True
+                    break
+            if not find_alias:
+                res += [[c]]
+    # print(classes, res)
+    return res
+    
+
+def merge_the_same_meaning_classes(classes_info_of_all_datasets):
+    # print(classes_info_of_all_datasets)
+
+    semantic_classes_of_all_datasets = []
+    all_aliases = []
+    for classes, aliases in classes_info_of_all_datasets.values():
+        all_aliases += aliases
+    for classes, aliases in classes_info_of_all_datasets.values():
+        semantic_classes_of_all_datasets += build_semantic_class_info(classes, all_aliases)
+        
+    # print(semantic_classes_of_all_datasets)
+    
+    grouped_classes_of_all_datasets = grouping(semantic_classes_of_all_datasets)#匹配过后的数据
+   
+    # print(grouped_classes_of_all_datasets)
+    
+    # final_grouped_classes_of_all_datasets = [Counter(c).most_common()[0][0] for c in grouped_classes_of_all_datasets]
+    # use most common class name; if the same common, use shortest class name!
+    final_grouped_classes_of_all_datasets = []
+    for c in grouped_classes_of_all_datasets:
+        counter = Counter(c).most_common()
+        max_times = counter[0][1]
+        candidate_class_names = []
+        for item, times in counter:
+            if times < max_times:
+                break
+            candidate_class_names += [item]
+        candidate_class_names.sort(key=lambda x: len(x))
+        
+        final_grouped_classes_of_all_datasets += [candidate_class_names[0]]
+    res = {}
+    res_map = {d: {} for d in classes_info_of_all_datasets.keys()}
+    
+    for dataset_name, (classes, _) in classes_info_of_all_datasets.items():
+        final_classes = []
+        for c in classes:
+            for grouped_names, final_name in zip(grouped_classes_of_all_datasets, final_grouped_classes_of_all_datasets):
+                if c in grouped_names:
+                    final_classes += [final_name]
+                    if final_name != c:
+                        res_map[dataset_name][c] = final_name
+                    break
+        res[dataset_name] = sorted(set(final_classes), key=final_classes.index)
+    return res, res_map
+
+
+if __name__ == '__main__':
+    cifar10_classes = ['airplane', 'automobile', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck']
+    cifar10_aliases = [['automobile', 'car']]
+    stl10_classes = ['airplane', 'bird', 'car', 'cat', 'deer', 'dog', 'horse', 'monkey', 'ship', 'truck']
+
+    final_classes_of_all_datasets, rename_map = merge_the_same_meaning_classes({
+        'CIFAR10': (cifar10_classes, cifar10_aliases),
+        'STL10': (stl10_classes, [])
+    })
+
+    print(final_classes_of_all_datasets, rename_map)

data/build/scenario.py

@@ -0,0 +1,466 @@
+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
+        

data/build_cl/build.py

@@ -0,0 +1,161 @@
+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 ..datasets.registery import static_dataset_registery
+from ..build.scenario import Scenario as DAScenario
+from copy import deepcopy
+from utils.common.log import logger
+import random
+from .scenario import _ABDatasetMetaInfo, Scenario
+        
+        
+def _check(source_datasets_meta_info: List[_ABDatasetMetaInfo], target_datasets_meta_info: List[_ABDatasetMetaInfo]):
+    # requirements for simplity
+    # 1. no same class in source datasets
+    
+    source_datasets_class = [i.classes for i in source_datasets_meta_info]
+    for ci1, c1 in enumerate(source_datasets_class):
+        for ci2, c2 in enumerate(source_datasets_class):
+            if ci1 == ci2:
+                continue
+            
+            c1_name = source_datasets_meta_info[ci1].name
+            c2_name = source_datasets_meta_info[ci2].name
+            intersection = set(c1).intersection(set(c2))
+            assert len(intersection) == 0, f'{c1_name} has intersection with {c2_name}: {intersection}'
+
+    
+def build_cl_scenario(
+    da_scenario: DAScenario,
+    target_datasets_name: List[str],
+    num_classes_per_task: int,
+    max_num_tasks: int,
+    data_dirs,
+    sanity_check=False
+):
+    config = deepcopy(locals())
+    
+    source_datasets_idx_map = {}
+    source_class_idx_max = 0
+    
+    for sd in da_scenario.config['source_datasets_name']:
+        da_scenario_idx_map = None
+        for k, v in da_scenario.all_datasets_idx_map.items():
+            if k.startswith(sd):
+                da_scenario_idx_map = v
+                break
+            
+        source_datasets_idx_map[sd] = da_scenario_idx_map
+        source_class_idx_max = max(source_class_idx_max, max(list(da_scenario_idx_map.values())))
+    
+    
+    target_class_idx_start = source_class_idx_max + 1
+    
+    target_datasets_meta_info = [_ABDatasetMetaInfo(d, *static_dataset_registery[d][1:], None, None) for d in target_datasets_name]
+    
+    task_datasets_seq = []
+    
+    num_tasks_per_dataset = {}
+    
+    for td_info_i, td_info in enumerate(target_datasets_meta_info):
+        
+        if td_info_i >= 1:
+            for _td_info_i, _td_info in enumerate(target_datasets_meta_info[0: td_info_i]):
+                if _td_info.name == td_info.name:
+                    # print(111)
+                    # class_idx_offset = sum([len(t.classes) for t in target_datasets_meta_info[0: td_info_i]])
+                    print(len(task_datasets_seq))
+                    
+                    task_index_offset = sum([v if __i < _td_info_i else 0 for __i, v in enumerate(num_tasks_per_dataset.values())])
+                    
+                    task_datasets_seq += task_datasets_seq[task_index_offset: task_index_offset + num_tasks_per_dataset[_td_info_i]]
+                    print(len(task_datasets_seq))
+                    break
+            continue
+        
+        td_classes = td_info.classes
+        num_tasks_per_dataset[td_info_i] = 0
+        
+        for ci in range(0, len(td_classes), num_classes_per_task):
+            task_i = ci // num_classes_per_task
+            task_datasets_seq += [_ABDatasetMetaInfo(
+                f'{td_info.name}|task-{task_i}|ci-{ci}-{ci + num_classes_per_task - 1}',
+                td_classes[ci: ci + num_classes_per_task],
+                td_info.task_type,
+                td_info.object_type,
+                td_info.class_aliases,
+                td_info.shift_type,
+                
+                td_classes[:ci] + td_classes[ci + num_classes_per_task: ],
+                {cii: cii + target_class_idx_start for cii in range(ci, ci + num_classes_per_task)}
+            )]
+            num_tasks_per_dataset[td_info_i] += 1
+            
+        if ci + num_classes_per_task < len(td_classes) - 1:
+            task_datasets_seq += [_ABDatasetMetaInfo(
+                f'{td_info.name}-task-{task_i + 1}|ci-{ci}-{ci + num_classes_per_task - 1}',
+                td_classes[ci: len(td_classes)],
+                td_info.task_type,
+                td_info.object_type,
+                td_info.class_aliases,
+                td_info.shift_type,
+                
+                td_classes[:ci],
+                {cii: cii + target_class_idx_start for cii in range(ci, len(td_classes))}
+            )]
+            num_tasks_per_dataset[td_info_i] += 1
+                
+        target_class_idx_start += len(td_classes)
+    
+    if len(task_datasets_seq) < max_num_tasks:
+        print(len(task_datasets_seq), max_num_tasks)
+        raise RuntimeError()
+    
+    task_datasets_seq = task_datasets_seq[0: max_num_tasks]
+    target_class_idx_start = max([max(list(td.idx_map.values())) + 1 for td in task_datasets_seq])
+    
+    scenario = Scenario(config, task_datasets_seq, target_class_idx_start, source_class_idx_max + 1, data_dirs)
+    
+    if sanity_check:
+        selected_tasks_index = []
+        for task_index, _ in enumerate(scenario.target_tasks_order):
+            cur_datasets = scenario.get_cur_task_train_datasets()
+            
+            if len(cur_datasets) < 300:
+                # empty_tasks_index += [task_index]
+                # while True:
+                    # replaced_task_index = random.randint(0, task_index - 1) # ensure no random
+                replaced_task_index = task_index // 2
+                assert replaced_task_index != task_index
+                while replaced_task_index in selected_tasks_index:
+                    replaced_task_index += 1
+                    
+                task_datasets_seq[task_index] = deepcopy(task_datasets_seq[replaced_task_index])
+                selected_tasks_index += [replaced_task_index]
+                
+                logger.warning(f'replace {task_index}-th task with {replaced_task_index}-th task')
+            
+            # print(task_index, [t.name for t in task_datasets_seq])
+            
+            scenario.next_task()
+            
+        # print([t.name for t in task_datasets_seq])
+            
+        if len(selected_tasks_index) > 0:
+            target_class_idx_start = max([max(list(td.idx_map.values())) + 1 for td in task_datasets_seq])
+            scenario = Scenario(config, task_datasets_seq, target_class_idx_start, source_class_idx_max + 1, data_dirs)
+            
+            for task_index, _ in enumerate(scenario.target_tasks_order):
+                cur_datasets = scenario.get_cur_task_train_datasets()
+                logger.info(f'task {task_index}, len {len(cur_datasets)}')
+                assert len(cur_datasets) > 0
+                
+                scenario.next_task()
+                
+            scenario = Scenario(config, task_datasets_seq, target_class_idx_start, source_class_idx_max + 1, data_dirs)
+
+    return scenario

data/build_cl/scenario.py

@@ -0,0 +1,146 @@
+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, MergedDataset, Scenario as DAScenario
+
+
+class _ABDatasetMetaInfo:
+    def __init__(self, name, classes, task_type, object_type, class_aliases, shift_type, ignore_classes, idx_map):
+        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
+        
+        self.ignore_classes = ignore_classes
+        self.idx_map = idx_map
+        
+    def __repr__(self) -> str:
+        return f'({self.name}, {self.classes}, {self.idx_map})'
+
+
+class Scenario:
+    def __init__(self, config, target_datasets_info: List[_ABDatasetMetaInfo], num_classes: int, num_source_classes: int, data_dirs):
+        self.config = config 
+        self.target_datasets_info = target_datasets_info
+        self.num_classes = num_classes
+        self.cur_task_index = 0
+        self.num_source_classes = num_source_classes
+        self.cur_class_offset = num_source_classes
+        self.data_dirs = data_dirs
+        
+        self.target_tasks_order = [i.name for i in self.target_datasets_info]
+        self.num_tasks_to_be_learn = sum([len(i.classes) for i in target_datasets_info])
+
+        logger.info(f'[scenario build] # classes: {num_classes}, # tasks to be learnt: {len(target_datasets_info)}, '
+                    f'# classes per task: {config["num_classes_per_task"]}')
+
+    def to_json(self):
+        config = copy.deepcopy(self.config)
+        config['da_scenario'] = config['da_scenario'].to_json()
+        target_datasets_info = [str(i) for i in self.target_datasets_info]
+        return dict(
+            config=config, target_datasets_info=target_datasets_info,
+            num_classes=self.num_classes
+        )
+        
+    def __str__(self):
+        return f'Scenario({self.to_json()})'
+    
+    def get_cur_class_offset(self):
+        return self.cur_class_offset
+    
+    def get_cur_num_class(self):
+        return len(self.target_datasets_info[self.cur_task_index].classes)
+    
+    def get_nc_per_task(self):
+        return len(self.target_datasets_info[0].classes)
+    
+    def next_task(self):
+        self.cur_class_offset += len(self.target_datasets_info[self.cur_task_index].classes)
+        self.cur_task_index += 1
+        
+        print(f'now, cur task: {self.cur_task_index}, cur_class_offset: {self.cur_class_offset}')
+        
+    def get_cur_task_datasets(self):
+        dataset_info = self.target_datasets_info[self.cur_task_index]
+        dataset_name = dataset_info.name.split('|')[0]
+        # print()
+        
+        # source_datasets_info = []
+        
+        res ={ **{split: get_dataset(dataset_name=dataset_name, 
+                    root_dir=self.data_dirs[dataset_name], 
+                    split=split, 
+                    transform=None, 
+                    ignore_classes=dataset_info.ignore_classes, 
+                    idx_map=dataset_info.idx_map) for split in ['train']}, 
+              
+              **{split: MergedDataset([get_dataset(dataset_name=dataset_name, 
+                    root_dir=self.data_dirs[dataset_name], 
+                    split=split, 
+                    transform=None, 
+                    ignore_classes=di.ignore_classes, 
+                    idx_map=di.idx_map) for di in self.target_datasets_info[0: self.cur_task_index + 1]]) 
+                 for split in ['val', 'test']}
+        }
+        
+        # if len(res['train']) < 200 or len(res['val']) < 200 or len(res['test']) < 200:
+        #     return None
+            
+        
+        if len(res['train']) < 1000:
+            res['train'] = MergedDataset([res['train']] * 5)
+            logger.info('aug train dataset')
+        if len(res['val']) < 1000:
+            res['val'] = MergedDataset(res['val'].datasets * 5)
+            logger.info('aug val dataset')
+        if len(res['test']) < 1000:
+            res['test'] = MergedDataset(res['test'].datasets * 5)
+            logger.info('aug test dataset')
+        # da_scenario: DAScenario = self.config['da_scenario']
+        # offline_datasets = da_scenario.get_offline_datasets()
+        
+        for k, v in res.items():
+            logger.info(f'{k} dataset: {len(v)}')
+        
+        # new_val_datasets = [
+        #     *[d['val'] for d in offline_datasets.values()],
+        #     res['val']
+        # ]
+        # res['val'] = MergedDataset(new_val_datasets)
+        
+        # new_test_datasets = [
+        #     *[d['test'] for d in offline_datasets.values()],
+        #     res['test']
+        # ]
+        # res['test'] = MergedDataset(new_test_datasets)
+        
+        return res
+    
+    def get_cur_task_train_datasets(self):
+        dataset_info = self.target_datasets_info[self.cur_task_index]
+        dataset_name = dataset_info.name.split('|')[0]
+        # print()
+        
+        # source_datasets_info = []
+        
+        res = get_dataset(dataset_name=dataset_name, 
+                    root_dir=self.data_dirs[dataset_name], 
+                    split='train', 
+                    transform=None, 
+                    ignore_classes=dataset_info.ignore_classes, 
+                    idx_map=dataset_info.idx_map)
+        
+        return res
+    
+    def get_online_cur_task_samples_for_training(self, num_samples):
+        dataset = self.get_cur_task_datasets()
+        dataset = dataset['train']
+        return next(iter(build_dataloader(dataset, num_samples, 0, True, None)))[0]

data/build_gen/build.py

@@ -0,0 +1,495 @@
+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)
+    

data/build_gen/merge_alias.py

@@ -0,0 +1,106 @@
+from re import L
+from typing import Dict, List
+from collections import Counter
+
+    
+def grouping(bondlist):
+    # reference: https://blog.csdn.net/YnagShanwen/article/details/111344386
+    groups = [] 
+    break1 = False    
+    while bondlist:
+        pair1 = bondlist.pop(0)        
+        a = 11111
+        b = 10000
+        while b != a:
+            a = b
+            for atomid in pair1:
+                for i,pair2 in enumerate(bondlist):            
+                    if atomid in pair2:
+                        pair1 = pair1 + pair2
+                        bondlist.pop(i)
+                        if not bondlist:
+                            break1 = True
+                        break
+                if break1:
+                    break
+            b = len(pair1)
+        groups.append(pair1)
+    return groups
+
+
+def build_semantic_class_info(classes: List[str], aliases: List[List[str]]):
+    res = []
+    for c in classes:
+        # print(res)
+        if len(aliases) == 0:
+            res += [[c]]
+        else:
+            find_alias = False
+            for alias in aliases:
+                if c in alias:
+                    res += [alias]
+                    find_alias = True
+                    break
+            if not find_alias:
+                res += [[c]]
+    # print(classes, res)
+    return res
+    
+
+def merge_the_same_meaning_classes(classes_info_of_all_datasets):
+    # print(classes_info_of_all_datasets)
+
+    semantic_classes_of_all_datasets = []
+    all_aliases = []
+    for classes, aliases in classes_info_of_all_datasets.values():
+        all_aliases += aliases
+    for classes, aliases in classes_info_of_all_datasets.values():
+        semantic_classes_of_all_datasets += build_semantic_class_info(classes, all_aliases)
+        
+    # print(semantic_classes_of_all_datasets)
+    
+    grouped_classes_of_all_datasets = grouping(semantic_classes_of_all_datasets)#匹配过后的数据
+   
+    # print(grouped_classes_of_all_datasets)
+    
+    # final_grouped_classes_of_all_datasets = [Counter(c).most_common()[0][0] for c in grouped_classes_of_all_datasets]
+    # use most common class name; if the same common, use shortest class name!
+    final_grouped_classes_of_all_datasets = []
+    for c in grouped_classes_of_all_datasets:
+        counter = Counter(c).most_common()
+        max_times = counter[0][1]
+        candidate_class_names = []
+        for item, times in counter:
+            if times < max_times:
+                break
+            candidate_class_names += [item]
+        candidate_class_names.sort(key=lambda x: len(x))
+        
+        final_grouped_classes_of_all_datasets += [candidate_class_names[0]]
+    res = {}
+    res_map = {d: {} for d in classes_info_of_all_datasets.keys()}
+    
+    for dataset_name, (classes, _) in classes_info_of_all_datasets.items():
+        final_classes = []
+        for c in classes:
+            for grouped_names, final_name in zip(grouped_classes_of_all_datasets, final_grouped_classes_of_all_datasets):
+                if c in grouped_names:
+                    final_classes += [final_name]
+                    if final_name != c:
+                        res_map[dataset_name][c] = final_name
+                    break
+        res[dataset_name] = sorted(set(final_classes), key=final_classes.index)
+    return res, res_map
+
+
+if __name__ == '__main__':
+    cifar10_classes = ['airplane', 'automobile', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck']
+    cifar10_aliases = [['automobile', 'car']]
+    stl10_classes = ['airplane', 'bird', 'car', 'cat', 'deer', 'dog', 'horse', 'monkey', 'ship', 'truck']
+
+    final_classes_of_all_datasets, rename_map = merge_the_same_meaning_classes({
+        'CIFAR10': (cifar10_classes, cifar10_aliases),
+        'STL10': (stl10_classes, [])
+    })
+
+    print(final_classes_of_all_datasets, rename_map)

data/build_gen/scenario.py

@@ -0,0 +1,473 @@
+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 ..datasets.dataset_split import train_val_split
+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']
+
+        # ori_datasets_map = {d: get_dataset(d, data_dirs[d], None, None, None, None) for d in source_datasets_name}
+        # res_source_datasets_map = {k: {split: train_val_split(copy.deepcopy(v), split, rate=0.97) for split in ['train', 'val']} for k, v in ori_datasets_map.items()}
+        # for ds in res_source_datasets_map.values():
+        #     for k, v in ds.items():
+        #         v.underlying_dataset.dataset.setSplit(k)
+        res_source_datasets_map = {d: {split: get_dataset(d, data_dirs[d], 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 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_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
+        

data/convert_all_load_to_single_load.py

@@ -0,0 +1,56 @@
+"""
+convert load-all-images-into-memory-before-training dataset
+to load-when-training-dataset
+
+
+"""
+
+
+from torchvision.datasets import CIFAR10, STL10, MNIST, USPS, SVHN
+import os
+import tqdm
+
+
+def convert(datasets_of_split, new_dir):
+    img_idx = {}
+
+    for d in datasets_of_split:
+        for x, y in tqdm.tqdm(d, total=len(d), dynamic_ncols=True):
+            # print(type(x), type(y))
+            # break
+            # y = str(y)
+            if y not in img_idx:
+                img_idx[y] = -1
+            img_idx[y] += 1
+
+            p = os.path.join(new_dir, f'{y:06d}', f'{img_idx[y]:06d}' + '.png')
+            os.makedirs(os.path.dirname(p), exist_ok=True)
+
+            x.save(p)
+
+
+if __name__ == '__main__':
+    # convert(
+    #     [CIFAR10('/data/zql/datasets/CIFAR10', True, download=True), CIFAR10('/data/zql/datasets/CIFAR10', False, download=True)],
+    #     '/data/zql/datasets/CIFAR10-single'
+    # )
+
+    # convert(
+    #     [STL10('/data/zql/datasets/STL10', 'train', download=False), STL10('/data/zql/datasets/STL10', 'test', download=False)],
+    #     '/data/zql/datasets/STL10-single'
+    # )
+
+    # convert(
+    #     [MNIST('/data/zql/datasets/MNIST', True, download=True), MNIST('/data/zql/datasets/MNIST', False, download=True)],
+    #     '/data/zql/datasets/MNIST-single'
+    # )
+
+    convert(
+        [SVHN('/data/zql/datasets/SVHN', 'train', download=True), SVHN('/data/zql/datasets/SVHN', 'test', download=True)],
+        '/data/zql/datasets/SVHN-single'
+    )
+
+    # convert(
+    #     [USPS('/data/zql/datasets/USPS', True, download=False), USPS('/data/zql/datasets/USPS', False, download=False)],
+    #     '/data/zql/datasets/USPS-single'
+    # )

data/convert_det_dataset_to_cls.py

@@ -0,0 +1,55 @@
+from data import ABDataset
+from utils.common.data_record import read_json, write_json
+from PIL import Image
+import os 
+from utils.common.file import ensure_dir
+import numpy as np
+from itertools import groupby
+from skimage import morphology, measure
+from PIL import Image
+from scipy import misc
+import tqdm
+from PIL import ImageFile
+ImageFile.LOAD_TRUNCATED_IMAGES = True
+import shutil
+
+
+def convert_det_dataset_to_det(coco_ann_json_path, data_dir, target_data_dir, min_img_size=224):
+    
+    coco_ann = read_json(coco_ann_json_path)
+    
+    img_id_to_path = {}
+    for img in coco_ann['images']:
+        img_id_to_path[img['id']] = os.path.join(data_dir, img['file_name'])
+    
+    classes_imgs_id_map = {}
+    for ann in tqdm.tqdm(coco_ann['annotations'], total=len(coco_ann['annotations']), dynamic_ncols=True):
+        img_id = ann['image_id']
+        img_path = img_id_to_path[img_id]
+        img = Image.open(img_path)
+
+        bbox = ann['bbox']
+        if bbox[2] < min_img_size or bbox[3] < min_img_size:
+            continue
+        
+        bbox = [bbox[0], bbox[1], bbox[0] + bbox[2], bbox[1] + bbox[3]]
+        
+        class_idx = str(ann['category_id'])
+        if class_idx not in classes_imgs_id_map.keys():
+            classes_imgs_id_map[class_idx] = 0
+        target_cropped_img_path = os.path.join(target_data_dir, class_idx, 
+                                               f'{classes_imgs_id_map[class_idx]}.{img_path.split(".")[-1]}')
+        classes_imgs_id_map[class_idx] += 1
+        
+        ensure_dir(target_cropped_img_path)
+        img.crop(bbox).save(target_cropped_img_path)
+        
+        
+        
+if __name__ == '__main__':
+    convert_det_dataset_to_det(
+        coco_ann_json_path='/data/zql/datasets/coco2017/train2017/coco_ann.json',
+        data_dir='/data/zql/datasets/coco2017/train2017',
+        target_data_dir='/data/zql/datasets/coco2017_for_cls_task',
+        min_img_size=224
+    )

data/convert_seg_dataset_to_cls.py

@@ -0,0 +1,324 @@
+from data import ABDataset
+from utils.common.data_record import read_json
+from PIL import Image
+import os 
+from utils.common.file import ensure_dir
+import numpy as np
+from itertools import groupby
+from skimage import morphology, measure
+from PIL import Image
+from scipy import misc
+import tqdm
+from PIL import ImageFile
+ImageFile.LOAD_TRUNCATED_IMAGES = True
+import shutil
+
+
+def convert_seg_dataset_to_cls(seg_imgs_path, seg_labels_path, target_cls_data_dir, ignore_classes_idx, thread_i, min_img_size=224, label_after_hook=lambda x: x):
+    """
+    Reference: https://blog.csdn.net/lizaijinsheng/article/details/119889946
+
+    NOTE: 
+    Background class should not be considered. 
+    However, if a seg dataset has only one valid class, so that the generated cls dataset also has only one class and 
+    the cls accuracy will be 100% forever. But we do not use the generated cls dataset alone, so it is ok.
+    """
+    assert len(seg_imgs_path) == len(seg_labels_path)
+    
+    classes_imgs_id_map = {}
+    
+    for seg_img_path, seg_label_path in tqdm.tqdm(zip(seg_imgs_path, seg_labels_path), total=len(seg_imgs_path), 
+                                                   dynamic_ncols=True, leave=False, desc=f'thread {thread_i}'):
+
+        try:
+            seg_img = Image.open(seg_img_path)
+            seg_label = Image.open(seg_label_path).convert('L')
+            seg_label = np.array(seg_label)
+            seg_label = label_after_hook(seg_label)
+        except Exception as e:
+            print(e)
+            print(f'file {seg_img_path} error, skip')
+            exit()
+        # seg_img = Image.open(seg_img_path)
+        # seg_label = Image.open(seg_label_path).convert('L')
+        # seg_label = np.array(seg_label)
+            
+        this_img_classes = set(seg_label.reshape(-1).tolist())
+        # print(this_img_classes)
+        
+        for class_idx in this_img_classes:
+            if class_idx in ignore_classes_idx:
+                continue
+            
+            if class_idx not in classes_imgs_id_map.keys():
+                classes_imgs_id_map[class_idx] = 0
+
+            mask = np.zeros((seg_label.shape[0], seg_label.shape[1]), dtype=np.uint8)
+            mask[seg_label == class_idx] = 1
+            mask_without_small = morphology.remove_small_objects(mask, min_size=10, connectivity=2)
+            label_image = measure.label(mask_without_small)
+
+            for region in measure.regionprops(label_image):
+                bbox = region.bbox # (top, left, bottom, right)
+                bbox = [bbox[1], bbox[0], bbox[3], bbox[2]]  # (left, top, right, bottom)
+                
+                width, height = bbox[2] - bbox[0], bbox[3] - bbox[1]
+                if width < min_img_size or height < min_img_size:
+                    continue
+                
+                target_cropped_img_path = os.path.join(target_cls_data_dir, str(class_idx), 
+                                                       f'{classes_imgs_id_map[class_idx]}.{seg_img_path.split(".")[-1]}')
+                ensure_dir(target_cropped_img_path)
+                seg_img.crop(bbox).save(target_cropped_img_path)    
+                # print(target_cropped_img_path)
+                # exit()
+                
+                classes_imgs_id_map[class_idx] += 1
+
+    num_cls_imgs = 0
+    for k, v in classes_imgs_id_map.items():
+        # print(f'# class {k}: {v + 1}')
+        num_cls_imgs += v
+    # print(f'total: {num_cls_imgs}')
+    
+    return classes_imgs_id_map
+    
+
+from concurrent.futures import ThreadPoolExecutor
+
+
+
+# def convert_seg_dataset_to_cls_multi_thread(seg_imgs_path, seg_labels_path, target_cls_data_dir, ignore_classes_idx, num_threads):
+#     if os.path.exists(target_cls_data_dir):
+#         shutil.rmtree(target_cls_data_dir)
+    
+#     assert len(seg_imgs_path) == len(seg_labels_path)
+#     n = len(seg_imgs_path) // num_threads
+    
+#     pool = ThreadPoolExecutor(max_workers=num_threads)
+#     # threads = []
+#     futures = []
+#     for thread_i in range(num_threads):
+#         # thread = threading.Thread(target=convert_seg_dataset_to_cls, 
+#         #                           args=(seg_imgs_path[thread_i * n: (thread_i + 1) * n], 
+#         #                                 seg_labels_path[thread_i * n: (thread_i + 1) * n], 
+#         #                                 target_cls_data_dir, ignore_classes_idx))
+#         # threads += [thread]
+#         future = pool.submit(convert_seg_dataset_to_cls, *(seg_imgs_path[thread_i * n: (thread_i + 1) * n], 
+#                                         seg_labels_path[thread_i * n: (thread_i + 1) * n], 
+#                                         target_cls_data_dir, ignore_classes_idx, thread_i))
+#         futures += [future]
+    
+#     futures += [
+#         pool.submit(convert_seg_dataset_to_cls, *(seg_imgs_path[(thread_i + 1) * n: ], 
+#                                         seg_labels_path[(thread_i + 1) * n: ], 
+#                                         target_cls_data_dir, ignore_classes_idx, thread_i))
+#     ]
+    
+#     for f in futures:
+#         f.done()
+    
+#     res = []
+#     for f in futures:
+#         res += [f.result()]
+#         print(res[-1])
+    
+#     res_dist = {}
+#     for r in res:
+#         for k, v in r.items():
+#             if k in res_dist.keys():
+#                 res_dist[k] += v 
+#             else:
+#                 res_dist[k] = v
+    
+#     print('results:')
+#     print(res_dist)
+    
+#     pool.shutdown()
+
+
+
+import random
+def random_crop_aug(target_dir):
+    for class_dir in os.listdir(target_dir):
+        class_dir = os.path.join(target_dir, class_dir)
+        
+        for img_path in os.listdir(class_dir):
+            img_path = os.path.join(class_dir, img_path)
+
+            img = Image.open(img_path)
+            
+            w, h = img.width, img.height
+            
+            for ri in range(5):
+                img.crop(
+                    [
+                        random.randint(0, w // 5),
+                        random.randint(0, h // 5),
+                        random.randint(w // 5 * 4, w),
+                        random.randint(h // 5 * 4, h)
+                    ]
+                ).save(
+                    os.path.join(os.path.dirname(img_path), f'randaug_{ri}_' + os.path.basename(img_path))
+                )
+                # print(img_path)
+                # exit()
+            
+
+if __name__ == '__main__':
+    # SuperviselyPerson
+    # root_dir = '/data/zql/datasets/supervisely_person/Supervisely Person Dataset'
+    
+    # images_path, labels_path = [], []
+    # for p in os.listdir(root_dir):
+    #     if p.startswith('ds'):
+    #         p1 = os.path.join(root_dir, p, 'img')
+    #         images_path += [(p, os.path.join(p1, n)) for n in os.listdir(p1)]
+    # for dsi, img_p in images_path:
+    #     target_p = os.path.join(root_dir, p, dsi, img_p.split('/')[-1])
+    #     labels_path += [target_p]
+    # images_path = [i[1] for i in images_path]
+    
+    # target_dir = '/data/zql/datasets/supervisely_person_for_cls_task'
+    # if os.path.exists(target_dir):
+    #     shutil.rmtree(target_dir)
+    # convert_seg_dataset_to_cls(
+    #     seg_imgs_path=images_path,
+    #     seg_labels_path=labels_path,
+    #     target_cls_data_dir=target_dir,
+    #     ignore_classes_idx=[0, 2],
+    #     # num_threads=8
+    #     thread_i=0
+    # )
+    
+    # random_crop_aug('/data/zql/datasets/supervisely_person_for_cls_task')
+    
+    
+    # GTA5
+    # root_dir = '/data/zql/datasets/GTA-ls-copy/GTA5'
+    # images_path, labels_path = [], []
+    # for p in os.listdir(os.path.join(root_dir, 'images')):
+    #     p = os.path.join(root_dir, 'images', p)
+    #     if not p.endswith('png'):
+    #         continue
+    #     images_path += [p]
+    #     labels_path += [p.replace('images', 'labels_gt')]
+
+    # target_dir = '/data/zql/datasets/gta5_for_cls_task'
+    # if os.path.exists(target_dir):
+    #     shutil.rmtree(target_dir)
+    
+    # convert_seg_dataset_to_cls(
+    #     seg_imgs_path=images_path,
+    #     seg_labels_path=labels_path,
+    #     target_cls_data_dir=target_dir,
+    #     ignore_classes_idx=[],
+    #     thread_i=0
+    # )
+    
+    # cityscapes
+    # root_dir = '/data/zql/datasets/cityscape/'
+    
+    # def _get_target_suffix(mode: str, target_type: str) -> str:
+    #     if target_type == 'instance':
+    #         return '{}_instanceIds.png'.format(mode)
+    #     elif target_type == 'semantic':
+    #         return '{}_labelIds.png'.format(mode)
+    #     elif target_type == 'color':
+    #         return '{}_color.png'.format(mode)
+    #     else:
+    #         return '{}_polygons.json'.format(mode)
+
+    
+    # images_path, labels_path = [], []
+    # split = 'train'
+    # images_dir = os.path.join(root_dir, 'leftImg8bit', split)
+    # targets_dir = os.path.join(root_dir, 'gtFine', split)
+    # for city in os.listdir(images_dir):
+    #     img_dir = os.path.join(images_dir, city)
+    #     target_dir = os.path.join(targets_dir, city)
+    #     for file_name in os.listdir(img_dir):
+    #         target_types = []
+    #         for t in ['semantic']:
+    #             target_name = '{}_{}'.format(file_name.split('_leftImg8bit')[0],
+    #                                             _get_target_suffix('gtFine', t))
+    #             target_types.append(os.path.join(target_dir, target_name))
+
+    #         images_path.append(os.path.join(img_dir, file_name))
+    #         labels_path.append(target_types[0])
+            
+    # print(images_path[0: 5], '\n', labels_path[0: 5])
+    
+    # target_dir = '/data/zql/datasets/cityscapes_for_cls_task'
+    # if os.path.exists(target_dir):
+    #     shutil.rmtree(target_dir)
+    # convert_seg_dataset_to_cls(
+    #     seg_imgs_path=images_path,
+    #     seg_labels_path=labels_path,
+    #     target_cls_data_dir=target_dir,
+    #     ignore_classes_idx=[],
+    #     # num_threads=8
+    #     thread_i=0
+    # )
+    
+    # import shutil
+    
+    # ignore_target_dir = '/data/zql/datasets/cityscapes_for_cls_task_ignored'
+    
+    # ignore_label = 255
+    # raw_idx_map_in_y_transform = {-1: ignore_label, 0: ignore_label, 1: ignore_label, 2: ignore_label,
+    #         3: ignore_label, 4: ignore_label, 5: ignore_label, 6: ignore_label,
+    #         7: 0, 8: 1, 9: ignore_label, 10: ignore_label, 11: 2, 12: 3, 13: 4,
+    #         14: ignore_label, 15: ignore_label, 16: ignore_label, 17: 5,
+    #         18: ignore_label, 19: 6, 20: 7, 21: 8, 22: 9, 23: 10, 24: 11, 25: 12, 26: 13, 27: 14,
+    #         28: 15, 29: ignore_label, 30: ignore_label, 31: 16, 32: 17, 33: 18}
+    # ignore_classes_idx = [k for k, v in raw_idx_map_in_y_transform.items() if v == ignore_label]
+    # ignore_classes_idx = sorted(ignore_classes_idx)
+    
+    # for class_dir in os.listdir(target_dir):
+    #     if int(class_dir) in ignore_classes_idx:
+    #         continue
+    #         shutil.move(
+    #             os.path.join(target_dir, class_dir),
+    #             os.path.join(ignore_target_dir, class_dir)
+    #         )
+    #     else:
+    #         shutil.move(
+    #             os.path.join(target_dir, class_dir),
+    #             os.path.join(target_dir, str(raw_idx_map_in_y_transform[int(class_dir)]))
+    #         )
+    #         continue
+    #     print(class_dir)
+    # exit()
+    
+    
+    
+    # baidu person
+    # root_dir = '/data/zql/datasets/baidu_person/clean_images/'
+    
+    # images_path, labels_path = [], []
+    # for p in os.listdir(os.path.join(root_dir, 'images')):
+    #     images_path += [os.path.join(root_dir, 'images', p)]
+    #     labels_path += [os.path.join(root_dir, 'profiles', p.split('.')[0] + '-profile.jpg')]
+    
+    # target_dir = '/data/zql/datasets/baiduperson_for_cls_task'
+    # # if os.path.exists(target_dir):
+    # #     shutil.rmtree(target_dir)
+        
+    # def label_after_hook(x):
+    #     x[x > 1] = 1
+    #     return x    
+    
+    # convert_seg_dataset_to_cls(
+    #     seg_imgs_path=images_path,
+    #     seg_labels_path=labels_path,
+    #     target_cls_data_dir=target_dir,
+    #     ignore_classes_idx=[1],
+    #     # num_threads=8
+    #     thread_i=1,
+    #     min_img_size=224,
+    #     label_after_hook=label_after_hook
+    # )
+    
+    
+    

data/convert_seg_dataset_to_det.py

@@ -0,0 +1,399 @@
+from data import ABDataset
+from utils.common.data_record import read_json, write_json
+from PIL import Image
+import os 
+from utils.common.file import ensure_dir
+import numpy as np
+from itertools import groupby
+from skimage import morphology, measure
+from PIL import Image
+from scipy import misc
+import tqdm
+from PIL import ImageFile
+ImageFile.LOAD_TRUNCATED_IMAGES = True
+import shutil
+
+
+def convert_seg_dataset_to_det(seg_imgs_path, seg_labels_path, root_dir, target_coco_ann_path, ignore_classes_idx, thread_i, min_img_size=224, label_after_hook=lambda x: x):
+    """
+    Reference: https://blog.csdn.net/lizaijinsheng/article/details/119889946
+
+    NOTE: 
+    Background class should not be considered. 
+    However, if a seg dataset has only one valid class, so that the generated cls dataset also has only one class and 
+    the cls accuracy will be 100% forever. But we do not use the generated cls dataset alone, so it is ok.
+    """
+    assert len(seg_imgs_path) == len(seg_labels_path)
+    
+    classes_imgs_id_map = {}
+    
+    coco_ann = {
+        'categories': [],
+        "type": "instances",
+        'images': [],
+        'annotations': []
+    }
+    
+    image_id = 0
+    ann_id = 0
+    
+    pbar = tqdm.tqdm(zip(seg_imgs_path, seg_labels_path), total=len(seg_imgs_path), 
+                                                   dynamic_ncols=True, leave=False, desc=f'thread {thread_i}')
+    for seg_img_path, seg_label_path in pbar:
+
+        try:
+            seg_img = Image.open(seg_img_path)
+            seg_label = Image.open(seg_label_path).convert('L')
+            seg_label = np.array(seg_label)
+            seg_label = label_after_hook(seg_label)
+        except Exception as e:
+            print(e)
+            print(f'file {seg_img_path} error, skip')
+            exit()
+        # seg_img = Image.open(seg_img_path)
+        # seg_label = Image.open(seg_label_path).convert('L')
+        # seg_label = np.array(seg_label)
+        
+        image_coco_info = {'file_name': os.path.relpath(seg_img_path, root_dir), 'height': seg_img.height, 'width': seg_img.width,
+                 'id':image_id}
+        image_id += 1
+        coco_ann['images'] += [image_coco_info]
+            
+        this_img_classes = set(seg_label.reshape(-1).tolist())
+        # print(this_img_classes)
+        
+        for class_idx in this_img_classes:
+            if class_idx in ignore_classes_idx:
+                continue
+            
+            if class_idx not in classes_imgs_id_map.keys():
+                classes_imgs_id_map[class_idx] = 0
+
+            mask = np.zeros((seg_label.shape[0], seg_label.shape[1]), dtype=np.uint8)
+            mask[seg_label == class_idx] = 1
+            mask_without_small = morphology.remove_small_objects(mask, min_size=10, connectivity=2)
+            label_image = measure.label(mask_without_small)
+
+            for region in measure.regionprops(label_image):
+                bbox = region.bbox # (top, left, bottom, right)
+                bbox = [bbox[1], bbox[0], bbox[3], bbox[2]]  # (left, top, right, bottom)
+                
+                width, height = bbox[2] - bbox[0], bbox[3] - bbox[1]
+                if width < min_img_size or height < min_img_size:
+                    continue
+                
+                # target_cropped_img_path = os.path.join(target_cls_data_dir, str(class_idx), 
+                #                                        f'{classes_imgs_id_map[class_idx]}.{seg_img_path.split(".")[-1]}')
+                # ensure_dir(target_cropped_img_path)
+                # seg_img.crop(bbox).save(target_cropped_img_path)    
+                # print(target_cropped_img_path)
+                # exit()
+                
+                ann_coco_info = {'area': width*height, 'iscrowd': 0, 'image_id':
+                   image_id - 1, 'bbox': [bbox[0], bbox[1], width, height],
+                   'category_id': class_idx, 
+                   'id': ann_id, 'ignore': 0,
+                   'segmentation': []}
+                ann_id += 1
+                
+                coco_ann['annotations'] += [ann_coco_info]
+                
+                classes_imgs_id_map[class_idx] += 1
+                
+                pbar.set_description(f'# ann: {ann_id}')
+
+    coco_ann['categories'] = [
+        {'id': ci, 'name': f'class_{c}_in_seg'} for ci, c in enumerate(classes_imgs_id_map.keys())
+    ]
+    c_to_ci_map = {c: ci for ci, c in enumerate(classes_imgs_id_map.keys())}
+    for ann in coco_ann['annotations']:
+        ann['category_id'] = c_to_ci_map[
+            ann['category_id']
+        ]
+    
+    write_json(target_coco_ann_path, coco_ann, indent=0, backup=True)
+    write_json(os.path.join(root_dir, 'coco_ann.json'), coco_ann, indent=0, backup=True)
+
+    num_cls_imgs = 0
+    for k, v in classes_imgs_id_map.items():
+        # print(f'# class {k}: {v + 1}')
+        num_cls_imgs += v
+    # print(f'total: {num_cls_imgs}')
+    
+    return classes_imgs_id_map
+    
+
+from concurrent.futures import ThreadPoolExecutor
+
+
+
+# def convert_seg_dataset_to_cls_multi_thread(seg_imgs_path, seg_labels_path, target_cls_data_dir, ignore_classes_idx, num_threads):
+#     if os.path.exists(target_cls_data_dir):
+#         shutil.rmtree(target_cls_data_dir)
+    
+#     assert len(seg_imgs_path) == len(seg_labels_path)
+#     n = len(seg_imgs_path) // num_threads
+    
+#     pool = ThreadPoolExecutor(max_workers=num_threads)
+#     # threads = []
+#     futures = []
+#     for thread_i in range(num_threads):
+#         # thread = threading.Thread(target=convert_seg_dataset_to_cls, 
+#         #                           args=(seg_imgs_path[thread_i * n: (thread_i + 1) * n], 
+#         #                                 seg_labels_path[thread_i * n: (thread_i + 1) * n], 
+#         #                                 target_cls_data_dir, ignore_classes_idx))
+#         # threads += [thread]
+#         future = pool.submit(convert_seg_dataset_to_cls, *(seg_imgs_path[thread_i * n: (thread_i + 1) * n], 
+#                                         seg_labels_path[thread_i * n: (thread_i + 1) * n], 
+#                                         target_cls_data_dir, ignore_classes_idx, thread_i))
+#         futures += [future]
+    
+#     futures += [
+#         pool.submit(convert_seg_dataset_to_cls, *(seg_imgs_path[(thread_i + 1) * n: ], 
+#                                         seg_labels_path[(thread_i + 1) * n: ], 
+#                                         target_cls_data_dir, ignore_classes_idx, thread_i))
+#     ]
+    
+#     for f in futures:
+#         f.done()
+    
+#     res = []
+#     for f in futures:
+#         res += [f.result()]
+#         print(res[-1])
+    
+#     res_dist = {}
+#     for r in res:
+#         for k, v in r.items():
+#             if k in res_dist.keys():
+#                 res_dist[k] += v 
+#             else:
+#                 res_dist[k] = v
+    
+#     print('results:')
+#     print(res_dist)
+    
+#     pool.shutdown()
+
+
+
+# import random
+# def random_crop_aug(target_dir):
+#     for class_dir in os.listdir(target_dir):
+#         class_dir = os.path.join(target_dir, class_dir)
+        
+#         for img_path in os.listdir(class_dir):
+#             img_path = os.path.join(class_dir, img_path)
+
+#             img = Image.open(img_path)
+            
+#             w, h = img.width, img.height
+            
+#             for ri in range(5):
+#                 img.crop(
+#                     [
+#                         random.randint(0, w // 5),
+#                         random.randint(0, h // 5),
+#                         random.randint(w // 5 * 4, w),
+#                         random.randint(h // 5 * 4, h)
+#                     ]
+#                 ).save(
+#                     os.path.join(os.path.dirname(img_path), f'randaug_{ri}_' + os.path.basename(img_path))
+#                 )
+#                 # print(img_path)
+#                 # exit()
+
+
+def post_ignore_classes(coco_ann_json_path):
+    # from data.datasets.object_detection.yolox_data_util.api import remap_dataset
+    # remap_dataset(coco_ann_json_path, [], {})
+    pass
+    
+    
+
+if __name__ == '__main__':
+    # SuperviselyPerson
+    # root_dir = '/data/zql/datasets/supervisely_person_full_20230635/Supervisely Person Dataset'
+    
+    # images_path, labels_path = [], []
+    # for p in os.listdir(root_dir):
+    #     if p.startswith('ds'):
+    #         p1 = os.path.join(root_dir, p, 'img')
+    #         images_path += [(p, os.path.join(p1, n)) for n in os.listdir(p1)]
+    # for dsi, img_p in images_path:
+    #     target_p = os.path.join(root_dir, p, dsi, img_p.split('/')[-1])
+    #     labels_path += [target_p]
+    # images_path = [i[1] for i in images_path]
+    
+    # target_coco_ann_path = '/data/zql/datasets/supervisely_person_for_det_task/coco_ann.json'
+    # if os.path.exists(target_coco_ann_path):
+    #     os.remove(target_coco_ann_path)
+    # convert_seg_dataset_to_det(
+    #     seg_imgs_path=images_path,
+    #     seg_labels_path=labels_path,
+    #     root_dir=root_dir,
+    #     target_coco_ann_path=target_coco_ann_path,
+    #     ignore_classes_idx=[0, 2],
+    #     # num_threads=8
+    #     thread_i=0
+    # )
+    
+    # random_crop_aug('/data/zql/datasets/supervisely_person_for_cls_task')
+    
+    
+    # GTA5
+    # root_dir = '/data/zql/datasets/GTA-ls-copy/GTA5'
+    # images_path, labels_path = [], []
+    # for p in os.listdir(os.path.join(root_dir, 'images')):
+    #     p = os.path.join(root_dir, 'images', p)
+    #     if not p.endswith('png'):
+    #         continue
+    #     images_path += [p]
+    #     labels_path += [p.replace('images', 'labels_gt')]
+
+    # target_coco_ann_path = '/data/zql/datasets/gta5_for_det_task/coco_ann.json'
+    # if os.path.exists(target_coco_ann_path):
+    #     os.remove(target_coco_ann_path)
+    
+    # """
+    # [
+    #     'road', 'sidewalk', 'building', 'wall',
+    #     'fence', 'pole', 'light', 'sign',
+    #     'vegetation', 'terrain', 'sky', 'people', # person
+    #     'rider', 'car', 'truck', 'bus', 'train',
+    #     'motocycle', 'bicycle'
+    # ]
+    # """
+    # need_classes_idx = [13, 15]
+    # convert_seg_dataset_to_det(
+    #     seg_imgs_path=images_path,
+    #     seg_labels_path=labels_path,
+    #     root_dir=root_dir,
+    #     target_coco_ann_path=target_coco_ann_path,
+    #     ignore_classes_idx=[i for i in range(20) if i not in need_classes_idx],
+    #     thread_i=0
+    # )
+    
+    # from data.datasets.object_detection.yolox_data_util.api import remap_dataset
+    # new_coco_ann_json_path = remap_dataset('/data/zql/datasets/GTA-ls-copy/GTA5/coco_ann.json', [-1], {0: 0, 1:-1, 2:-1, 3: 1, 4:-1, 5:-1})
+    # print(new_coco_ann_json_path)
+    
+    # cityscapes
+    # root_dir = '/data/zql/datasets/cityscape/'
+    
+    # def _get_target_suffix(mode: str, target_type: str) -> str:
+    #     if target_type == 'instance':
+    #         return '{}_instanceIds.png'.format(mode)
+    #     elif target_type == 'semantic':
+    #         return '{}_labelIds.png'.format(mode)
+    #     elif target_type == 'color':
+    #         return '{}_color.png'.format(mode)
+    #     else:
+    #         return '{}_polygons.json'.format(mode)
+
+    
+    # images_path, labels_path = [], []
+    # split = 'train'
+    # images_dir = os.path.join(root_dir, 'leftImg8bit', split)
+    # targets_dir = os.path.join(root_dir, 'gtFine', split)
+    # for city in os.listdir(images_dir):
+    #     img_dir = os.path.join(images_dir, city)
+    #     target_dir = os.path.join(targets_dir, city)
+    #     for file_name in os.listdir(img_dir):
+    #         target_types = []
+    #         for t in ['semantic']:
+    #             target_name = '{}_{}'.format(file_name.split('_leftImg8bit')[0],
+    #                                             _get_target_suffix('gtFine', t))
+    #             target_types.append(os.path.join(target_dir, target_name))
+
+    #         images_path.append(os.path.join(img_dir, file_name))
+    #         labels_path.append(target_types[0])
+            
+    # # print(images_path[0: 5], '\n', labels_path[0: 5])
+    
+    # target_coco_ann_path = '/data/zql/datasets/cityscape/coco_ann.json'
+    # # if os.path.exists(target_dir):
+    # #     shutil.rmtree(target_dir)
+        
+    # need_classes_idx = [26, 28]
+    # convert_seg_dataset_to_det(
+    #     seg_imgs_path=images_path,
+    #     seg_labels_path=labels_path,
+    #     root_dir=root_dir,
+    #     target_coco_ann_path=target_coco_ann_path,
+    #     ignore_classes_idx=[i for i in range(80) if i not in need_classes_idx],
+    #     # num_threads=8
+    #     thread_i=0
+    # )
+    
+    # import shutil
+    
+    # ignore_target_dir = '/data/zql/datasets/cityscapes_for_cls_task_ignored'
+    
+    # ignore_label = 255
+    # raw_idx_map_in_y_transform = {-1: ignore_label, 0: ignore_label, 1: ignore_label, 2: ignore_label,
+    #         3: ignore_label, 4: ignore_label, 5: ignore_label, 6: ignore_label,
+    #         7: 0, 8: 1, 9: ignore_label, 10: ignore_label, 11: 2, 12: 3, 13: 4,
+    #         14: ignore_label, 15: ignore_label, 16: ignore_label, 17: 5,
+    #         18: ignore_label, 19: 6, 20: 7, 21: 8, 22: 9, 23: 10, 24: 11, 25: 12, 26: 13, 27: 14,
+    #         28: 15, 29: ignore_label, 30: ignore_label, 31: 16, 32: 17, 33: 18}
+    # ignore_classes_idx = [k for k, v in raw_idx_map_in_y_transform.items() if v == ignore_label]
+    # ignore_classes_idx = sorted(ignore_classes_idx)
+    
+    # for class_dir in os.listdir(target_dir):
+    #     if int(class_dir) in ignore_classes_idx:
+    #         continue
+    #         shutil.move(
+    #             os.path.join(target_dir, class_dir),
+    #             os.path.join(ignore_target_dir, class_dir)
+    #         )
+    #     else:
+    #         shutil.move(
+    #             os.path.join(target_dir, class_dir),
+    #             os.path.join(target_dir, str(raw_idx_map_in_y_transform[int(class_dir)]))
+    #         )
+    #         continue
+    #     print(class_dir)
+    # exit()
+    
+    
+    
+    # baidu person
+    # root_dir = '/data/zql/datasets/baidu_person/clean_images/'
+    
+    # images_path, labels_path = [], []
+    # for p in os.listdir(os.path.join(root_dir, 'images')):
+    #     images_path += [os.path.join(root_dir, 'images', p)]
+    #     labels_path += [os.path.join(root_dir, 'profiles', p.split('.')[0] + '-profile.jpg')]
+    
+    # target_dir = '/data/zql/datasets/baiduperson_for_cls_task'
+    # # if os.path.exists(target_dir):
+    # #     shutil.rmtree(target_dir)
+        
+    # def label_after_hook(x):
+    #     x[x > 1] = 1
+    #     return x    
+    
+    # convert_seg_dataset_to_det(
+    #     seg_imgs_path=images_path,
+    #     seg_labels_path=labels_path,
+    #     root_dir=root_dir,
+    #     target_coco_ann_path='/data/zql/datasets/baidu_person/clean_images/coco_ann_zql.json',
+    #     ignore_classes_idx=[1],
+    #     # num_threads=8
+    #     thread_i=1,
+    #     min_img_size=224,
+    #     label_after_hook=label_after_hook
+    # )
+    
+    
+    # from data.visualize import visualize_classes_in_object_detection
+    # from data import get_dataset
+    # d = get_dataset('CityscapesDet', '/data/zql/datasets/cityscape/', 'val', None, [], None)
+    # visualize_classes_in_object_detection(d, {'car': 0, 'bus': 1}, {}, 'debug.png')
+    
+    # d = get_dataset('GTA5Det', '/data/zql/datasets/GTA-ls-copy/GTA5', 'val', None, [], None)
+    # visualize_classes_in_object_detection(d, {'car': 0, 'bus': 1}, {}, 'debug.png')
+    
+    # d = get_dataset('BaiduPersonDet', '/data/zql/datasets/baidu_person/clean_images/', 'val', None, [], None)
+    # visualize_classes_in_object_detection(d, {'person': 0}, {}, 'debug.png')

data/dataloader.py

@@ -0,0 +1,131 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+# domainbed/lib/fast_data_loader.py
+
+import torch
+from .datasets.ab_dataset import ABDataset
+
+
+class _InfiniteSampler(torch.utils.data.Sampler):
+    """Wraps another Sampler to yield an infinite stream."""
+
+    def __init__(self, sampler):
+        self.sampler = sampler
+
+    def __iter__(self):
+        while True:
+            for batch in self.sampler:
+                yield batch
+
+
+class InfiniteDataLoader:
+    def __init__(self, dataset, weights, batch_size, num_workers, collate_fn=None):
+        super().__init__()
+
+        if weights:
+            sampler = torch.utils.data.WeightedRandomSampler(
+                weights, replacement=True, num_samples=batch_size
+            )
+        else:
+            sampler = torch.utils.data.RandomSampler(dataset, replacement=True)
+
+        batch_sampler = torch.utils.data.BatchSampler(
+            sampler, batch_size=batch_size, drop_last=True
+        )
+
+        if collate_fn is not None:
+            self._infinite_iterator = iter(
+                torch.utils.data.DataLoader(
+                    dataset,
+                    num_workers=num_workers,
+                    batch_sampler=_InfiniteSampler(batch_sampler),
+                    pin_memory=False,
+                    collate_fn=collate_fn
+                )
+            )
+        else:
+            self._infinite_iterator = iter(
+                torch.utils.data.DataLoader(
+                    dataset,
+                    num_workers=num_workers,
+                    batch_sampler=_InfiniteSampler(batch_sampler),
+                    pin_memory=False
+                )
+            )
+        self.dataset = dataset
+
+    def __iter__(self):
+        while True:
+            yield next(self._infinite_iterator)
+
+    def __len__(self):
+        raise ValueError
+
+
+class FastDataLoader:
+    """
+    DataLoader wrapper with slightly improved speed by not respawning worker
+    processes at every epoch.
+    """
+
+    def __init__(self, dataset, batch_size, num_workers, shuffle=False, collate_fn=None):
+        super().__init__()
+        
+        self.num_workers = num_workers
+
+        if shuffle:
+            sampler = torch.utils.data.RandomSampler(dataset, replacement=False)
+        else:
+            sampler = torch.utils.data.SequentialSampler(dataset)
+
+        batch_sampler = torch.utils.data.BatchSampler(
+            sampler,
+            batch_size=batch_size,
+            drop_last=False,
+        )
+        if collate_fn is not None:
+            self._infinite_iterator = iter(
+                torch.utils.data.DataLoader(
+                    dataset,
+                    num_workers=num_workers,
+                    batch_sampler=_InfiniteSampler(batch_sampler),
+                    pin_memory=False,
+                    collate_fn=collate_fn
+                )
+            )
+        else:
+            self._infinite_iterator = iter(
+                torch.utils.data.DataLoader(
+                    dataset,
+                    num_workers=num_workers,
+                    batch_sampler=_InfiniteSampler(batch_sampler),
+                    pin_memory=False,
+                )
+            )
+
+        self.dataset = dataset
+        self.batch_size = batch_size
+        self._length = len(batch_sampler)
+
+    def __iter__(self):
+        for _ in range(len(self)):
+            yield next(self._infinite_iterator)
+
+    def __len__(self):
+        return self._length
+
+
+def build_dataloader(dataset: ABDataset, batch_size: int, num_workers: int, infinite: bool, shuffle_when_finite: bool, collate_fn=None):
+    assert batch_size <= len(dataset), len(dataset)
+    if infinite:
+        dataloader = InfiniteDataLoader(
+            dataset, None, batch_size, num_workers=num_workers, collate_fn=collate_fn)
+    else:
+        dataloader = FastDataLoader(
+            dataset, batch_size, num_workers, shuffle=shuffle_when_finite, collate_fn=collate_fn)
+
+    return dataloader
+
+
+def get_a_batch_dataloader(dataset: ABDataset, batch_size: int, num_workers: int, infinite: bool, shuffle_when_finite: bool):
+    pass
+    

data/dataset.py

@@ -0,0 +1,43 @@
+import importlib
+from typing import Type
+import torch
+from torch.utils.data import TensorDataset
+from torch.utils.data.dataloader import DataLoader
+
+from .datasets.ab_dataset import ABDataset
+
+from .datasets import * # import all datasets
+from .datasets.registery import static_dataset_registery
+
+
+def get_dataset(dataset_name, root_dir, split, transform=None, ignore_classes=[], idx_map=None) -> ABDataset:
+    dataset_cls = static_dataset_registery[dataset_name][0]
+    dataset = dataset_cls(root_dir, split, transform, ignore_classes, idx_map)
+
+    return dataset
+
+
+def get_num_limited_dataset(dataset: ABDataset, num_samples: int, discard_label=True):
+    dataloader = iter(DataLoader(dataset, num_samples // 2, shuffle=True))
+    x, y = [], []
+    cur_num_samples = 0
+    while True:
+        batch = next(dataloader)
+        cur_x, cur_y = batch[0], batch[1]
+        
+        x += [cur_x]
+        y += [cur_y]
+        cur_num_samples += cur_x.size(0)
+        
+        if cur_num_samples >= num_samples:
+            break
+        
+    x, y = torch.cat(x)[0: num_samples], torch.cat(y)[0: num_samples]
+    if discard_label:
+        new_dataset = TensorDataset(x)
+    else:
+        new_dataset = TensorDataset(x, y)
+    
+    dataset.dataset = new_dataset
+    
+    return dataset

data/datasets/__init__.py

@@ -0,0 +1,12 @@
+from .image_classification import *
+from .object_detection import *
+from .semantic_segmentation import *
+from .action_recognition import *
+
+from .sentiment_classification import *
+from .text_generation import *
+from .machine_translation import *
+from .pos_tagging import *
+
+from .mm_image_classification import *
+from .visual_question_answering import *