| 2 | |
| 3 | |
| 4 | class ChunkedGenerator: |
| 5 | def __init__(self, batch_size, cameras, poses_3d, poses_2d, poses_2d_GT, |
| 6 | chunk_length=1, pad=0, causal_shift=0, |
| 7 | shuffle=False, random_seed=1234, |
| 8 | augment=False, reverse_aug= False,kps_left=None, kps_right=None, joints_left=None, joints_right=None, |
| 9 | endless=False, out_all = False): |
| 10 | assert poses_3d is None or len(poses_3d) == len(poses_2d), (len(poses_3d), len(poses_2d)) |
| 11 | assert cameras is None or len(cameras) == len(poses_2d) |
| 12 | |
| 13 | pairs, bounds_1, bounds_2, augment_vectors, reverse_augment_vectors = [], [], [], [], [] |
| 14 | |
| 15 | self.saved_index = {} |
| 16 | start_index = 0 |
| 17 | |
| 18 | for key in poses_2d.keys(): |
| 19 | assert poses_3d is None or poses_3d[key].shape[0] == poses_3d[key].shape[0] |
| 20 | n_chunks = (poses_2d[key].shape[0] + chunk_length - 1) // chunk_length |
| 21 | offset = (n_chunks * chunk_length - poses_2d[key].shape[0]) // 2 |
| 22 | bounds = np.arange(n_chunks + 1) * chunk_length - offset |
| 23 | augment_vector = np.full(len(bounds) - 1, False, dtype=bool) |
| 24 | reverse_augment_vector = np.full(len(bounds) - 1, False, dtype=bool) |
| 25 | keys = np.tile(np.array(key).reshape([1,3]),(len(bounds) - 1, 1)) |
| 26 | |
| 27 | pairs += list(keys) |
| 28 | bounds_1 += list(bounds[:-1]) |
| 29 | bounds_2 += list(bounds[1:]) |
| 30 | augment_vectors += list(augment_vector) |
| 31 | reverse_augment_vectors += list(reverse_augment_vector) |
| 32 | |
| 33 | if reverse_aug: |
| 34 | pairs += list(keys) |
| 35 | bounds_1 += list(bounds[:-1]) |
| 36 | bounds_2 += list(bounds[1:]) |
| 37 | augment_vectors += list(augment_vector) |
| 38 | reverse_augment_vectors += list(~reverse_augment_vector) |
| 39 | |
| 40 | if augment: |
| 41 | if reverse_aug: |
| 42 | pairs += list(keys) |
| 43 | bounds_1 += list(bounds[:-1]) |
| 44 | bounds_2 += list(bounds[1:]) |
| 45 | augment_vectors += list(~augment_vector) |
| 46 | reverse_augment_vectors += list(~reverse_augment_vector) |
| 47 | |
| 48 | else: |
| 49 | pairs += list(keys) |
| 50 | bounds_1 += list(bounds[:-1]) |
| 51 | bounds_2 += list(bounds[1:]) |
| 52 | augment_vectors += list(~augment_vector) |
| 53 | reverse_augment_vectors += list(reverse_augment_vector) |
| 54 | |
| 55 | end_index = start_index + poses_3d[key].shape[0] |
| 56 | self.saved_index[key] = [start_index,end_index] |
| 57 | start_index = start_index + poses_3d[key].shape[0] |
| 58 | |
| 59 | self.pairs = np.array(pairs) |
| 60 | self.bounds_1 = np.array(bounds_1) |
| 61 | self.bounds_2 = np.array(bounds_2) |