Base motion dataset. Args: data_prefix (str): the prefix of data path. pipeline (list): a list of dict, where each element represents a operation defined in `diffplanner.datasets.pipelines`. ann_file (str | None, optional): the annotation file. When ann_file i
| 12 | |
| 13 | @DATASETS.register_module() |
| 14 | class BaseMotionDataset(Dataset): |
| 15 | """Base motion dataset. |
| 16 | Args: |
| 17 | data_prefix (str): the prefix of data path. |
| 18 | pipeline (list): a list of dict, where each element represents |
| 19 | a operation defined in `diffplanner.datasets.pipelines`. |
| 20 | ann_file (str | None, optional): the annotation file. When ann_file is |
| 21 | str, the subclass is expected to read from the ann_file. When |
| 22 | ann_file is None, the subclass is expected to read according |
| 23 | to data_prefix. |
| 24 | test_mode (bool): in train mode or test mode. Default: None. |
| 25 | dataset_name (str | None, optional): the name of dataset. It is used |
| 26 | to identify the type of evaluation metric. Default: None. |
| 27 | """ |
| 28 | |
| 29 | def __init__(self, |
| 30 | data_prefix: str, |
| 31 | pipeline: list, |
| 32 | dataset_name: Optional[Union[str, None]] = None, |
| 33 | fixed_length: Optional[Union[int, None]] = None, |
| 34 | ann_file: Optional[Union[str, None]] = None, |
| 35 | motion_dir: Optional[Union[str, None]] = None, |
| 36 | eval_cfg: Optional[Union[dict, None]] = None, |
| 37 | test_mode: Optional[bool] = False): |
| 38 | super(BaseMotionDataset, self).__init__() |
| 39 | |
| 40 | self.data_prefix = data_prefix |
| 41 | self.pipeline = Compose(pipeline) |
| 42 | self.dataset_name = dataset_name |
| 43 | self.fixed_length = fixed_length |
| 44 | self.ann_file = os.path.join(data_prefix, 'datasets', dataset_name, ann_file) |
| 45 | self.motion_dir = os.path.join(data_prefix, 'datasets', dataset_name, motion_dir) |
| 46 | self.eval_cfg = copy.deepcopy(eval_cfg) |
| 47 | self.test_mode = test_mode |
| 48 | |
| 49 | self.load_annotations() |
| 50 | if self.test_mode: |
| 51 | self.prepare_evaluation() |
| 52 | |
| 53 | def load_anno(self, name): |
| 54 | motion_path = os.path.join(self.motion_dir, name + '.npy') |
| 55 | motion_data = np.load(motion_path) |
| 56 | return {'motion': motion_data} |
| 57 | |
| 58 | |
| 59 | def load_annotations(self): |
| 60 | """Load annotations from ``ann_file`` to ``data_infos``""" |
| 61 | self.data_infos = [] |
| 62 | for line in open(self.ann_file, 'r').readlines()[:]: |
| 63 | line = line.strip() |
| 64 | self.data_infos.append(self.load_anno(line)) |
| 65 | |
| 66 | |
| 67 | def prepare_data(self, idx: int): |
| 68 | """"Prepare raw data for the f'{idx'}-th data.""" |
| 69 | results = copy.deepcopy(self.data_infos[idx]) |
| 70 | results['dataset_name'] = self.dataset_name |
| 71 | results['sample_idx'] = idx |
nothing calls this directly
no outgoing calls
no test coverage detected