Upload 1804 files

.gitattributes

@@ -33,3 +33,9 @@ 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

.gitignore

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

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,12 @@
+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 .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/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,11 @@
+from .image_classification import *
+from .object_detection import *
+from .semantic_segmentation import *
+from .action_recognition import *
+
+from .sentiment_classification import *
+from .machine_translation import *
+from .pos_tagging import *
+
+from .mm_image_classification import *
+from .visual_question_answering import *

data/datasets/ab_dataset.py

@@ -0,0 +1,48 @@
+from abc import ABC, abstractmethod
+from typing import Dict, List, Optional
+from torchvision.transforms import Compose
+
+
+class ABDataset(ABC):
+    def __init__(self, root_dir, split, transform=None, ignore_classes=[], idx_map=None):
+        
+        self.root_dir = root_dir
+        self.split = split
+        self.transform = transform
+        self.ignore_classes = ignore_classes
+        self.idx_map = idx_map
+        
+        self.dataset = None
+        
+        # injected by @dataset_register
+        self.name = None
+        self.classes = None
+        self.raw_classes = None
+        self.class_aliases = None
+        self.shift_type = None
+        self.task_type = None # ['Image Classification', 'Object Detection', ...]
+        self.object_type = None # ['generic object', 'digit and letter', ...]
+    
+    @abstractmethod
+    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]]):
+        raise NotImplementedError
+    
+    def build(self):
+        if not hasattr(self, 'classes'):
+            raise AttributeError('attr `classes` is injected by `@dataset_register()`. '
+                                 'Your dataset class should be wrapped with @dataset_register().')
+        self.dataset = self.create_dataset(self.root_dir, self.split, self.transform, 
+                                           self.classes, self.ignore_classes, self.idx_map)
+        self.raw_classes = self.classes
+        self.classes = [i for i in self.classes if i not in self.ignore_classes]
+    
+    def __getitem__(self, idx):
+        if self.dataset is None:
+            raise AttributeError('Real dataset is build in `@dataset_register()`. '
+                                 'Your dataset class should be wrapped with @dataset_register().')
+        return self.dataset[idx]
+    
+    def __len__(self):
+        return len(self.dataset)
+    

data/datasets/action_recognition/__init__.py

@@ -0,0 +1,4 @@
+from .ucf101 import UCF101
+from .hmdb51 import HMDB51
+# from .kinetics400 import Kinetics400
+from .ixmas import IXMAS

data/datasets/action_recognition/common_dataset.py

@@ -0,0 +1,152 @@
+import os
+from pathlib import Path
+
+import random
+
+import numpy as np
+import pickle as pk
+import cv2
+from tqdm import tqdm
+from PIL import Image
+
+import torchvision.transforms as transforms
+import torch
+
+# from prefetch_generator import BackgroundGenerator
+from torch.utils.data import DataLoader, Dataset
+
+
+class VideoDataset(Dataset):
+
+    def __init__(self, directory_list, local_rank=0, enable_GPUs_num=0, distributed_load=False, resize_shape=[224, 224] , mode='train', clip_len=32, crop_size = 168):
+        
+        self.clip_len, self.crop_size, self.resize_shape = clip_len, crop_size, resize_shape
+        self.mode = mode
+
+        self.fnames, labels = [],[]
+        # get the directory of the specified split
+        for directory in directory_list:
+            folder = Path(directory)
+            print("Load dataset from folder : ", folder)
+            for label in sorted(os.listdir(folder)):
+                for fname in os.listdir(os.path.join(folder, label)) if mode=="train" else os.listdir(os.path.join(folder, label))[:10]:
+                    self.fnames.append(os.path.join(folder, label, fname))
+                    labels.append(label)
+        # print(labels)
+        random_list = list(zip(self.fnames, labels))
+        random.shuffle(random_list)
+        self.fnames[:], labels[:] = zip(*random_list)
+        self.labels = labels
+
+        # self.fnames = self.fnames[:240]
+
+        if mode == 'train' and distributed_load:
+            single_num_ = len(self.fnames)//enable_GPUs_num
+            self.fnames = self.fnames[local_rank*single_num_:((local_rank+1)*single_num_)]
+            labels = labels[local_rank*single_num_:((local_rank+1)*single_num_)]
+
+        # prepare a mapping between the label names (strings) and indices (ints)
+        self.label2index = {label:index for index, label in enumerate(sorted(set(labels)))} 
+        # convert the list of label names into an array of label indices
+        self.label_array = np.array([self.label2index[label] for label in labels], dtype=int)
+                
+    def __getitem__(self, index):
+        # loading and preprocessing. TODO move them to transform classess
+        buffer = self.loadvideo(self.fnames[index])
+        
+        height_index = np.random.randint(buffer.shape[2] - self.crop_size)
+        width_index = np.random.randint(buffer.shape[3] - self.crop_size)
+
+        return buffer[:,:,height_index:height_index + self.crop_size, width_index:width_index + self.crop_size], self.label_array[index]
+
+
+    def __len__(self):
+        return len(self.fnames)
+
+
+    def loadvideo(self, fname):
+        # initialize a VideoCapture object to read video data into a numpy array
+        self.transform = transforms.Compose([
+                transforms.Resize([self.resize_shape[0], self.resize_shape[1]]),
+                transforms.ToTensor(),
+                transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
+                ])
+
+        flip, flipCode = 1, random.choice([-1,0,1]) if np.random.random() < 0.5 and self.mode=="train" else 0
+
+        try:
+            video_stream = cv2.VideoCapture(fname)
+            frame_count = int(video_stream.get(cv2.CAP_PROP_FRAME_COUNT))
+        except RuntimeError:
+            index = np.random.randint(self.__len__())
+            video_stream = cv2.VideoCapture(self.fnames[index])
+            frame_count = int(video_stream.get(cv2.CAP_PROP_FRAME_COUNT))
+
+        while frame_count<self.clip_len+2:
+            index = np.random.randint(self.__len__())
+            video_stream = cv2.VideoCapture(self.fnames[index])
+            frame_count = int(video_stream.get(cv2.CAP_PROP_FRAME_COUNT))
+
+        speed_rate = np.random.randint(1, 3) if frame_count > self.clip_len*2+2 else 1
+        time_index = np.random.randint(frame_count - self.clip_len * speed_rate)
+
+        start_idx, end_idx, final_idx = time_index, time_index+(self.clip_len*speed_rate), frame_count-1
+        count, sample_count, retaining = 0, 0, True
+
+        # create a buffer. Must have dtype float, so it gets converted to a FloatTensor by Pytorch later
+        buffer = np.empty((self.clip_len, 3, self.resize_shape[0], self.resize_shape[1]), np.dtype('float32'))
+        
+        while (count <= end_idx and retaining):
+            retaining, frame = video_stream.read()
+            if count < start_idx:
+                count += 1
+                continue
+            if count % speed_rate == speed_rate-1 and count >= start_idx and sample_count < self.clip_len:
+                if flip:
+                    frame = cv2.flip(frame, flipCode=flipCode)
+                try:
+                    buffer[sample_count] = self.transform(Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)))
+                except cv2.error as err:
+                    continue
+                sample_count += 1
+            count += 1
+        video_stream.release()
+
+        return buffer.transpose((1, 0, 2, 3))
+
+
+if __name__ == '__main__':
+
+    datapath = ['/data/datasets/ucf101/videos']
+    
+    dataset = VideoDataset(datapath, 
+                            resize_shape=[224, 224],
+                            mode='validation')
+    x, y = dataset[0]
+    # x: (3, num_frames, w, h)
+    print(x.shape, y.shape, y)
+    
+    # dataloader = DataLoader(dataset, batch_size=8, shuffle=True, num_workers=24, pin_memory=True)
+
+    # bar = tqdm(total=len(dataloader), ncols=80)
+
+    # prefetcher = DataPrefetcher(BackgroundGenerator(dataloader), 0)
+    # batch = prefetcher.next()
+    # iter_id = 0
+    # while batch is not None:
+    #     iter_id += 1
+    #     bar.update(1)
+    #     if iter_id >= len(dataloader):
+    #         break
+
+    #     batch = prefetcher.next()
+    #     print(batch[0].shape)
+    #     print("label: ", batch[1])
+
+    # '''
+    # for step, (buffer, labels) in enumerate(BackgroundGenerator(dataloader)):
+    #     print(buffer.shape)
+    #     print("label: ", labels)
+    #     bar.update(1)
+    # '''
+

data/datasets/action_recognition/hmdb51.py

@@ -0,0 +1,45 @@
+from ..data_aug import cityscapes_like_image_train_aug, cityscapes_like_image_test_aug, cityscapes_like_label_aug
+# from torchvision.datasets import Cityscapes as RawCityscapes
+from ..ab_dataset import ABDataset
+from ..dataset_split import train_val_test_split
+import numpy as np
+from typing import Dict, List, Optional
+from torchvision.transforms import Compose, Lambda
+import os
+
+from .common_dataset import VideoDataset
+from ..registery import dataset_register
+
+
+@dataset_register(
+    name='HMDB51', 
+    classes=['brush_hair', 'cartwheel', 'catch', 'chew', 'clap', 'climb', 'climb_stairs', 'dive', 'draw_sword', 'dribble', 'drink', 'eat', 'fall_floor', 'fencing', 'flic_flac', 'golf', 'handstand', 'hit', 'hug', 'jump', 'kick', 'kick_ball', 'kiss', 'laugh', 'pick', 'pour', 'pullup', 'punch', 'push', 'pushup', 'ride_bike', 'ride_horse', 'run', 'shake_hands', 'shoot_ball', 'shoot_bow', 'shoot_gun', 'sit', 'situp', 'smile', 'smoke', 'somersault', 'stand', 'swing_baseball', 'sword', 'sword_exercise', 'talk', 'throw', 'turn', 'walk', 'wave'],
+    task_type='Action Recognition',
+    object_type='Web Video',
+    # class_aliases=[['automobile', 'car']],
+    class_aliases=[],
+    shift_type=None
+)
+class HMDB51(ABDataset): # just for demo now
+    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]]):
+        # if transform is None:
+        #     x_transform = cityscapes_like_image_train_aug() if split == 'train' else cityscapes_like_image_test_aug()
+        #     y_transform = cityscapes_like_label_aug()
+        #     self.transform = x_transform
+        # else:
+        #     x_transform, y_transform = transform
+        
+        dataset = VideoDataset([root_dir], mode='train')
+        
+        if len(ignore_classes) > 0: 
+            for ignore_class in ignore_classes:
+                ci = classes.index(ignore_class)
+                dataset.fnames = [img for img, label in zip(dataset.fnames, dataset.label_array) if label != ci]
+                dataset.label_array = [label for label in dataset.label_array if label != ci]
+                
+        if idx_map is not None:
+            dataset.label_array = [idx_map[label] for label in dataset.label_array]
+        
+        dataset = train_val_test_split(dataset, split)
+        return dataset

data/datasets/action_recognition/ixmas.py

@@ -0,0 +1,45 @@
+from ..data_aug import cityscapes_like_image_train_aug, cityscapes_like_image_test_aug, cityscapes_like_label_aug
+# from torchvision.datasets import Cityscapes as RawCityscapes
+from ..ab_dataset import ABDataset
+from ..dataset_split import train_val_test_split
+import numpy as np
+from typing import Dict, List, Optional
+from torchvision.transforms import Compose, Lambda
+import os
+
+from .common_dataset import VideoDataset
+from ..registery import dataset_register
+
+
+@dataset_register(
+    name='IXMAS', 
+    classes=['check_watch', 'cross_arms', 'get_up', 'kick', 'pick_up', 'point', 'punch', 'scratch_head', 'sit_down', 'turn_around', 'walk', 'wave'],
+    task_type='Action Recognition',
+    object_type='Web Video',
+    # class_aliases=[['automobile', 'car']],
+    class_aliases=[],
+    shift_type=None
+)
+class IXMAS(ABDataset): # just for demo now
+    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]]):
+        # if transform is None:
+        #     x_transform = cityscapes_like_image_train_aug() if split == 'train' else cityscapes_like_image_test_aug()
+        #     y_transform = cityscapes_like_label_aug()
+        #     self.transform = x_transform
+        # else:
+        #     x_transform, y_transform = transform
+        
+        dataset = VideoDataset([root_dir], mode='train')
+        
+        if len(ignore_classes) > 0: 
+            for ignore_class in ignore_classes:
+                ci = classes.index(ignore_class)
+                dataset.fnames = [img for img, label in zip(dataset.fnames, dataset.label_array) if label != ci]
+                dataset.label_array = [label for label in dataset.label_array if label != ci]
+                
+        if idx_map is not None:
+            dataset.label_array = [idx_map[label] for label in dataset.label_array]
+        
+        dataset = train_val_test_split(dataset, split)
+        return dataset

data/datasets/action_recognition/kinetics400.py

@@ -0,0 +1,51 @@
+from ..data_aug import cityscapes_like_image_train_aug, cityscapes_like_image_test_aug, cityscapes_like_label_aug
+# from torchvision.datasets import Cityscapes as RawCityscapes
+from ..ab_dataset import ABDataset
+from ..dataset_split import train_val_split, train_val_test_split
+import numpy as np
+from typing import Dict, List, Optional
+from torchvision.transforms import Compose, Lambda
+import os
+
+from .common_dataset import VideoDataset
+from ..registery import dataset_register
+
+
+@dataset_register(
+    name='Kinetics400', 
+    classes=['abseiling', 'air drumming', 'answering questions', 'applauding', 'applying cream', 'archery', 'arm wrestling', 'arranging flowers', 'assembling computer', 'auctioning', 'baby waking up', 'baking cookies', 'balloon blowing', 'bandaging', 'barbequing', 'bartending', 'beatboxing', 'bee keeping', 'belly dancing', 'bench pressing', 'bending back', 'bending metal', 'biking through snow', 'blasting sand', 'blowing glass', 'blowing leaves', 'blowing nose', 'blowing out candles', 'bobsledding', 'bookbinding', 'bouncing on trampoline', 'bowling', 'braiding hair', 'breading or breadcrumbing', 'breakdancing', 'brush painting', 'brushing hair', 'brushing teeth', 'building cabinet', 'building shed', 'bungee jumping', 'busking', 'canoeing or kayaking', 'capoeira', 'carrying baby', 'cartwheeling', 'carving pumpkin', 'catching fish', 'catching or throwing baseball', 'catching or throwing frisbee', 'catching or throwing softball', 'celebrating', 'changing oil', 'changing wheel', 'checking tires', 'cheerleading', 'chopping wood', 'clapping', 'clay pottery making', 'clean and jerk', 'cleaning floor', 'cleaning gutters', 'cleaning pool', 'cleaning shoes', 'cleaning toilet', 'cleaning windows', 'climbing a rope', 'climbing ladder', 'climbing tree', 'contact juggling', 'cooking chicken', 'cooking egg', 'cooking on campfire', 'cooking sausages', 'counting money', 'country line dancing', 'cracking neck', 'crawling baby', 'crossing river', 'crying', 'curling hair', 'cutting nails', 'cutting pineapple', 'cutting watermelon', 'dancing ballet', 'dancing charleston', 'dancing gangnam style', 'dancing macarena', 'deadlifting', 'decorating the christmas tree', 'digging', 'dining', 'disc golfing', 'diving cliff', 'dodgeball', 'doing aerobics', 'doing laundry', 'doing nails', 'drawing', 'dribbling basketball', 'drinking', 'drinking beer', 'drinking shots', 'driving car', 'driving tractor', 'drop kicking', 'drumming fingers', 'dunking basketball', 'dying hair', 'eating burger', 'eating cake', 'eating carrots', 'eating chips', 'eating doughnuts', 'eating hotdog', 'eating ice cream', 'eating spaghetti', 'eating watermelon', 'egg hunting', 'exercising arm', 'exercising with an exercise ball', 'extinguishing fire', 'faceplanting', 'feeding birds', 'feeding fish', 'feeding goats', 'filling eyebrows', 'finger snapping', 'fixing hair', 'flipping pancake', 'flying kite', 'folding clothes', 'folding napkins', 'folding paper', 'front raises', 'frying vegetables', 'garbage collecting', 'gargling', 'getting a haircut', 'getting a tattoo', 'giving or receiving award', 'golf chipping', 'golf driving', 'golf putting', 'grinding meat', 'grooming dog', 'grooming horse', 'gymnastics tumbling', 'hammer throw', 'headbanging', 'headbutting', 'high jump', 'high kick', 'hitting baseball', 'hockey stop', 'holding snake', 'hopscotch', 'hoverboarding', 'hugging', 'hula hooping', 'hurdling', 'hurling (sport)', 'ice climbing', 'ice fishing', 'ice skating', 'ironing', 'javelin throw', 'jetskiing', 'jogging', 'juggling balls', 'juggling fire', 'juggling soccer ball', 'jumping into pool', 'jumpstyle dancing', 'kicking field goal', 'kicking soccer ball', 'kissing', 'kitesurfing', 'knitting', 'krumping', 'laughing', 'laying bricks', 'long jump', 'lunge', 'making a cake', 'making a sandwich', 'making bed', 'making jewelry', 'making pizza', 'making snowman', 'making sushi', 'making tea', 'marching', 'massaging back', 'massaging feet', 'massaging legs', "massaging person's head", 'milking cow', 'mopping floor', 'motorcycling', 'moving furniture', 'mowing lawn', 'news anchoring', 'opening bottle', 'opening present', 'paragliding', 'parasailing', 'parkour', 'passing American football (in game)', 'passing American football (not in game)', 'peeling apples', 'peeling potatoes', 'petting animal (not cat)', 'petting cat', 'picking fruit', 'planting trees', 'plastering', 'playing accordion', 'playing badminton', 'playing bagpipes', 'playing basketball', 'playing bass guitar', 'playing cards', 'playing cello', 'playing chess', 'playing clarinet', 'playing controller', 'playing cricket', 'playing cymbals', 'playing didgeridoo', 'playing drums', 'playing flute', 'playing guitar', 'playing harmonica', 'playing harp', 'playing ice hockey', 'playing keyboard', 'playing kickball', 'playing monopoly', 'playing organ', 'playing paintball', 'playing piano', 'playing poker', 'playing recorder', 'playing saxophone', 'playing squash or racquetball', 'playing tennis', 'playing trombone', 'playing trumpet', 'playing ukulele', 'playing violin', 'playing volleyball', 'playing xylophone', 'pole vault', 'presenting weather forecast', 'pull ups', 'pumping fist', 'pumping gas', 'punching bag', 'punching person (boxing)', 'push up', 'pushing car', 'pushing cart', 'pushing wheelchair', 'reading book', 'reading newspaper', 'recording music', 'riding a bike', 'riding camel', 'riding elephant', 'riding mechanical bull', 'riding mountain bike', 'riding mule', 'riding or walking with horse', 'riding scooter', 'riding unicycle', 'ripping paper', 'robot dancing', 'rock climbing', 'rock scissors paper', 'roller skating', 'running on treadmill', 'sailing', 'salsa dancing', 'sanding floor', 'scrambling eggs', 'scuba diving', 'setting table', 'shaking hands', 'shaking head', 'sharpening knives', 'sharpening pencil', 'shaving head', 'shaving legs', 'shearing sheep', 'shining shoes', 'shooting basketball', 'shooting goal (soccer)', 'shot put', 'shoveling snow', 'shredding paper', 'shuffling cards', 'side kick', 'sign language interpreting', 'singing', 'situp', 'skateboarding', 'ski jumping', 'skiing (not slalom or crosscountry)', 'skiing crosscountry', 'skiing slalom', 'skipping rope', 'skydiving', 'slacklining', 'slapping', 'sled dog racing', 'smoking', 'smoking hookah', 'snatch weight lifting', 'sneezing', 'sniffing', 'snorkeling', 'snowboarding', 'snowkiting', 'snowmobiling', 'somersaulting', 'spinning poi', 'spray painting', 'spraying', 'springboard diving', 'squat', 'sticking tongue out', 'stomping grapes', 'stretching arm', 'stretching leg', 'strumming guitar', 'surfing crowd', 'surfing water', 'sweeping floor', 'swimming backstroke', 'swimming breast stroke', 'swimming butterfly stroke', 'swing dancing', 'swinging legs', 'swinging on something', 'sword fighting', 'tai chi', 'taking a shower', 'tango dancing', 'tap dancing', 'tapping guitar', 'tapping pen', 'tasting beer', 'tasting food', 'testifying', 'texting', 'throwing axe', 'throwing ball', 'throwing discus', 'tickling', 'tobogganing', 'tossing coin', 'tossing salad', 'training dog', 'trapezing', 'trimming or shaving beard', 'trimming trees', 'triple jump', 'tying bow tie', 'tying knot (not on a tie)', 'tying tie', 'unboxing', 'unloading truck', 'using computer', 'using remote controller (not gaming)', 'using segway', 'vault', 'waiting in line', 'walking the dog', 'washing dishes', 'washing feet', 'washing hair', 'washing hands', 'water skiing', 'water sliding', 'watering plants', 'waxing back', 'waxing chest', 'waxing eyebrows', 'waxing legs', 'weaving basket', 'welding', 'whistling', 'windsurfing', 'wrapping present', 'wrestling', 'writing', 'yawning', 'yoga', 'zumba'],
+    task_type='Action Recognition',
+    object_type='Web Video',
+    # class_aliases=[['automobile', 'car']],
+    class_aliases=[],
+    shift_type=None
+)
+class Kinetics400(ABDataset): # just for demo now
+    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]]):
+        # if transform is None:
+        #     x_transform = cityscapes_like_image_train_aug() if split == 'train' else cityscapes_like_image_test_aug()
+        #     y_transform = cityscapes_like_label_aug()
+        #     self.transform = x_transform
+        # else:
+        #     x_transform, y_transform = transform
+        
+        if split == 'test':
+            root_dir = os.path.join(root_dir, 'videos_val')
+        else:
+            root_dir = os.path.join(root_dir, 'videos_train')
+        # print(root_dir)
+        dataset = VideoDataset([root_dir], mode='train')
+        
+        if len(ignore_classes) > 0: 
+            for ignore_class in ignore_classes:
+                ci = classes.index(ignore_class)
+                dataset.fnames = [img for img, label in zip(dataset.fnames, dataset.label_array) if label != ci]
+                dataset.label_array = [label for label in dataset.label_array if label != ci]
+                
+        if idx_map is not None:
+            dataset.label_array = [idx_map[label] for label in dataset.label_array]
+        
+        if split != 'test':
+            dataset = train_val_split(dataset, split)
+        return dataset

data/datasets/action_recognition/ucf101.py

@@ -0,0 +1,45 @@
+from ..data_aug import cityscapes_like_image_train_aug, cityscapes_like_image_test_aug, cityscapes_like_label_aug
+# from torchvision.datasets import Cityscapes as RawCityscapes
+from ..ab_dataset import ABDataset
+from ..dataset_split import train_val_test_split
+import numpy as np
+from typing import Dict, List, Optional
+from torchvision.transforms import Compose, Lambda
+import os
+
+from .common_dataset import VideoDataset
+from ..registery import dataset_register
+
+
+@dataset_register(
+    name='UCF101', 
+    classes=['apply_eye_makeup', 'apply_lipstick', 'archery', 'baby_crawling', 'balance_beam', 'band_marching', 'baseball_pitch', 'basketball', 'basketball_dunk', 'bench_press', 'biking', 'billiards', 'blow_dry_hair', 'blowing_candles', 'body_weight_squats', 'bowling', 'boxing_punching_bag', 'boxing_speed_bag', 'breast_stroke', 'brushing_teeth', 'clean_and_jerk', 'cliff_diving', 'cricket_bowling', 'cricket_shot', 'cutting_in_kitchen', 'diving', 'drumming', 'fencing', 'field_hockey_penalty', 'floor_gymnastics', 'frisbee_catch', 'front_crawl', 'golf_swing', 'haircut', 'hammer_throw', 'hammering', 'handstand_pushups', 'handstand_walking', 'head_massage', 'high_jump', 'horse_race', 'horse_riding', 'hula_hoop', 'ice_dancing', 'javelin_throw', 'juggling_balls', 'jump_rope', 'jumping_jack', 'kayaking', 'knitting', 'long_jump', 'lunges', 'military_parade', 'mixing', 'mopping_floor', 'nunchucks', 'parallel_bars', 'pizza_tossing', 'playing_cello', 'playing_daf', 'playing_dhol', 'playing_flute', 'playing_guitar', 'playing_piano', 'playing_sitar', 'playing_tabla', 'playing_violin', 'pole_vault', 'pommel_horse', 'pull_ups', 'punch', 'push_ups', 'rafting', 'rock_climbing_indoor', 'rope_climbing', 'rowing', 'salsa_spin', 'shaving_beard', 'shotput', 'skate_boarding', 'skiing', 'skijet', 'sky_diving', 'soccer_juggling', 'soccer_penalty', 'still_rings', 'sumo_wrestling', 'surfing', 'swing', 'table_tennis_shot', 'tai_chi', 'tennis_swing', 'throw_discus', 'trampoline_jumping', 'typing', 'uneven_bars', 'volleyball_spiking', 'walking_with_dog', 'wall_pushups', 'writing_on_board', 'yo_yo'],
+    task_type='Action Recognition',
+    object_type='Web Video',
+    # class_aliases=[['automobile', 'car']],
+    class_aliases=[],
+    shift_type=None
+)
+class UCF101(ABDataset): # just for demo now
+    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]]):
+        # if transform is None:
+        #     x_transform = cityscapes_like_image_train_aug() if split == 'train' else cityscapes_like_image_test_aug()
+        #     y_transform = cityscapes_like_label_aug()
+        #     self.transform = x_transform
+        # else:
+        #     x_transform, y_transform = transform
+        
+        dataset = VideoDataset([root_dir], mode='train')
+        
+        if len(ignore_classes) > 0: 
+            for ignore_class in ignore_classes:
+                ci = classes.index(ignore_class)
+                dataset.fnames = [img for img, label in zip(dataset.fnames, dataset.label_array) if label != ci]
+                dataset.label_array = [label for label in dataset.label_array if label != ci]
+                
+        if idx_map is not None:
+            dataset.label_array = [idx_map[label] for label in dataset.label_array]
+        
+        dataset = train_val_test_split(dataset, split)
+        return dataset

data/datasets/data_aug.py

@@ -0,0 +1,93 @@
+from torchvision import transforms
+import torch
+
+
+def one_d_image_train_aug(to_3_channels=False):
+    mean, std = (0.1307, 0.1307, 0.1307), (0.3081, 0.3081, 0.3081)
+    return transforms.Compose([
+        transforms.Resize(32),
+        # transforms.RandomCrop(32, padding=4),
+        transforms.ToTensor(),
+        transforms.Lambda((lambda x: torch.cat([x] * 3)) if to_3_channels else (lambda x: x)),
+        transforms.Normalize(mean, std)
+    ])
+
+
+def one_d_image_test_aug(to_3_channels=False):
+    mean, std = (0.1307, 0.1307, 0.1307), (0.3081, 0.3081, 0.3081)
+    return transforms.Compose([
+        transforms.Resize(32),
+        transforms.ToTensor(),
+        transforms.Lambda((lambda x: torch.cat([x] * 3)) if to_3_channels else (lambda x: x)),
+        transforms.Normalize(mean, std)
+    ])
+
+
+def cifar_like_image_train_aug(mean=None, std=None):
+    if mean is None:
+        mean, std = (0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)
+    return transforms.Compose([
+        transforms.Resize(40), # NOTE: this is critical!!! or you may crop a small part of an image
+        transforms.RandomCrop(32, padding=4),
+        transforms.RandomHorizontalFlip(),
+        transforms.ToTensor(),
+        transforms.Normalize(mean, std)
+    ])
+
+
+def cifar_like_image_test_aug(mean=None, std=None):
+    if mean is None:
+        mean, std = (0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)
+    return transforms.Compose([
+        transforms.Resize(32),
+        transforms.ToTensor(),
+        transforms.Normalize(mean, std)
+    ])
+
+def imagenet_like_image_train_aug():
+    mean, std = [0.485, 0.456, 0.406], [0.229, 0.224, 0.225]
+    return transforms.Compose([
+        transforms.Resize((256, 256)),
+        transforms.RandomCrop((224, 224), padding=4),
+        transforms.RandomHorizontalFlip(),
+        transforms.ToTensor(),
+        transforms.Normalize(mean, std)
+    ])
+
+
+def imagenet_like_image_test_aug():
+    mean, std = [0.485, 0.456, 0.406], [0.229, 0.224, 0.225]
+    return transforms.Compose([
+        transforms.Resize((224, 224)),
+        transforms.ToTensor(),
+        transforms.Normalize(mean, std)
+    ])
+
+
+def cityscapes_like_image_train_aug():
+    return transforms.Compose([
+        transforms.Resize((224, 224)),
+        transforms.ToTensor(),
+        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+    ])
+    
+def cityscapes_like_image_test_aug():
+    return transforms.Compose([
+        transforms.Resize((224, 224)),
+        transforms.ToTensor(),
+        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+    ])
+    
+def cityscapes_like_label_aug():
+    import numpy as np
+    return transforms.Compose([
+        transforms.Resize((224, 224)),
+        transforms.Lambda(lambda x: torch.from_numpy(np.array(x)).long())
+    ])
+
+
+def pil_image_to_tensor(img_size=224):
+    return transforms.Compose([
+        transforms.Resize((img_size, img_size)), 
+        transforms.ToTensor()
+    ])

data/datasets/dataset_cache.py

@@ -0,0 +1,40 @@
+from typing import List, Optional, Dict
+import os
+import torch
+from utils.common.log import logger
+import hashlib
+
+
+def get_dataset_cache_path(root_dir: str, 
+                         classes: List[str], ignore_classes: List[str], idx_map: Optional[Dict[int, int]]):
+    
+    def _hash(o):
+        if isinstance(o, list):
+            o = sorted(o)
+        elif isinstance(o, dict):
+            o = {k: o[k] for k in sorted(o)}
+        elif isinstance(o, set):
+            o = sorted(list(o))
+        # else:
+        #     print(type(o))
+        
+        obj = hashlib.md5()
+        obj.update(str(o).encode('utf-8'))
+        return obj.hexdigest()
+    
+    cache_key = _hash(f'zql_data_{_hash(root_dir)}_{_hash(classes)}_{_hash(ignore_classes)}_{_hash(idx_map)}.cache')
+    
+    # print(root_dir, classes, ignore_classes, idx_map)
+    # print('cache key', cache_key)
+    
+    cache_file_path = os.path.join('/tmp', f'./zql_data_cache_{cache_key}.cache')
+    return cache_file_path
+
+
+def cache_dataset_status(status, cache_file_path, dataset_name):
+    logger.info(f'cache dataset status: {dataset_name}')
+    torch.save(status, cache_file_path)
+    
+def read_cached_dataset_status(cache_file_path, dataset_name):
+    logger.info(f'read dataset cache: {dataset_name}')
+    return torch.load(cache_file_path)

data/datasets/dataset_split.py

@@ -0,0 +1,81 @@
+import torch
+import os
+import numpy as np
+from .ab_dataset import ABDataset
+
+
+class _SplitDataset(torch.utils.data.Dataset):
+    """Used by split_dataset"""
+
+    def __init__(self, underlying_dataset, keys):
+        super(_SplitDataset, self).__init__()
+        self.underlying_dataset = underlying_dataset
+        self.keys = keys
+
+    def __getitem__(self, key):
+        return self.underlying_dataset[self.keys[key]]
+
+    def __len__(self):
+        return len(self.keys)
+
+
+def split_dataset(dataset, n, seed=0, transform=None):
+    
+    if isinstance(dataset, ABDataset):
+        if dataset.task_type == 'Object Detection':
+            return split_dataset_det(dataset, n, seed)
+        if dataset.task_type == 'MM Object Detection':
+            return split_dataset_det_mm(dataset, n, seed, transform=transform)
+    
+    """
+    Return a pair of datasets corresponding to a random split of the given
+    dataset, with n datapoints in the first dataset and the rest in the last,
+    using the given random seed
+    """
+    assert n <= len(dataset), f'{n}_{len(dataset)}'
+
+    cache_p = f'{n}_{seed}_{len(dataset)}'
+    cache_p = os.path.join(os.path.expanduser(
+        '~'), '.domain_benchmark_split_dataset_cache_' + str(cache_p))
+    if os.path.exists(cache_p):
+        keys_1, keys_2 = torch.load(cache_p)
+    else:
+        keys = list(range(len(dataset)))
+        np.random.RandomState(seed).shuffle(keys)
+        keys_1 = keys[:n]
+        keys_2 = keys[n:]
+        torch.save((keys_1, keys_2), cache_p)
+    
+    return _SplitDataset(dataset, keys_1), _SplitDataset(dataset, keys_2)
+
+
+def train_val_split(dataset, split):
+    assert split in ['train', 'val']
+    if split == 'train':
+        return split_dataset(dataset, int(len(dataset) * 0.8))[0]
+    else:
+        return split_dataset(dataset, int(len(dataset) * 0.8))[1]
+
+    
+def train_val_test_split(dataset, split):
+    assert split in ['train', 'val', 'test']
+
+    train_set, test_set = split_dataset(dataset, int(len(dataset) * 0.8))
+    train_set, val_set = split_dataset(train_set, int(len(train_set) * 0.8))
+    
+    return {'train': train_set, 'val': val_set, 'test': test_set}[split]
+
+
+def split_dataset_det(dataset: ABDataset, n, seed=0):
+    coco_ann_json_path = dataset.ann_json_file_path_for_split
+    from .object_detection.yolox_data_util.api import coco_split, get_default_yolox_coco_dataset
+    split_coco_ann_json_path = coco_split(coco_ann_json_path, ratio=n / len(dataset))[0]
+    # print(n, len(dataset))
+    return get_default_yolox_coco_dataset(dataset.root_dir, split_coco_ann_json_path, train=dataset.split == 'train'), None
+
+def split_dataset_det_mm(dataset: ABDataset, n, seed=0, transform=None):
+    coco_ann_json_path = dataset.ann_json_file_path_for_split
+    from .object_detection.yolox_data_util.api import coco_split, get_yolox_coco_dataset_with_caption
+    split_coco_ann_json_path = coco_split(coco_ann_json_path, ratio=n / len(dataset))[0]
+    # print(n, len(dataset))
+    return get_yolox_coco_dataset_with_caption(dataset.root_dir, split_coco_ann_json_path, transform=transform, train=dataset.split == 'train', classes=dataset.classes), None

data/datasets/image_classification/__init__.py

@@ -0,0 +1,24 @@
+from .mnist import MNIST
+from .usps import USPS
+from .svhn import SVHN
+from .emnist import EMNIST
+from .cifar10 import CIFAR10
+from .stl10 import STL10
+from .imagenet import ImageNet
+from .imagenet_a import ImageNetA
+from .caltech256 import Caltech256
+from .domainnet_real import DomainNetReal
+from .synsigns import SYNSIGNS
+from .gtsrb import GTSRB
+
+from .cifar10_single import CIFAR10Single
+from .stl10_single import STL10Single
+from .mnist_single import MNISTSingle
+from .usps_single import USPSSingle
+from .svhn_single import SVHNSingle
+
+from .baidu_person_cls import BaiduPersonCls
+from .cityscapes_cls import CityscapesCls
+from .gta5_cls import GTA5Cls
+from .supervisely_person_cls import SuperviselyPersonCls
+from .coco_cls import COCOCls