Prepares predictions and ground truth pose to compute metrics. Only keeps ground truth and predicted assemblies with at least 2 valid keypoints. Sets the coordinates for all keypoints that aren't visible (for ground truth, visibility <= 0 and for predictions score <= 0) to ``np.nan``.
(
ground_truth: dict[str, np.ndarray],
predictions: dict[str, np.ndarray],
)
| 137 | |
| 138 | |
| 139 | def prepare_evaluation_data( |
| 140 | ground_truth: dict[str, np.ndarray], |
| 141 | predictions: dict[str, np.ndarray], |
| 142 | ) -> list[tuple[np.ndarray, np.ndarray]]: |
| 143 | """Prepares predictions and ground truth pose to compute metrics. |
| 144 | |
| 145 | Only keeps ground truth and predicted assemblies with at least 2 valid keypoints. |
| 146 | Sets the coordinates for all keypoints that aren't visible (for ground truth, |
| 147 | visibility <= 0 and for predictions score <= 0) to ``np.nan``. |
| 148 | |
| 149 | Sorts valid predictions by score. |
| 150 | |
| 151 | Args: |
| 152 | ground_truth: For each image, the GT of shape (n_idv, n_bpt, 3). |
| 153 | predictions: For each image, the pose predictions of shape (n_pred, n_bpt, 3). |
| 154 | |
| 155 | Returns: |
| 156 | A list containing (ground truth pose, predicted pose) for each image in the |
| 157 | dataset, where the predicted pose is sorted from highest to lowest score. |
| 158 | """ |
| 159 | pose_data = [] |
| 160 | for image, gt in ground_truth.items(): |
| 161 | gt = gt.copy() |
| 162 | gt[gt[..., 2] <= 0] = np.nan |
| 163 | |
| 164 | # only keep ground truth pose with at least one keypoint |
| 165 | gt_mask = np.any(np.all(~np.isnan(gt), axis=-1), axis=-1) |
| 166 | gt = gt[gt_mask] |
| 167 | |
| 168 | pred = predictions[image][..., :3].copy() # PAF have 5 values; keep xy + score |
| 169 | pred[pred[..., 2] < 0] = np.nan |
| 170 | |
| 171 | # only keep predicted pose with at least two keypoints |
| 172 | pred_mask = np.any(np.all(~np.isnan(pred), axis=-1), axis=-1) |
| 173 | pred = pred[pred_mask] |
| 174 | |
| 175 | scores = np.nanmean(pred[:, :, 2], axis=-1) |
| 176 | pred_order = np.argsort(-scores, kind="mergesort") |
| 177 | pose_data.append((gt, pred[pred_order])) |
| 178 | |
| 179 | return pose_data |