Matches assemblies to ground truth predictions. Returns: int: the total number of valid ground truth assemblies list[MatchedPrediction]: a list containing all valid predictions, potentially matched to ground truth assemblies.
(
predictions: list[Assembly],
ground_truth: list[Assembly],
sigma: float,
margin: int = 0,
symmetric_kpts: list[tuple[int, int]] | None = None,
greedy_matching: bool = False,
greedy_oks_threshold: float = 0.0,
)
| 948 | |
| 949 | |
| 950 | def match_assemblies( |
| 951 | predictions: list[Assembly], |
| 952 | ground_truth: list[Assembly], |
| 953 | sigma: float, |
| 954 | margin: int = 0, |
| 955 | symmetric_kpts: list[tuple[int, int]] | None = None, |
| 956 | greedy_matching: bool = False, |
| 957 | greedy_oks_threshold: float = 0.0, |
| 958 | ) -> tuple[int, list[MatchedPrediction]]: |
| 959 | """Matches assemblies to ground truth predictions. |
| 960 | |
| 961 | Returns: |
| 962 | int: the total number of valid ground truth assemblies |
| 963 | list[MatchedPrediction]: a list containing all valid predictions, potentially |
| 964 | matched to ground truth assemblies. |
| 965 | """ |
| 966 | # Only consider assemblies of at least two keypoints |
| 967 | predictions = [a for a in predictions if len(a) > 1] |
| 968 | ground_truth = [a for a in ground_truth if len(a) > 1] |
| 969 | num_ground_truth = len(ground_truth) |
| 970 | |
| 971 | # Sort predictions by score |
| 972 | inds_pred = np.argsort([ins.affinity if ins.n_links else ins.confidence for ins in predictions])[::-1] |
| 973 | predictions = np.asarray(predictions)[inds_pred] |
| 974 | |
| 975 | # indices of unmatched ground truth assemblies |
| 976 | matched = [ |
| 977 | MatchedPrediction( |
| 978 | prediction=p, |
| 979 | score=(p.affinity if p.n_links else p.confidence), |
| 980 | ground_truth=None, |
| 981 | oks=0.0, |
| 982 | ) |
| 983 | for p in predictions |
| 984 | ] |
| 985 | |
| 986 | # Greedy assembly matching like in pycocotools |
| 987 | if greedy_matching: |
| 988 | matched_gt_indices = set() |
| 989 | for idx, pred in enumerate(predictions): |
| 990 | oks = [ |
| 991 | calc_object_keypoint_similarity( |
| 992 | pred.xy, |
| 993 | gt.xy, |
| 994 | sigma, |
| 995 | margin, |
| 996 | symmetric_kpts, |
| 997 | ) |
| 998 | for gt in ground_truth |
| 999 | ] |
| 1000 | if np.all(np.isnan(oks)): |
| 1001 | continue |
| 1002 | |
| 1003 | ind_best = np.nanargmax(oks) |
| 1004 | |
| 1005 | # if this gt already matched, and not a crowd, continue |
| 1006 | if ind_best in matched_gt_indices: |
| 1007 | continue |
no test coverage detected