Soft OKS NMS implementations. Args: kpts_db thr: retain oks overlap < thr. max_dets: max number of detections to keep. sigmas: Keypoint labelling uncertainty. Returns: np.ndarray: indexes to keep.
(kpts_db, thr, max_dets=20, sigmas=None, vis_thr=None)
| 148 | |
| 149 | |
| 150 | def soft_oks_nms(kpts_db, thr, max_dets=20, sigmas=None, vis_thr=None): |
| 151 | """Soft OKS NMS implementations. |
| 152 | |
| 153 | Args: |
| 154 | kpts_db |
| 155 | thr: retain oks overlap < thr. |
| 156 | max_dets: max number of detections to keep. |
| 157 | sigmas: Keypoint labelling uncertainty. |
| 158 | |
| 159 | Returns: |
| 160 | np.ndarray: indexes to keep. |
| 161 | """ |
| 162 | if len(kpts_db) == 0: |
| 163 | return [] |
| 164 | |
| 165 | scores = np.array([k['score'] for k in kpts_db]) |
| 166 | kpts = np.array([k['keypoints'].flatten() for k in kpts_db]) |
| 167 | areas = np.array([k['area'] for k in kpts_db]) |
| 168 | |
| 169 | order = scores.argsort()[::-1] |
| 170 | scores = scores[order] |
| 171 | |
| 172 | keep = np.zeros(max_dets, dtype=np.intp) |
| 173 | keep_cnt = 0 |
| 174 | while len(order) > 0 and keep_cnt < max_dets: |
| 175 | i = order[0] |
| 176 | |
| 177 | oks_ovr = oks_iou(kpts[i], kpts[order[1:]], areas[i], areas[order[1:]], |
| 178 | sigmas, vis_thr) |
| 179 | |
| 180 | order = order[1:] |
| 181 | scores = _rescore(oks_ovr, scores[1:], thr) |
| 182 | |
| 183 | tmp = scores.argsort()[::-1] |
| 184 | order = order[tmp] |
| 185 | scores = scores[tmp] |
| 186 | |
| 187 | keep[keep_cnt] = i |
| 188 | keep_cnt += 1 |
| 189 | |
| 190 | keep = keep[:keep_cnt] |
| 191 | |
| 192 | return keep |