|
|
import numpy as np |
|
|
import torch |
|
|
from torch.utils.data.dataset import Dataset |
|
|
import os |
|
|
from configs.paths import dataset_root |
|
|
import copy |
|
|
from tqdm import tqdm |
|
|
from .base import BASE |
|
|
|
|
|
class BEDLAM(BASE): |
|
|
def __init__(self, split='train_6fps',**kwargs): |
|
|
super(BEDLAM, self).__init__(**kwargs) |
|
|
assert split in ['train_1fps','train_3fps','train_6fps','validation_6fps'] |
|
|
assert not self.kid_offset |
|
|
|
|
|
self.ds_name = 'bedlam' |
|
|
self.dataset_path = os.path.join(dataset_root,'bedlam') |
|
|
annots_path = os.path.join(self.dataset_path,f'bedlam_smpl_{split}.npz') |
|
|
self.annots = np.load(annots_path, allow_pickle=True)['annots'][()] |
|
|
self.img_names = list(self.annots.keys()) |
|
|
self.split = 'train' if 'train' in split else 'validation' |
|
|
|
|
|
def __len__(self): |
|
|
return len(self.img_names) |
|
|
|
|
|
def cnt_instances(self): |
|
|
ins_cnt = 0 |
|
|
for idx in tqdm(range(len(self))): |
|
|
img_id = idx |
|
|
img_name = self.img_names[img_id] |
|
|
|
|
|
ins_cnt += len(self.annots[img_name]['shape']) |
|
|
|
|
|
|
|
|
print(f'TOTAL: {ins_cnt}') |
|
|
|
|
|
def get_raw_data(self, idx): |
|
|
|
|
|
img_id = idx%len(self.img_names) |
|
|
img_name = self.img_names[img_id] |
|
|
|
|
|
annots = copy.deepcopy(self.annots[img_name]) |
|
|
img_path = os.path.join(self.dataset_path,self.split,img_name) |
|
|
|
|
|
cam_intrinsics = torch.from_numpy(annots['cam_int']).unsqueeze(0) |
|
|
cam_rot = torch.from_numpy(np.stack(annots['cam_rot'])) |
|
|
cam_trans = torch.from_numpy(np.stack(annots['cam_trans'])) |
|
|
|
|
|
betas = torch.from_numpy(np.stack(annots['shape'])) |
|
|
poses = torch.from_numpy(np.stack(annots['pose_world'])) |
|
|
transl = torch.from_numpy(np.stack(annots['trans_world'])) |
|
|
|
|
|
raw_data={'img_path': img_path, |
|
|
'ds': 'bedlam', |
|
|
'pnum': len(betas), |
|
|
'betas': betas.float(), |
|
|
'poses': poses.float(), |
|
|
'transl': transl.float(), |
|
|
'cam_rot': cam_rot.float(), |
|
|
'cam_trans': cam_trans.float(), |
|
|
'cam_intrinsics':cam_intrinsics.float(), |
|
|
'3d_valid': True, |
|
|
'age_valid': False, |
|
|
'detect_all_people':True |
|
|
} |
|
|
|
|
|
if self.mode == 'eval': |
|
|
raw_data['occ_level'] = torch.zeros(len(betas),dtype=int) |
|
|
|
|
|
return raw_data |
|
|
|
|
|
|
|
|
|
