Spaces:
Sleeping
Sleeping
| from lib.kits.basic import * | |
| from lib.data.datasets.hsmr_v1.mocap_dataset import MoCapDataset | |
| from lib.data.datasets.hsmr_v1.wds_loader import load_tars_as_wds | |
| import webdataset as wds | |
| class MixedWebDataset(wds.WebDataset): | |
| def __init__(self) -> None: | |
| super(wds.WebDataset, self).__init__() | |
| class DataModule(pl.LightningDataModule): | |
| def __init__(self, name, cfg): | |
| super().__init__() | |
| self.name = name | |
| self.cfg = cfg | |
| self.cfg_eval = self.cfg.pop('eval', None) | |
| self.cfg_train = self.cfg.pop('train', None) | |
| self.cfg_mocap = self.cfg.pop('mocap', None) | |
| def setup(self, stage=None): | |
| if stage in ['test', None, '_debug_eval'] and self.cfg_eval is not None: | |
| # get_logger().info('Evaluation dataset will be enabled.') | |
| self._setup_eval() | |
| if stage in ['fit', None, '_debug_train'] and self.cfg_train is not None: | |
| # get_logger().info('Training dataset will be enabled.') | |
| self._setup_train() | |
| if stage in ['fit', None, '_debug_mocap'] and self.cfg_mocap is not None: | |
| # get_logger().info('Mocap dataset will be enabled.') | |
| self._setup_mocap() | |
| def train_dataloader(self): | |
| img_dataset = torch.utils.data.DataLoader( | |
| dataset = self.train_dataset, | |
| **self.cfg_train.dataloader, | |
| ) | |
| ret = {'img_ds' : img_dataset } | |
| if self.cfg_mocap is not None: | |
| mocap_dataset = torch.utils.data.DataLoader( | |
| dataset = self.mocap_dataset, | |
| **self.cfg_mocap.dataloader, | |
| ) | |
| ret['mocap_ds'] = mocap_dataset | |
| return ret | |
| # ========== Internal Modules to Setup Datasets ========== | |
| def _setup_train(self): | |
| names, datasets, weights = [], [], [] | |
| ld_cfg = self.cfg_train.cfg # cfg for initializing wds loading pipeline | |
| for ds_cfg in self.cfg_train.datasets: | |
| dataset = load_tars_as_wds( | |
| ld_cfg, | |
| ds_cfg.item.urls, | |
| ds_cfg.item.epoch_size | |
| ) | |
| names.append(ds_cfg.name) | |
| datasets.append(dataset) | |
| weights.append(ds_cfg.weight) | |
| # get_logger().info(f"Dataset '{ds_cfg.name}' loaded.") | |
| # Normalize the weights and mix the datasets. | |
| weights = to_numpy(weights) | |
| weights = weights / weights.sum() | |
| self.train_datasets = datasets | |
| self.train_dataset = MixedWebDataset() | |
| self.train_dataset.append(wds.RandomMix(datasets, weights)) | |
| self.train_dataset = self.train_dataset.with_epoch(50_000).shuffle(1000, initial=1000) | |
| def _setup_mocap(self): | |
| self.mocap_dataset = MoCapDataset(**self.cfg_mocap.cfg) | |
| def _setup_eval(self, selected_ds_names:Optional[List[str]]=None): | |
| from lib.data.datasets.skel_hmr2_fashion.image_dataset import ImageDataset | |
| hack_cfg = { | |
| 'IMAGE_SIZE' : 256, | |
| 'IMAGE_MEAN' : [0.485, 0.456, 0.406], | |
| 'IMAGE_STD' : [0.229, 0.224, 0.225], | |
| 'BBOX_SHAPE' : [192, 256], | |
| 'augm' : self.cfg.image_augmentation, | |
| 'SUPPRESS_KP_CONF_THRESH' : 0.3, | |
| 'FILTER_NUM_KP' : 4, | |
| 'FILTER_NUM_KP_THRESH' : 0.0, | |
| 'FILTER_REPROJ_THRESH' : 31000, | |
| 'SUPPRESS_BETAS_THRESH' : 3.0, | |
| 'SUPPRESS_BAD_POSES' : False, | |
| 'POSES_BETAS_SIMULTANEOUS': True, | |
| 'FILTER_NO_POSES' : False, | |
| 'BETAS_REG' : True, | |
| } | |
| self.eval_datasets = {} | |
| for dataset_cfg in self.cfg_eval.datasets: | |
| if selected_ds_names is not None and dataset_cfg.name not in selected_ds_names: | |
| continue | |
| dataset = ImageDataset( | |
| cfg = hack_cfg, | |
| dataset_file = dataset_cfg.item.dataset_file, | |
| img_dir = dataset_cfg.item.img_root, | |
| train = False, | |
| ) | |
| dataset._kp_list_ = dataset_cfg.item.kp_list | |
| self.eval_datasets[dataset_cfg.name] = dataset | |