| 281 | |
| 282 | |
| 283 | class SkeletonTracker(BaseTracker): |
| 284 | def __init__(self, n_bodyparts): |
| 285 | super().__init__(dim=n_bodyparts * 2, dim_z=n_bodyparts) |
| 286 | self.kf.Q[self.kf.dim_z :, self.kf.dim_z :] *= 10 |
| 287 | self.kf.R[self.kf.dim_z :, self.kf.dim_z :] *= 0.01 |
| 288 | self.kf.P[self.kf.dim_z :, self.kf.dim_z :] *= 1000 |
| 289 | |
| 290 | def update(self, pose): |
| 291 | flat = pose.reshape((-1, 1)) |
| 292 | empty = np.isnan(flat).squeeze() |
| 293 | if empty.any(): |
| 294 | H = self.kf.H.copy() |
| 295 | H[empty] = 0 |
| 296 | flat[empty] = 0 |
| 297 | self.kf.update(flat, H=H) |
| 298 | else: |
| 299 | super().update(flat) |
| 300 | |
| 301 | @BaseTracker.state.setter |
| 302 | def state(self, pose): |
| 303 | curr_pose = pose.copy() |
| 304 | empty = np.isnan(curr_pose).all(axis=1) |
| 305 | if empty.any(): |
| 306 | fill = np.nanmean(pose, axis=0) |
| 307 | curr_pose[empty] = fill |
| 308 | super(SkeletonTracker, type(self)).state.fset(self, curr_pose.reshape((-1, 1))) |
| 309 | |
| 310 | |
| 311 | class BoxTracker(BaseTracker): |