Spaces:
Sleeping
Sleeping
| from __future__ import absolute_import | |
| from __future__ import print_function | |
| from __future__ import division | |
| import torch | |
| import numpy as np | |
| from skimage.util.shape import view_as_windows | |
| from configs import constants as _C | |
| from .utils.normalizer import Normalizer | |
| from ..utils.imutils import transform | |
| class BaseDataset(torch.utils.data.Dataset): | |
| def __init__(self, cfg, training=True): | |
| super(BaseDataset, self).__init__() | |
| self.epoch = 0 | |
| self.training = training | |
| self.n_joints = _C.KEYPOINTS.NUM_JOINTS | |
| self.n_frames = cfg.DATASET.SEQLEN + 1 | |
| self.keypoints_normalizer = Normalizer(cfg) | |
| def prepare_video_batch(self): | |
| r = self.epoch % 4 | |
| self.video_indices = [] | |
| vid_name = self.labels['vid'] | |
| if isinstance(vid_name, torch.Tensor): vid_name = vid_name.numpy() | |
| video_names_unique, group = np.unique( | |
| vid_name, return_index=True) | |
| perm = np.argsort(group) | |
| group_perm = group[perm] | |
| indices = np.split( | |
| np.arange(0, self.labels['vid'].shape[0]), group_perm[1:] | |
| ) | |
| for idx in range(len(video_names_unique)): | |
| indexes = indices[idx] | |
| if indexes.shape[0] < self.n_frames: continue | |
| chunks = view_as_windows( | |
| indexes, (self.n_frames), step=self.n_frames // 4 | |
| ) | |
| start_finish = chunks[r::4, (0, -1)].tolist() | |
| self.video_indices += start_finish | |
| self.epoch += 1 | |
| def __len__(self): | |
| if self.training: | |
| return len(self.video_indices) | |
| else: | |
| return len(self.labels['kp2d']) | |
| def __getitem__(self, index): | |
| return self.get_single_sequence(index) | |
| def get_single_sequence(self, index): | |
| NotImplementedError('get_single_sequence is not implemented') | |
| def get_naive_intrinsics(self, res): | |
| # Assume 45 degree FOV | |
| img_w, img_h = res | |
| self.focal_length = (img_w * img_w + img_h * img_h) ** 0.5 | |
| self.cam_intrinsics = torch.eye(3).repeat(1, 1, 1).float() | |
| self.cam_intrinsics[:, 0, 0] = self.focal_length | |
| self.cam_intrinsics[:, 1, 1] = self.focal_length | |
| self.cam_intrinsics[:, 0, 2] = img_w/2. | |
| self.cam_intrinsics[:, 1, 2] = img_h/2. | |
| def j2d_processing(self, kp, bbox): | |
| center = bbox[..., :2] | |
| scale = bbox[..., -1:] | |
| nparts = kp.shape[0] | |
| for i in range(nparts): | |
| kp[i, 0:2] = transform(kp[i, 0:2] + 1, center, scale, | |
| [224, 224]) | |
| kp[:, :2] = 2. * kp[:, :2] / 224 - 1. | |
| kp = kp.astype('float32') | |
| return kp |