Add comprehensive tests for Assembler and Assembly classes

Introduces detailed unit tests for the Assembler, Assembly, Joint, and Link classes in dlclive.core.inferenceutils. The new tests cover metadata parsing, detection flattening, link extraction, assembly building, Mahalanobis distance calculation, I/O helpers, and various Assembly operations, improving test coverage and reliability.

Author: C-Achard

dlclive/core/inferenceutils.py

@@ -8,6 +8,9 @@
 #
 # Licensed under GNU Lesser General Public License v3.0
 #
+
+
+# NOTE DUPLICATED from deeplabcut/core/inferenceutils.py
 from __future__ import annotations
 
 import heapq
@@ -17,9 +20,10 @@
 import pickle
 import warnings
 from collections import defaultdict
+from collections.abc import Iterable
 from dataclasses import dataclass
 from math import erf, sqrt
-from typing import Any, Iterable, Tuple
+from typing import Any
 
 import networkx as nx
 import numpy as np
@@ -41,7 +45,7 @@ def _conv_square_to_condensed_indices(ind_row, ind_col, n):
     return n * ind_col - ind_col * (ind_col + 1) // 2 + ind_row - 1 - ind_col
 
 
-Position = Tuple[float, float]
+Position = tuple[float, float]
 
 
 @dataclass(frozen=True)
@@ -61,9 +65,7 @@ def __init__(self, j1, j2, affinity=1):
         self._length = sqrt((j1.pos[0] - j2.pos[0]) ** 2 + (j1.pos[1] - j2.pos[1]) ** 2)
 
     def __repr__(self):
-        return (
-            f"Link {self.idx}, affinity={self.affinity:.2f}, length={self.length:.2f}"
-        )
+        return f"Link {self.idx}, affinity={self.affinity:.2f}, length={self.length:.2f}"
 
     @property
     def confidence(self):
@@ -155,7 +157,7 @@ def soft_identity(self):
         unq, idx, cnt = np.unique(data[:, 3], return_inverse=True, return_counts=True)
         avg = np.bincount(idx, weights=data[:, 2]) / cnt
         soft = softmax(avg)
-        return dict(zip(unq.astype(int), soft))
+        return dict(zip(unq.astype(int), soft, strict=False))
 
     @property
     def affinity(self):
@@ -261,9 +263,7 @@ def __init__(
         self.max_overlap = max_overlap
         self._has_identity = "identity" in self[0]
         if identity_only and not self._has_identity:
-            warnings.warn(
-                "The network was not trained with identity; setting `identity_only` to False."
-            )
+            warnings.warn("The network was not trained with identity; setting `identity_only` to False.", stacklevel=2)
         self.identity_only = identity_only & self._has_identity
         self.nan_policy = nan_policy
         self.force_fusion = force_fusion
@@ -344,15 +344,15 @@ def calibrate(self, train_data_file):
             pass
         n_bpts = len(df.columns.get_level_values("bodyparts").unique())
         if n_bpts == 1:
-            warnings.warn("There is only one keypoint; skipping calibration...")
+            warnings.warn("There is only one keypoint; skipping calibration...", stacklevel=2)
             return
 
         xy = df.to_numpy().reshape((-1, n_bpts, 2))
         frac_valid = np.mean(~np.isnan(xy), axis=(1, 2))
         # Only keeps skeletons that are more than 90% complete
         xy = xy[frac_valid >= 0.9]
         if not xy.size:
-            warnings.warn("No complete poses were found. Skipping calibration...")
+            warnings.warn("No complete poses were found. Skipping calibration...", stacklevel=2)
             return
 
         # TODO Normalize dists by longest length?
@@ -368,13 +368,9 @@ def calibrate(self, train_data_file):
             self.safe_edge = True
         except np.linalg.LinAlgError:
             # Covariance matrix estimation fails due to numerical singularities
-            warnings.warn(
-                "The assembler could not be robustly calibrated. Continuing without it..."
-            )
+            warnings.warn("The assembler could not be robustly calibrated. Continuing without it...", stacklevel=2)
 
-    def calc_assembly_mahalanobis_dist(
-        self, assembly, return_proba=False, nan_policy="little"
-    ):
+    def calc_assembly_mahalanobis_dist(self, assembly, return_proba=False, nan_policy="little"):
         if self._kde is None:
             raise ValueError("Assembler should be calibrated first with training data.")
 
@@ -428,10 +424,10 @@ def _flatten_detections(data_dict):
             ids = [np.ones(len(arr), dtype=int) * -1 for arr in confidence]
         else:
             ids = [arr.argmax(axis=1) for arr in ids]
-        for i, (coords, conf, id_) in enumerate(zip(coordinates, confidence, ids)):
+        for i, (coords, conf, id_) in enumerate(zip(coordinates, confidence, ids, strict=False)):
             if not np.any(coords):
                 continue
-            for xy, p, g in zip(coords, conf, id_):
+            for xy, p, g in zip(coords, conf, id_, strict=False):
                 joint = Joint(tuple(xy), p.item(), i, ind, g)
                 ind += 1
                 yield joint
@@ -453,9 +449,7 @@ def extract_best_links(self, joints_dict, costs, trees=None):
             aff[np.isnan(aff)] = 0
 
             if trees:
-                vecs = np.vstack(
-                    [[*det_s.pos, *det_t.pos] for det_s in dets_s for det_t in dets_t]
-                )
+                vecs = np.vstack([[*det_s.pos, *det_t.pos] for det_s in dets_s for det_t in dets_t])
                 dists = []
                 for n, tree in enumerate(trees, start=1):
                     d, _ = tree.query(vecs)
@@ -464,45 +458,34 @@ def extract_best_links(self, joints_dict, costs, trees=None):
                 aff *= w.reshape(aff.shape)
 
             if self.greedy:
-                conf = np.asarray(
-                    [
-                        [det_s.confidence * det_t.confidence for det_t in dets_t]
-                        for det_s in dets_s
-                    ]
-                )
-                rows, cols = np.where(
-                    (conf >= self.pcutoff * self.pcutoff) & (aff >= self.min_affinity)
-                )
+                conf = np.asarray([[det_s.confidence * det_t.confidence for det_t in dets_t] for det_s in dets_s])
+                rows, cols = np.where((conf >= self.pcutoff * self.pcutoff) & (aff >= self.min_affinity))
                 candidates = sorted(
-                    zip(rows, cols, aff[rows, cols], lengths[rows, cols]),
+                    zip(rows, cols, aff[rows, cols], lengths[rows, cols], strict=False),
                     key=lambda x: x[2],
                     reverse=True,
                 )
                 i_seen = set()
                 j_seen = set()
-                for i, j, w, l in candidates:
+                for i, j, w, _l in candidates:
                     if i not in i_seen and j not in j_seen:
                         i_seen.add(i)
                         j_seen.add(j)
                         links.append(Link(dets_s[i], dets_t[j], w))
                         if len(i_seen) == self.max_n_individuals:
                             break
             else:  # Optimal keypoint pairing
-                inds_s = sorted(
-                    range(len(dets_s)), key=lambda x: dets_s[x].confidence, reverse=True
-                )[: self.max_n_individuals]
-                inds_t = sorted(
-                    range(len(dets_t)), key=lambda x: dets_t[x].confidence, reverse=True
-                )[: self.max_n_individuals]
-                keep_s = [
-                    ind for ind in inds_s if dets_s[ind].confidence >= self.pcutoff
+                inds_s = sorted(range(len(dets_s)), key=lambda x: dets_s[x].confidence, reverse=True)[
+                    : self.max_n_individuals
                 ]
-                keep_t = [
-                    ind for ind in inds_t if dets_t[ind].confidence >= self.pcutoff
+                inds_t = sorted(range(len(dets_t)), key=lambda x: dets_t[x].confidence, reverse=True)[
+                    : self.max_n_individuals
                 ]
+                keep_s = [ind for ind in inds_s if dets_s[ind].confidence >= self.pcutoff]
+                keep_t = [ind for ind in inds_t if dets_t[ind].confidence >= self.pcutoff]
                 aff = aff[np.ix_(keep_s, keep_t)]
                 rows, cols = linear_sum_assignment(aff, maximize=True)
-                for row, col in zip(rows, cols):
+                for row, col in zip(rows, cols, strict=False):
                     w = aff[row, col]
                     if w >= self.min_affinity:
                         links.append(Link(dets_s[keep_s[row]], dets_t[keep_t[col]], w))
@@ -538,19 +521,17 @@ def push_to_stack(i):
             if new_ind in assembled:
                 continue
             if safe_edge:
-                d_old = self.calc_assembly_mahalanobis_dist(
-                    assembly, nan_policy=nan_policy
-                )
+                d_old = self.calc_assembly_mahalanobis_dist(assembly, nan_policy=nan_policy)
                 success = assembly.add_link(best, store_dict=True)
                 if not success:
                     assembly._dict = dict()
                     continue
                 d = self.calc_assembly_mahalanobis_dist(assembly, nan_policy=nan_policy)
                 if d < d_old:
                     push_to_stack(new_ind)
-                    if tabu:  
-                        _, _, link = heapq.heappop(tabu)  
-                        heapq.heappush(stack, (-link.affinity, next(counter), link)) 
+                    if tabu:
+                        _, _, link = heapq.heappop(tabu)
+                        heapq.heappush(stack, (-link.affinity, next(counter), link))
                 else:
                     heapq.heappush(tabu, (d - d_old, next(counter), best))
                     assembly.__dict__.update(assembly._dict)
@@ -593,9 +574,7 @@ def build_assemblies(self, links):
                 continue
             assembly = Assembly(self.n_multibodyparts)
             assembly.add_link(link)
-            self._fill_assembly(
-                assembly, lookup, assembled, self.safe_edge, self.nan_policy
-            )
+            self._fill_assembly(assembly, lookup, assembled, self.safe_edge, self.nan_policy)
             for assembly_link in assembly._links:
                 i, j = assembly_link.idx
                 lookup[i].pop(j)
@@ -607,10 +586,7 @@ def build_assemblies(self, links):
         n_extra = len(assemblies) - self.max_n_individuals
         if n_extra > 0:
             if self.safe_edge:
-                ds_old = [
-                    self.calc_assembly_mahalanobis_dist(assembly)
-                    for assembly in assemblies
-                ]
+                ds_old = [self.calc_assembly_mahalanobis_dist(assembly) for assembly in assemblies]
                 while len(assemblies) > self.max_n_individuals:
                     ds = []
                     for i, j in itertools.combinations(range(len(assemblies)), 2):
@@ -665,7 +641,7 @@ def build_assemblies(self, links):
                             for idx in store[j]._idx:
                                 store[idx] = store[i]
                     except KeyError:
-                        # Some links may reference indices that were never added to `store`;  
+                        # Some links may reference indices that were never added to `store`;
                         # in that case we intentionally skip merging for this link
                         pass
 
@@ -742,10 +718,7 @@ def _assemble(self, data_dict, ind_frame):
             for _, group in groups:
                 ass = Assembly(self.n_multibodyparts)
                 for joint in sorted(group, key=lambda x: x.confidence, reverse=True):
-                    if (
-                        joint.confidence >= self.pcutoff
-                        and joint.label < self.n_multibodyparts
-                    ):
+                    if joint.confidence >= self.pcutoff and joint.label < self.n_multibodyparts:
                         ass.add_joint(joint)
                 if len(ass):
                     assemblies.append(ass)
@@ -774,24 +747,18 @@ def _assemble(self, data_dict, ind_frame):
             assembled.update(assembled_)
 
         # Remove invalid assemblies
-        discarded = set(
-            joint
-            for joint in joints
-            if joint.idx not in assembled and np.isfinite(joint.confidence)
-        )
+        discarded = set(joint for joint in joints if joint.idx not in assembled and np.isfinite(joint.confidence))
         for assembly in assemblies[::-1]:
             if 0 < assembly.n_links < self.min_n_links or not len(assembly):
                 for link in assembly._links:
                     discarded.update((link.j1, link.j2))
                 assemblies.remove(assembly)
         if 0 < self.max_overlap < 1:  # Non-maximum pose suppression
             if self._kde is not None:
-                scores = [
-                    -self.calc_assembly_mahalanobis_dist(ass) for ass in assemblies
-                ]
+                scores = [-self.calc_assembly_mahalanobis_dist(ass) for ass in assemblies]
             else:
                 scores = [ass._affinity for ass in assemblies]
-            lst = list(zip(scores, assemblies))
+            lst = list(zip(scores, assemblies, strict=False))
             assemblies = []
             while lst:
                 temp = max(lst, key=lambda x: x[0])
@@ -857,9 +824,7 @@ def wrapped(i):
             n_frames = len(self.metadata["imnames"])
             with multiprocessing.Pool(n_processes) as p:
                 with tqdm(total=n_frames) as pbar:
-                    for i, (assemblies, unique) in p.imap_unordered(
-                        wrapped, range(n_frames), chunksize=chunk_size
-                    ):
+                    for i, (assemblies, unique) in p.imap_unordered(wrapped, range(n_frames), chunksize=chunk_size):
                         if assemblies:
                             self.assemblies[i] = assemblies
                         if unique is not None:
@@ -878,9 +843,7 @@ def parse_metadata(data):
         params["joint_names"] = data["metadata"]["all_joints_names"]
         params["num_joints"] = len(params["joint_names"])
         params["paf_graph"] = data["metadata"]["PAFgraph"]
-        params["paf"] = data["metadata"].get(
-            "PAFinds", np.arange(len(params["joint_names"]))
-        )
+        params["paf"] = data["metadata"].get("PAFinds", np.arange(len(params["joint_names"])))
         params["bpts"] = params["ibpts"] = range(params["num_joints"])
         params["imnames"] = [fn for fn in list(data) if fn != "metadata"]
         return params
@@ -970,11 +933,7 @@ def calc_object_keypoint_similarity(
     else:
         oks = []
         xy_preds = [xy_pred]
-        combos = (
-            pair
-            for l in range(len(symmetric_kpts))
-            for pair in itertools.combinations(symmetric_kpts, l + 1)
-        )
+        combos = (pair for l in range(len(symmetric_kpts)) for pair in itertools.combinations(symmetric_kpts, l + 1))
         for pairs in combos:
             # Swap corresponding keypoints
             tmp = xy_pred.copy()
@@ -1011,9 +970,7 @@ def match_assemblies(
     num_ground_truth = len(ground_truth)
 
     # Sort predictions by score
-    inds_pred = np.argsort(
-        [ins.affinity if ins.n_links else ins.confidence for ins in predictions]
-    )[::-1]
+    inds_pred = np.argsort([ins.affinity if ins.n_links else ins.confidence for ins in predictions])[::-1]
     predictions = np.asarray(predictions)[inds_pred]
 
     # indices of unmatched ground truth assemblies
@@ -1074,7 +1031,7 @@ def match_assemblies(
                 if ~np.isnan(oks):
                     mat[i, j] = oks
         rows, cols = linear_sum_assignment(mat, maximize=True)
-        for row, col in zip(rows, cols):
+        for row, col in zip(rows, cols, strict=False):
             matched[row].ground_truth = ground_truth[col]
             matched[row].oks = mat[row, col]
             _ = inds_true.remove(col)
@@ -1087,7 +1044,7 @@ def parse_ground_truth_data_file(h5_file):
     try:
         df.drop("single", axis=1, level="individuals", inplace=True)
     except KeyError:
-        # Ignore if the "single" individual column is absent 
+        # Ignore if the "single" individual column is absent
         pass
     # Cast columns of dtype 'object' to float to avoid TypeError
     # further down in _parse_ground_truth_data.
@@ -1120,15 +1077,13 @@ def find_outlier_assemblies(dict_of_assemblies, criterion="area", qs=(5, 95)):
         raise ValueError(f"Invalid criterion {criterion}.")
 
     if len(qs) != 2:
-        raise ValueError(
-            "Two percentiles (for lower and upper bounds) should be given."
-        )
+        raise ValueError("Two percentiles (for lower and upper bounds) should be given.")
 
     tuples = []
     for frame_ind, assemblies in dict_of_assemblies.items():
         for assembly in assemblies:
             tuples.append((frame_ind, getattr(assembly, criterion)))
-    frame_inds, vals = zip(*tuples)
+    frame_inds, vals = zip(*tuples, strict=False)
     vals = np.asarray(vals)
     lo, up = np.percentile(vals, qs, interpolation="nearest")
     inds = np.flatnonzero((vals < lo) | (vals > up)).tolist()
@@ -1226,9 +1181,7 @@ def evaluate_assembly_greedy(
         oks = np.asarray([match.oks for match in all_matched])[sorted_pred_indices]
 
         # Compute prediction and recall
-        p, r = _compute_precision_and_recall(
-            total_gt_assemblies, oks, oks_t, recall_thresholds
-        )
+        p, r = _compute_precision_and_recall(total_gt_assemblies, oks, oks_t, recall_thresholds)
         precisions.append(p)
         recalls.append(r)
 
@@ -1246,12 +1199,14 @@ def evaluate_assembly(
     ass_pred_dict,
     ass_true_dict,
     oks_sigma=0.072,
-    oks_thresholds=np.linspace(0.5, 0.95, 10),
+    oks_thresholds=None,
     margin=0,
     symmetric_kpts=None,
     greedy_matching=False,
     with_tqdm: bool = True,
 ):
+    if oks_thresholds is None:
+        oks_thresholds = np.linspace(0.5, 0.95, 10)
     if greedy_matching:
         return evaluate_assembly_greedy(
             ass_true_dict,
@@ -1299,9 +1254,7 @@ def evaluate_assembly(
     precisions = []
     recalls = []
     for t in oks_thresholds:
-        p, r = _compute_precision_and_recall(
-            total_gt_assemblies, oks, t, recall_thresholds
-        )
+        p, r = _compute_precision_and_recall(total_gt_assemblies, oks, t, recall_thresholds)
         precisions.append(p)
         recalls.append(r)