| 429 | yield joint |
| 430 | |
| 431 | def extract_best_links(self, joints_dict, costs, trees=None): |
| 432 | links = [] |
| 433 | for ind in self.paf_inds: |
| 434 | s, t = self.graph[ind] |
| 435 | dets_s = joints_dict.get(s, None) |
| 436 | dets_t = joints_dict.get(t, None) |
| 437 | if dets_s is None or dets_t is None: |
| 438 | continue |
| 439 | if ind not in costs: |
| 440 | continue |
| 441 | lengths = costs[ind]["distance"] |
| 442 | if np.isinf(lengths).all(): |
| 443 | continue |
| 444 | aff = costs[ind][self.method].copy() |
| 445 | aff[np.isnan(aff)] = 0 |
| 446 | |
| 447 | if trees: |
| 448 | vecs = np.vstack([[*det_s.pos, *det_t.pos] for det_s in dets_s for det_t in dets_t]) |
| 449 | dists = [] |
| 450 | for n, tree in enumerate(trees, start=1): |
| 451 | d, _ = tree.query(vecs) |
| 452 | dists.append(np.exp(-self._gamma * n * d)) |
| 453 | w = np.mean(dists, axis=0) |
| 454 | aff *= w.reshape(aff.shape) |
| 455 | |
| 456 | if self.greedy: |
| 457 | conf = np.asarray([[det_s.confidence * det_t.confidence for det_t in dets_t] for det_s in dets_s]) |
| 458 | rows, cols = np.where((conf >= self.pcutoff * self.pcutoff) & (aff >= self.min_affinity)) |
| 459 | candidates = sorted( |
| 460 | zip(rows, cols, aff[rows, cols], lengths[rows, cols], strict=False), |
| 461 | key=lambda x: x[2], |
| 462 | reverse=True, |
| 463 | ) |
| 464 | i_seen = set() |
| 465 | j_seen = set() |
| 466 | for i, j, w, _l in candidates: |
| 467 | if i not in i_seen and j not in j_seen: |
| 468 | i_seen.add(i) |
| 469 | j_seen.add(j) |
| 470 | links.append(Link(dets_s[i], dets_t[j], w)) |
| 471 | if len(i_seen) == self.max_n_individuals: |
| 472 | break |
| 473 | else: # Optimal keypoint pairing |
| 474 | inds_s = sorted(range(len(dets_s)), key=lambda x: dets_s[x].confidence, reverse=True)[ |
| 475 | : self.max_n_individuals |
| 476 | ] |
| 477 | inds_t = sorted(range(len(dets_t)), key=lambda x: dets_t[x].confidence, reverse=True)[ |
| 478 | : self.max_n_individuals |
| 479 | ] |
| 480 | keep_s = [ind for ind in inds_s if dets_s[ind].confidence >= self.pcutoff] |
| 481 | keep_t = [ind for ind in inds_t if dets_t[ind].confidence >= self.pcutoff] |
| 482 | aff = aff[np.ix_(keep_s, keep_t)] |
| 483 | rows, cols = linear_sum_assignment(aff, maximize=True) |
| 484 | for row, col in zip(rows, cols, strict=False): |
| 485 | w = aff[row, col] |
| 486 | if w >= self.min_affinity: |
| 487 | links.append(Link(dets_s[keep_s[row]], dets_t[keep_t[col]], w)) |
| 488 | return links |