| 33 | """ |
| 34 | |
| 35 | def __init__(self, |
| 36 | data_root: Union[dict, List], |
| 37 | ann_file: Union[dict, List, str], |
| 38 | verbose: bool = False, |
| 39 | color_setting: str = None, |
| 40 | thickness: float = 0.01): |
| 41 | |
| 42 | if isinstance(ann_file, dict): |
| 43 | ann_file = list(ann_file.values()) |
| 44 | elif isinstance(ann_file, str): |
| 45 | ann_file = [ann_file] |
| 46 | self.ann_files = ann_file |
| 47 | |
| 48 | if isinstance(data_root, str): |
| 49 | data_root = [data_root] |
| 50 | if isinstance(data_root, list): |
| 51 | self.data_root = dict() |
| 52 | for dataset in DATASETS: |
| 53 | self.data_root[dataset] = None |
| 54 | for root in data_root: |
| 55 | for dataset in DATASETS: |
| 56 | if dataset.lower() in root.lower(): |
| 57 | self.data_root[dataset] = root |
| 58 | break |
| 59 | if isinstance(data_root, dict): |
| 60 | self.data_root = data_root |
| 61 | |
| 62 | self.verbose = verbose |
| 63 | self.thickness = thickness |
| 64 | |
| 65 | if self.verbose: |
| 66 | print('Dataset root') |
| 67 | for dataset in DATASETS: |
| 68 | print(dataset, ':', self.data_root.get(dataset, None)) |
| 69 | |
| 70 | if self.verbose: |
| 71 | print('Loading') |
| 72 | self.metainfo = None |
| 73 | data_list = [] |
| 74 | for file in self.ann_files: |
| 75 | if isinstance(file, list): |
| 76 | data_list += file |
| 77 | continue |
| 78 | elif isinstance(file, dict): |
| 79 | if 'data_list' in file: |
| 80 | data = file |
| 81 | else: |
| 82 | data_list.append(file) |
| 83 | continue |
| 84 | elif isinstance(file, str): |
| 85 | with open(file, 'rb') as f: |
| 86 | data = pickle.load(f) |
| 87 | if self.metainfo is None: |
| 88 | self.metainfo = data['metainfo'] |
| 89 | else: |
| 90 | assert self.metainfo == data['metainfo'] |
| 91 | data_list += data['data_list'] |
| 92 | |