Spaces:
Starting
on
L4
Starting
on
L4
File size: 1,200 Bytes
5ac1897 |
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 |
from lib.kits.basic import *
class MoCapDataset:
def __init__(self, dataset_file:str, pve_threshold:Optional[float]=None):
'''
Dataset class used for loading a dataset of unpaired SMPL parameter annotations
### Args
- dataset_file: str
- Path to the dataset npz file.
- pve_threshold: float
- Threshold for PVE quality filtering.
'''
data = np.load(dataset_file)
if pve_threshold is not None:
pve = data['pve_max']
mask = pve < pve_threshold
else:
mask = np.ones(len(data['poses']), dtype=np.bool)
self.poses = data['poses'].astype(np.float32)[mask, 3:]
self.betas = data['betas'].astype(np.float32)[mask]
self.length = len(self.poses)
get_logger().info(f'Loaded {self.length} items among {len(pve)} samples filtered from {dataset_file} (using threshold = {pve_threshold})')
def __getitem__(self, idx: int) -> Dict:
poses = self.poses[idx].copy()
betas = self.betas[idx].copy()
item = {'poses_body': poses, 'betas': betas}
return item
def __len__(self) -> int:
return self.length
|