From 4bc3fce0b160748f3cd5e6d43c292878152f0d78 Mon Sep 17 00:00:00 2001 From: crosstyan Date: Sat, 3 May 2025 17:29:50 +0800 Subject: [PATCH] feat: Add minimum affinity filter to affinity result grouping - Introduced a `min_affinity` parameter to the `affinity_result_by_tracking` function, allowing users to specify a threshold for filtering affinity results. - Updated the logic to skip results with affinities below the specified minimum, enhancing the relevance of grouped detections. - Improved function documentation to include details about the new parameter and its purpose. --- playground.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/playground.py b/playground.py index ca6fb4b..e6ef055 100644 --- a/playground.py +++ b/playground.py @@ -1197,13 +1197,22 @@ display(affinities) # %% def affinity_result_by_tracking( results: Iterable[AffinityResult], + min_affinity: float = 0.0, ) -> dict[TrackingID, list[Detection]]: """ Group affinity results by target ID. + + Args: + results: the affinity results to group + min_affinity: the minimum affinity to consider + Returns: + a dictionary mapping tracking IDs to a list of detections """ res: dict[TrackingID, list[Detection]] = defaultdict(list) for affinity_result in results: - for _affinity, t, d in affinity_result.tracking_association(): + for affinity, t, d in affinity_result.tracking_association(): + if affinity < min_affinity: + continue res[t.id].append(d) return res