(self, seq_i, start_3d, end_3d, flip, reverse)
| 109 | return self.state |
| 110 | |
| 111 | def get_batch(self, seq_i, start_3d, end_3d, flip, reverse): |
| 112 | subject,action,cam_index = seq_i |
| 113 | seq_name = (subject,action,int(cam_index)) |
| 114 | |
| 115 | start_2d = start_3d - self.pad - self.causal_shift |
| 116 | end_2d = end_3d + self.pad - self.causal_shift |
| 117 | |
| 118 | seq_2d = self.poses_2d[seq_name].copy() |
| 119 | seq_2d_GT = self.poses_2d_GT[seq_name].copy() |
| 120 | |
| 121 | low_2d = max(start_2d, 0) |
| 122 | high_2d = min(end_2d, seq_2d.shape[0]) |
| 123 | pad_left_2d = low_2d - start_2d |
| 124 | pad_right_2d = end_2d - high_2d |
| 125 | |
| 126 | if pad_left_2d != 0 or pad_right_2d != 0: |
| 127 | self.batch_2d = np.pad(seq_2d[low_2d:high_2d], ((pad_left_2d, pad_right_2d), (0, 0), (0, 0)), 'edge') |
| 128 | self.batch_2d_GT = np.pad(seq_2d_GT[low_2d:high_2d], ((pad_left_2d, pad_right_2d), (0, 0), (0, 0)), 'edge') |
| 129 | |
| 130 | else: |
| 131 | self.batch_2d = seq_2d[low_2d:high_2d] |
| 132 | self.batch_2d_GT = seq_2d_GT[low_2d:high_2d] |
| 133 | |
| 134 | if flip: |
| 135 | self.batch_2d[ :, :, 0] *= -1 |
| 136 | self.batch_2d_GT[ :, :, 0] *= -1 |
| 137 | |
| 138 | self.batch_2d[ :, self.kps_left + self.kps_right] = self.batch_2d[ :, |
| 139 | self.kps_right + self.kps_left] |
| 140 | self.batch_2d_GT[ :, self.kps_left + self.kps_right] = self.batch_2d_GT[ :, |
| 141 | self.kps_right + self.kps_left] |
| 142 | if reverse: |
| 143 | self.batch_2d = self.batch_2d[::-1].copy() |
| 144 | self.batch_2d_GT = self.batch_2d_GT[::-1].copy() |
| 145 | |
| 146 | if self.poses_3d is not None: |
| 147 | seq_3d = self.poses_3d[seq_name].copy() |
| 148 | if self.out_all: |
| 149 | low_3d = low_2d |
| 150 | high_3d = high_2d |
| 151 | pad_left_3d = pad_left_2d |
| 152 | pad_right_3d = pad_right_2d |
| 153 | else: |
| 154 | low_3d = max(start_3d, 0) |
| 155 | high_3d = min(end_3d, seq_3d.shape[0]) |
| 156 | pad_left_3d = low_3d - start_3d |
| 157 | pad_right_3d = end_3d - high_3d |
| 158 | |
| 159 | if pad_left_3d != 0 or pad_right_3d != 0: |
| 160 | self.batch_3d = np.pad(seq_3d[low_3d:high_3d], ((pad_left_3d, pad_right_3d), (0, 0), (0, 0)), 'edge') |
| 161 | else: |
| 162 | self.batch_3d = seq_3d[low_3d:high_3d] |
| 163 | |
| 164 | if flip: |
| 165 | self.batch_3d[ :, :, 0] *= -1 |
| 166 | self.batch_3d[ :, self.joints_left + self.joints_right] = \ |
| 167 | self.batch_3d[ :, self.joints_right + self.joints_left] |
| 168 | if reverse: |
no outgoing calls
no test coverage detected