MCPcopy Create free account
hub / github.com/MotrixLab/FineMoGen / BaseMotionDataset

Class BaseMotionDataset

mogen/datasets/base_dataset.py:18–124  ·  view source on GitHub ↗

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 `mogen.datasets.pipelines`. ann_file (str | None, optional): the annotation file. When ann_file is

Source from the content-addressed store, hash-verified

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

Callers

nothing calls this directly

Calls

no outgoing calls

Tested by

no test coverage detected