From 8b325ef7d7dcad1596d591551471f678d93d05b1 Mon Sep 17 00:00:00 2001
From: Niels Poulsen <niels.poulsen@epfl.ch>
Date: Wed, 11 Dec 2024 18:13:20 +0100
Subject: [PATCH 1/4] prune PAF graph implementation for PyTorch models

---
 .../pose_estimation_pytorch/apis/evaluate.py  |   7 +
 .../apis/prune_paf_graph.py                   | 268 ++++++++++++++++++
 .../pose_estimation_pytorch/data/__init__.py  |  11 +
 .../models/predictors/paf_predictor.py        |   9 +
 .../models/predictors/utils.py                | 193 -------------
 5 files changed, 295 insertions(+), 193 deletions(-)
 create mode 100644 deeplabcut/pose_estimation_pytorch/apis/prune_paf_graph.py
 delete mode 100644 deeplabcut/pose_estimation_pytorch/models/predictors/utils.py

diff --git a/deeplabcut/pose_estimation_pytorch/apis/evaluate.py b/deeplabcut/pose_estimation_pytorch/apis/evaluate.py
index 14460ace21..7da1233080 100755
--- a/deeplabcut/pose_estimation_pytorch/apis/evaluate.py
+++ b/deeplabcut/pose_estimation_pytorch/apis/evaluate.py
@@ -21,6 +21,7 @@
 from tqdm import tqdm
 
 import deeplabcut.core.metrics as metrics
+import deeplabcut.pose_estimation_pytorch.apis.prune_paf_graph as prune_paf_graph
 from deeplabcut.core.weight_init import WeightInitialization
 from deeplabcut.pose_estimation_pytorch import utils
 from deeplabcut.pose_estimation_pytorch.apis.utils import (
@@ -436,6 +437,12 @@ def evaluate_snapshot(
         detector_snapshot: Only for TD models. If defined, evaluation metrics are
             computed using the detections made by this snapshot
     """
+    head_type = loader.model_cfg["model"]["heads"]["bodypart"]["type"]
+    if head_type == "DLCRNetHead":
+        prune_paf_graph.benchmark_paf_graphs(
+            loader=loader, snapshot_path=snapshot.path, verbose=False,
+        )
+
     parameters = loader.get_dataset_parameters()
     pcutoff = cfg.get("pcutoff", 0.6)
 
diff --git a/deeplabcut/pose_estimation_pytorch/apis/prune_paf_graph.py b/deeplabcut/pose_estimation_pytorch/apis/prune_paf_graph.py
new file mode 100644
index 0000000000..9d49e0f650
--- /dev/null
+++ b/deeplabcut/pose_estimation_pytorch/apis/prune_paf_graph.py
@@ -0,0 +1,268 @@
+#
+# DeepLabCut Toolbox (deeplabcut.org)
+# © A. & M.W. Mathis Labs
+# https://github.com/DeepLabCut/DeepLabCut
+#
+# Please see AUTHORS for contributors.
+# https://github.com/DeepLabCut/DeepLabCut/blob/main/AUTHORS
+#
+# Licensed under GNU Lesser General Public License v3.0
+#
+from __future__ import annotations
+
+from collections import defaultdict
+from pathlib import Path
+
+import networkx as nx
+import numpy as np
+import torch
+from tqdm import tqdm
+
+import deeplabcut.core.metrics as metrics
+import deeplabcut.pose_estimation_pytorch.apis.utils as utils
+import deeplabcut.pose_estimation_pytorch.data as data
+import deeplabcut.pose_estimation_pytorch.models.predictors as predictors
+import deeplabcut.utils.auxiliaryfunctions as auxiliaryfunctions
+from deeplabcut.core.crossvalutils import find_closest_neighbors
+from deeplabcut.pose_estimation_pytorch.models import PoseModel
+from deeplabcut.pose_estimation_pytorch.models.predictors.paf_predictor import Graph
+
+
+@torch.no_grad()
+def benchmark_paf_graphs(
+    loader: data.Loader,
+    snapshot_path: Path,
+    verbose: bool = False,
+) -> list[dict]:
+    runner = utils.get_pose_inference_runner(loader.model_cfg, snapshot_path)
+    device = runner.device
+    preprocessor = runner.preprocessor
+    model = runner.model
+    predictor = model.heads.bodypart.predictor
+
+    # only benchmark the PAF graph if the PAF indices contain all edges
+    if len(predictor.edges_to_keep) < len(predictor.graph):
+        return [dict(edges_to_keep=predictor.edges_to_keep)]
+
+    model.to(device)
+    model.eval()
+
+    if not isinstance(predictor, predictors.PartAffinityFieldPredictor):
+        raise ValueError(f"Predictor should be a PartAffinityFieldPredictor.")
+
+    if verbose:
+        print("-------------------------------------------------")
+        print("Benchmarking different Part-Affinity Field Graphs")
+        print("  (1/3) Obtaining the best graph candidates")
+
+    gt_train = loader.ground_truth_keypoints("train")
+    best_paf_edges, _ = get_n_best_paf_graphs(
+        model, gt_train, preprocessor, device, predictor.graph, n_graphs=10,
+    )
+
+    if verbose:
+        print("  (2/3) Running test inference")
+
+    gt_test = loader.ground_truth_keypoints("test")
+    images_test = [img_path for img_path in gt_test]
+
+    predictions = {graph_id: {} for graph_id in range(len(best_paf_edges))}
+    with torch.no_grad():
+        for image_path in tqdm(images_test):
+            image, _ = preprocessor(image_path, {})
+            outputs = model(image.to(device))
+            for graph_id, edges in enumerate(best_paf_edges):
+                predictor.set_paf_edges_to_keep(edges)
+                pred_pose = model.get_predictions(outputs)["bodypart"]["poses"]
+                predictions[graph_id][image_path] = pred_pose.cpu().numpy()[0]
+
+    if verbose:
+        print("  (3/3) Evaluating Graphs")
+
+    results = []
+    for graph_id, pred_pose in predictions.items():
+        edges_to_keep = [int(i) for i in best_paf_edges[graph_id]]
+        graph_metrics = metrics.compute_metrics(
+            gt_test,
+            pred_pose,
+            single_animal=False,
+            pcutoff=0.6,
+        )
+        results.append(
+            dict(
+                edges_to_keep=edges_to_keep,
+                key_metric=graph_metrics["mAP"],
+                metrics=graph_metrics,
+            )
+        )
+
+        if verbose:
+            print("    ---")
+            print(f"    |Graph {graph_id}: {len(edges_to_keep)} edges")
+            print(f"    |   mAP: {graph_metrics['mAP']}")
+            print(f"    |   mAR: {graph_metrics['mAR']}")
+            print(f"    |   edges: {edges_to_keep}")
+            print()
+
+    results = list(sorted(results, key=lambda r: 1 - r["key_metric"]))
+    if len(results) > 0:
+        best_results = results[0]
+        best_edges = best_results["edges_to_keep"]
+        graph_metrics = best_results["metrics"]
+
+        if verbose:
+            print("Selecting the following Graph")
+            print(60 * "-")
+            print(f"|Graph {graph_id}: {len(best_edges)} edges")
+            print(f"|   mAP: {graph_metrics['mAP']}")
+            print(f"|   mAR: {graph_metrics['mAR']}")
+            print(f"|   edges: {best_edges}")
+            print()
+
+        # update the edges to keep in the PyTorch configuration file
+        head_update = dict(predictor=dict(edges_to_keep=best_edges))
+        loader.update_model_cfg(dict(model=dict(heads=dict(bodypart=head_update))))
+
+        # update the edges indices
+        test_config = loader.model_folder.parent / "test" / "pose_cfg.yaml"
+        auxiliaryfunctions.edit_config(
+            str(test_config), dict(paf_best=best_edges)
+        )
+
+    return results
+
+
+def _calc_separability(
+    vals_left: np.ndarray,
+    vals_right: np.ndarray,
+    n_bins: int = 101,
+    metric: str = "jeffries",
+    max_sensitivity: bool = False,
+) -> tuple[float, float]:
+    if metric not in ("jeffries", "auc"):
+        raise ValueError("`metric` should be either 'jeffries' or 'auc'.")
+
+    bins = np.linspace(0, 1, n_bins)
+    hist_left = np.histogram(vals_left, bins=bins)[0]
+    hist_left = hist_left / hist_left.sum()
+    hist_right = np.histogram(vals_right, bins=bins)[0]
+    hist_right = hist_right / hist_right.sum()
+    tpr = np.cumsum(hist_right)
+    if metric == "jeffries":
+        sep = np.sqrt(
+            2 * (1 - np.sum(np.sqrt(hist_left * hist_right)))
+        )  # Jeffries-Matusita distance
+    else:
+        sep = np.trapz(np.cumsum(hist_left), tpr)
+    if max_sensitivity:
+        threshold = bins[max(1, np.argmax(tpr > 0))]
+    else:
+        threshold = bins[np.argmin(1 - np.cumsum(hist_left) + tpr)]
+    return sep, threshold
+
+
+@torch.no_grad()
+def compute_within_between_paf_costs(
+    model: PoseModel,
+    ground_truth: dict[str, np.ndarray],
+    preprocessor: data.Preprocessor,
+    device: str,
+) -> tuple[defaultdict, defaultdict]:
+    predictor = model.heads.bodypart.predictor
+    images = [img_path for img_path in ground_truth]
+
+    within = defaultdict(list)
+    between = defaultdict(list)
+    for image_path in tqdm(images):
+        image, _ = preprocessor(image_path, {})
+        outputs = model(image.to(device))
+        preds = model.get_predictions(outputs)["bodypart"]["preds"][0]
+        gt_pose_with_vis = ground_truth[image_path].transpose((1, 0, 2))
+
+        # mask non-visible keypoints
+        gt_pose = gt_pose_with_vis[..., :2].copy()
+        gt_pose[gt_pose_with_vis[..., 2] <= 0] = np.nan
+
+        if np.isnan(gt_pose).all():
+            continue
+
+        coords_pred = preds["coordinates"][0]
+        costs_pred = preds["costs"]
+
+        # Get animal IDs and corresponding indices in the arrays of detections
+        lookup = dict()
+        for i, (coord_pred, coord_gt) in enumerate(zip(coords_pred, gt_pose)):
+            inds = np.flatnonzero(np.all(~np.isnan(coord_pred), axis=1))
+            inds_gt = np.flatnonzero(np.all(~np.isnan(coord_gt), axis=1))
+            if inds.size and inds_gt.size:
+                neighbors = find_closest_neighbors(
+                    coord_gt[inds_gt], coord_pred[inds], k=3
+                )
+                found = neighbors != -1
+                lookup[i] = dict(zip(inds_gt[found], inds[neighbors[found]]))
+
+        for k, v in costs_pred.items():
+            paf = v["m1"]
+            mask_within = np.zeros(paf.shape, dtype=bool)
+            s, t = predictor.graph[k]
+            if s not in lookup or t not in lookup:
+                continue
+            lu_s = lookup[s]
+            lu_t = lookup[t]
+            common_id = set(lu_s).intersection(lu_t)
+            for id_ in common_id:
+                mask_within[lu_s[id_], lu_t[id_]] = True
+            within_vals = paf[mask_within]
+            between_vals = paf[~mask_within]
+            within[k].extend(within_vals)
+            between[k].extend(between_vals)
+
+    return within, between
+
+
+def get_n_best_paf_graphs(
+    model: PoseModel,
+    ground_truth: dict[str, np.ndarray],
+    preprocessor: data.Preprocessor,
+    device: str,
+    full_graph: Graph,
+    root_edges: list[int] | None = None,
+    n_graphs: int = 10,
+    metric: str = "auc",
+) -> tuple[list[list[int]], dict[int, float]]:
+    return_preds = model.heads.bodypart.predictor.return_preds
+    model.heads.bodypart.predictor.return_preds = True
+
+    within_train, between_train = compute_within_between_paf_costs(
+        model, ground_truth, preprocessor, device
+    )
+    existing_edges = list(set(k for k, v in within_train.items() if v))
+
+    scores, _ = zip(
+        *[
+            _calc_separability(between_train[n], within_train[n], metric=metric)
+            for n in existing_edges
+        ]
+    )
+
+    # Find minimal skeleton
+    G = nx.Graph()
+    for edge, score in zip(existing_edges, scores):
+        if np.isfinite(score):
+            G.add_edge(*full_graph[edge], weight=score)
+
+    order = np.asarray(existing_edges)[np.argsort(scores)[::-1]]
+    if root_edges is None:
+        root_edges = []
+        for edge in nx.maximum_spanning_edges(G, data=False):
+            root_edges.append(full_graph.index(sorted(edge)))
+
+    n_edges = len(existing_edges) - len(root_edges)
+    lengths = np.linspace(0, n_edges, min(n_graphs, n_edges + 1), dtype=int)[1:]
+    order = order[np.isin(order, root_edges, invert=True)]
+    best_edges = [root_edges]
+    for length in lengths:
+        best_edges.append(root_edges + list(order[:length]))
+
+    model.heads.bodypart.predictor.return_preds = return_preds
+    return best_edges, dict(zip(existing_edges, scores))
diff --git a/deeplabcut/pose_estimation_pytorch/data/__init__.py b/deeplabcut/pose_estimation_pytorch/data/__init__.py
index 9d3875347f..a578b101d4 100644
--- a/deeplabcut/pose_estimation_pytorch/data/__init__.py
+++ b/deeplabcut/pose_estimation_pytorch/data/__init__.py
@@ -15,4 +15,15 @@
     PoseDatasetParameters,
     PoseDataset,
 )
+from deeplabcut.pose_estimation_pytorch.data.postprocessor import (
+    build_bottom_up_postprocessor,
+    build_detector_postprocessor,
+    build_top_down_postprocessor,
+    Postprocessor,
+)
+from deeplabcut.pose_estimation_pytorch.data.preprocessor import (
+    build_bottom_up_preprocessor,
+    build_top_down_preprocessor,
+    Preprocessor,
+)
 from deeplabcut.pose_estimation_pytorch.data.transforms import build_transforms
diff --git a/deeplabcut/pose_estimation_pytorch/models/predictors/paf_predictor.py b/deeplabcut/pose_estimation_pytorch/models/predictors/paf_predictor.py
index 0c9f757c7a..f7c64ef52c 100644
--- a/deeplabcut/pose_estimation_pytorch/models/predictors/paf_predictor.py
+++ b/deeplabcut/pose_estimation_pytorch/models/predictors/paf_predictor.py
@@ -364,3 +364,12 @@ def compute_peaks_and_costs(
             peaks_and_costs.append(dict_)
 
         return peaks_and_costs
+
+    def set_paf_edges_to_keep(self, edge_indices: list[int]) -> None:
+        """Sets the PAF edge indices to use to assemble individuals
+
+        Args:
+            edge_indices: The indices of edges in the graph to keep.
+        """
+        self.edges_to_keep = edge_indices
+        self.assembler.paf_inds = edge_indices
diff --git a/deeplabcut/pose_estimation_pytorch/models/predictors/utils.py b/deeplabcut/pose_estimation_pytorch/models/predictors/utils.py
deleted file mode 100644
index b0df0ecdb4..0000000000
--- a/deeplabcut/pose_estimation_pytorch/models/predictors/utils.py
+++ /dev/null
@@ -1,193 +0,0 @@
-#
-# DeepLabCut Toolbox (deeplabcut.org)
-# © A. & M.W. Mathis Labs
-# https://github.com/DeepLabCut/DeepLabCut
-#
-# Please see AUTHORS for contributors.
-# https://github.com/DeepLabCut/DeepLabCut/blob/main/AUTHORS
-#
-# Licensed under GNU Lesser General Public License v3.0
-#
-from __future__ import annotations
-
-from collections import defaultdict
-
-import albumentations as A
-import networkx as nx
-import numpy as np
-import torch
-from numpy.typing import ArrayLike, NDArray
-from torch.utils.data import DataLoader
-from tqdm import tqdm
-
-import deeplabcut.core.metrics as metrics
-from deeplabcut.core.crossvalutils import find_closest_neighbors
-from deeplabcut.pose_estimation_pytorch import Loader
-from deeplabcut.pose_estimation_pytorch.models import PoseModel
-from deeplabcut.pose_estimation_pytorch.models.predictors.paf_predictor import Graph
-
-
-def _calc_separability(
-    vals_left: ArrayLike,
-    vals_right: ArrayLike,
-    n_bins: int = 101,
-    metric: str = "jeffries",
-    max_sensitivity: bool = False,
-) -> tuple[float, float]:
-    if metric not in ("jeffries", "auc"):
-        raise ValueError("`metric` should be either 'jeffries' or 'auc'.")
-
-    bins = np.linspace(0, 1, n_bins)
-    hist_left = np.histogram(vals_left, bins=bins)[0]
-    hist_left = hist_left / hist_left.sum()
-    hist_right = np.histogram(vals_right, bins=bins)[0]
-    hist_right = hist_right / hist_right.sum()
-    tpr = np.cumsum(hist_right)
-    if metric == "jeffries":
-        sep = np.sqrt(
-            2 * (1 - np.sum(np.sqrt(hist_left * hist_right)))
-        )  # Jeffries-Matusita distance
-    else:
-        sep = np.trapz(np.cumsum(hist_left), tpr)
-    if max_sensitivity:
-        threshold = bins[max(1, np.argmax(tpr > 0))]
-    else:
-        threshold = bins[np.argmin(1 - np.cumsum(hist_left) + tpr)]
-    return sep, threshold
-
-
-def get_n_best_paf_graphs(
-    model: PoseModel,
-    train_dataloader: DataLoader,
-    full_graph: Graph,
-    root_edges: list[int] | None = None,
-    n_graphs: int = 10,
-    metric: str = "auc",
-    device: str = "cuda",
-) -> tuple[list[list[int]], dict[int, float]]:
-    within_train, between_train = compute_within_between_paf_costs(
-        model, train_dataloader, device
-    )
-    existing_edges = list(set(k for k, v in within_train.items() if v))
-
-    scores, _ = zip(
-        *[
-            _calc_separability(between_train[n], within_train[n], metric=metric)
-            for n in existing_edges
-        ]
-    )
-
-    # Find minimal skeleton
-    G = nx.Graph()
-    for edge, score in zip(existing_edges, scores):
-        if np.isfinite(score):
-            G.add_edge(*full_graph[edge], weight=score)
-
-    order = np.asarray(existing_edges)[np.argsort(scores)[::-1]]
-    if root_edges is None:
-        root_edges = []
-        for edge in nx.maximum_spanning_edges(G, data=False):
-            root_edges.append(full_graph.index(sorted(edge)))
-
-    n_edges = len(existing_edges) - len(root_edges)
-    lengths = np.linspace(0, n_edges, min(n_graphs, n_edges + 1), dtype=int)[1:]
-    order = order[np.isin(order, root_edges, invert=True)]
-    best_edges = [root_edges]
-    for length in lengths:
-        best_edges.append(root_edges + list(order[:length]))
-    return best_edges, dict(zip(existing_edges, scores))
-
-
-def compute_within_between_paf_costs(
-    model: PoseModel, dataloader: DataLoader, device: str = "cuda"
-) -> tuple[defaultdict[list]]:
-    model.to(device)
-    predictor = model.heads.bodypart.predictor
-    within = defaultdict(list)
-    between = defaultdict(list)
-    with torch.no_grad():
-        for batch in tqdm(dataloader):
-            inputs = batch["image"].to(device)
-            preds = model.get_predictions(model(inputs))["bodypart"]
-
-            for coords_gt, preds_ in zip(
-                batch["annotations"]["keypoints"], preds["preds"]
-            ):
-                coords_gt = coords_gt.permute(1, 0, 2).detach().cpu().numpy()
-                if np.isnan(coords_gt).all():
-                    continue
-
-                coords_pred = preds_["coordinates"][0]
-                costs_pred = preds_["costs"]
-
-                # Get animal IDs and corresponding indices in the arrays of detections
-                lookup = dict()
-                for i, (coord_pred, coord_gt) in enumerate(zip(coords_pred, coords_gt)):
-                    inds = np.flatnonzero(np.all(~np.isnan(coord_pred), axis=1))
-                    inds_gt = np.flatnonzero(np.all(~np.isnan(coord_gt), axis=1))
-                    if inds.size and inds_gt.size:
-                        neighbors = find_closest_neighbors(
-                            coord_gt[inds_gt], coord_pred[inds], k=3
-                        )
-                        found = neighbors != -1
-                        lookup[i] = dict(zip(inds_gt[found], inds[neighbors[found]]))
-
-                for k, v in costs_pred.items():
-                    paf = v["m1"]
-                    mask_within = np.zeros(paf.shape, dtype=bool)
-                    s, t = predictor.graph[k]
-                    if s not in lookup or t not in lookup:
-                        continue
-                    lu_s = lookup[s]
-                    lu_t = lookup[t]
-                    common_id = set(lu_s).intersection(lu_t)
-                    for id_ in common_id:
-                        mask_within[lu_s[id_], lu_t[id_]] = True
-                    within_vals = paf[mask_within]
-                    between_vals = paf[~mask_within]
-                    within[k].extend(within_vals)
-                    between[k].extend(between_vals)
-    return within, between
-
-
-def benchmark_paf_graphs(
-    model: PoseModel,
-    loader: Loader,
-    transform: A.BaseCompose,
-    batch_size: int = 8,
-    device: str = "cuda",
-) -> tuple[list[dict[str, float]], list[dict[str, NDArray]], list[list[int]]]:
-    predictor = model.heads.bodypart.predictor
-    train_dataset = loader.create_dataset(mode="train", task="BU", transform=transform)
-    valid_dataset = loader.create_dataset(mode="test", task="BU", transform=transform)
-    train_dataloader = DataLoader(train_dataset, batch_size=batch_size, shuffle=False)
-    valid_dataloader = DataLoader(valid_dataset, batch_size=batch_size, shuffle=False)
-    best_paf_edges, _ = get_n_best_paf_graphs(
-        model, train_dataloader, predictor.graph, device=device
-    )
-    poses_gt = loader.ground_truth_keypoints("test")
-    results = []
-    poses = []
-    for edges in best_paf_edges:
-        predictor.edges_to_keep = predictor.assembler.paf_inds = edges
-        paths = []
-        poses_ = []
-        with torch.no_grad():
-            for batch in tqdm(valid_dataloader):
-                paths.extend(batch["path"])
-                inputs = batch["image"].to(device)
-                # FIXME We can do better than the repetition below
-                preds = model.get_predictions(model(inputs))["bodypart"]
-                poses_.extend(preds["poses"])
-        poses_ = torch.stack(poses_).detach().cpu().numpy()
-        poses_ = dict(zip(paths, poses_))
-        poses.append(poses_)
-        results.append(
-            metrics.compute_metrics(
-                poses_gt,
-                poses_,
-                single_animal=train_dataset.parameters.max_num_animals == 1,
-                pcutoff=0.6,
-            )
-        )
-    return results, poses, best_paf_edges

From 9106a970e89e99a332c7bc2f91079745cc3eb5fd Mon Sep 17 00:00:00 2001
From: Niels Poulsen <niels.poulsen@epfl.ch>
Date: Thu, 12 Dec 2024 11:39:54 +0100
Subject: [PATCH 2/4] add docs

---
 .../apis/prune_paf_graph.py                   | 26 +++++++++++++++++--
 1 file changed, 24 insertions(+), 2 deletions(-)

diff --git a/deeplabcut/pose_estimation_pytorch/apis/prune_paf_graph.py b/deeplabcut/pose_estimation_pytorch/apis/prune_paf_graph.py
index 9d49e0f650..8eedb8f831 100644
--- a/deeplabcut/pose_estimation_pytorch/apis/prune_paf_graph.py
+++ b/deeplabcut/pose_estimation_pytorch/apis/prune_paf_graph.py
@@ -33,7 +33,28 @@ def benchmark_paf_graphs(
     loader: data.Loader,
     snapshot_path: Path,
     verbose: bool = False,
+    overwrite: bool = False,
+    update_config: bool = True,
 ) -> list[dict]:
+    """Prunes the PAF graph to maximize performance
+
+    Args:
+        loader: The loader for the model to prune.
+        snapshot_path: The path to the snapshot with which to prune the model.
+        verbose: Verbose pruning of the model.
+        overwrite: Whether to overwrite the graph if it was already pruned.
+        update_config: Whether to update the model configuration with the pruned graph.
+
+    Returns:
+        A list of dictionaries containing results for each pruned graph.
+
+        If the graph was already pruned, a single element is returned with an
+        "edges_to_keep" key,  containing the indices of edges to keep in the graph.
+
+        Otherwise, a list of graphs that were evaluated is returned, with "key_metric",
+        "edges_to_keep" and "metrics" keys. The list is sorted by "key_metric" (which
+        is pose mAP).
+    """
     runner = utils.get_pose_inference_runner(loader.model_cfg, snapshot_path)
     device = runner.device
     preprocessor = runner.preprocessor
@@ -41,7 +62,7 @@ def benchmark_paf_graphs(
     predictor = model.heads.bodypart.predictor
 
     # only benchmark the PAF graph if the PAF indices contain all edges
-    if len(predictor.edges_to_keep) < len(predictor.graph):
+    if not overwrite and len(predictor.edges_to_keep) < len(predictor.graph):
         return [dict(edges_to_keep=predictor.edges_to_keep)]
 
     model.to(device)
@@ -105,7 +126,8 @@ def benchmark_paf_graphs(
             print()
 
     results = list(sorted(results, key=lambda r: 1 - r["key_metric"]))
-    if len(results) > 0:
+
+    if update_config and len(results) > 0:
         best_results = results[0]
         best_edges = best_results["edges_to_keep"]
         graph_metrics = best_results["metrics"]

From cb46540a0a22a40c2c8f7fd3b1820415bddbb1ac Mon Sep 17 00:00:00 2001
From: Niels Poulsen <niels.poulsen@epfl.ch>
Date: Thu, 12 Dec 2024 15:03:09 +0100
Subject: [PATCH 3/4] fix: PAF predictor preds key

---
 .../models/predictors/paf_predictor.py              | 13 ++++++++-----
 1 file changed, 8 insertions(+), 5 deletions(-)

diff --git a/deeplabcut/pose_estimation_pytorch/models/predictors/paf_predictor.py b/deeplabcut/pose_estimation_pytorch/models/predictors/paf_predictor.py
index f7c64ef52c..f891524fa9 100644
--- a/deeplabcut/pose_estimation_pytorch/models/predictors/paf_predictor.py
+++ b/deeplabcut/pose_estimation_pytorch/models/predictors/paf_predictor.py
@@ -152,11 +152,14 @@ def forward(
             heatmaps, self.nms_radius, threshold=0.01
         )
         if ~torch.any(peaks):
-            return {
-                "poses": -torch.ones(
-                    (batch_size, self.num_animals, self.num_multibodyparts, 5)
-                )
-            }
+            poses = -torch.ones(
+                (batch_size, self.num_animals, self.num_multibodyparts, 5)
+            )
+            results = dict(poses=poses)
+            if self.return_preds:
+                results["preds"] = [dict(coordinates=[[]], costs=[])],
+
+            return results
 
         locrefs = locrefs.reshape(batch_size, n_channels, 2, height, width)
         locrefs = locrefs * self.locref_stdev

From b80b6eafea84713b078625bc5f7fb0197492988a Mon Sep 17 00:00:00 2001
From: Niels Poulsen <niels.poulsen@epfl.ch>
Date: Thu, 12 Dec 2024 15:22:25 +0100
Subject: [PATCH 4/4] verbose improvement

---
 deeplabcut/pose_estimation_pytorch/apis/prune_paf_graph.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/deeplabcut/pose_estimation_pytorch/apis/prune_paf_graph.py b/deeplabcut/pose_estimation_pytorch/apis/prune_paf_graph.py
index 8eedb8f831..726674a8d7 100644
--- a/deeplabcut/pose_estimation_pytorch/apis/prune_paf_graph.py
+++ b/deeplabcut/pose_estimation_pytorch/apis/prune_paf_graph.py
@@ -135,7 +135,7 @@ def benchmark_paf_graphs(
         if verbose:
             print("Selecting the following Graph")
             print(60 * "-")
-            print(f"|Graph {graph_id}: {len(best_edges)} edges")
+            print(f"|Graph with {len(best_edges)} edges")
             print(f"|   mAP: {graph_metrics['mAP']}")
             print(f"|   mAR: {graph_metrics['mAR']}")
             print(f"|   edges: {best_edges}")