Spaces:
Runtime error
Runtime error
File size: 3,685 Bytes
cc0dd3c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 |
# Copyright (c) OpenMMLab. All rights reserved.
from copy import deepcopy
from typing import Any, Callable, List, Tuple, Union
from mmengine.dataset import BaseDataset
from mmengine.registry import build_from_cfg
from mmpose.registry import DATASETS
from .datasets.utils import parse_pose_metainfo
@DATASETS.register_module()
class CombinedDataset(BaseDataset):
"""A wrapper of combined dataset.
Args:
metainfo (dict): The meta information of combined dataset.
datasets (list): The configs of datasets to be combined.
pipeline (list, optional): Processing pipeline. Defaults to [].
"""
def __init__(self,
metainfo: dict,
datasets: list,
pipeline: List[Union[dict, Callable]] = [],
**kwargs):
self.datasets = []
for cfg in datasets:
dataset = build_from_cfg(cfg, DATASETS)
self.datasets.append(dataset)
self._lens = [len(dataset) for dataset in self.datasets]
self._len = sum(self._lens)
super(CombinedDataset, self).__init__(pipeline=pipeline, **kwargs)
self._metainfo = parse_pose_metainfo(metainfo)
@property
def metainfo(self):
return deepcopy(self._metainfo)
def __len__(self):
return self._len
def _get_subset_index(self, index: int) -> Tuple[int, int]:
"""Given a data sample's global index, return the index of the sub-
dataset the data sample belongs to, and the local index within that
sub-dataset.
Args:
index (int): The global data sample index
Returns:
tuple[int, int]:
- subset_index (int): The index of the sub-dataset
- local_index (int): The index of the data sample within
the sub-dataset
"""
if index >= len(self) or index < -len(self):
raise ValueError(
f'index({index}) is out of bounds for dataset with '
f'length({len(self)}).')
if index < 0:
index = index + len(self)
subset_index = 0
while index >= self._lens[subset_index]:
index -= self._lens[subset_index]
subset_index += 1
return subset_index, index
def prepare_data(self, idx: int) -> Any:
"""Get data processed by ``self.pipeline``.The source dataset is
depending on the index.
Args:
idx (int): The index of ``data_info``.
Returns:
Any: Depends on ``self.pipeline``.
"""
data_info = self.get_data_info(idx)
return self.pipeline(data_info)
def get_data_info(self, idx: int) -> dict:
"""Get annotation by index.
Args:
idx (int): Global index of ``CombinedDataset``.
Returns:
dict: The idx-th annotation of the datasets.
"""
subset_idx, sample_idx = self._get_subset_index(idx)
# Get data sample processed by ``subset.pipeline``
data_info = self.datasets[subset_idx][sample_idx]
# Add metainfo items that are required in the pipeline and the model
metainfo_keys = [
'upper_body_ids', 'lower_body_ids', 'flip_pairs',
'dataset_keypoint_weights', 'flip_indices'
]
for key in metainfo_keys:
data_info[key] = deepcopy(self._metainfo[key])
return data_info
def full_init(self):
"""Fully initialize all sub datasets."""
if self._fully_initialized:
return
for dataset in self.datasets:
dataset.full_init()
self._fully_initialized = True
|