deeplabcut.pose_estimation_pytorch.apis.prune_paf_graph

Functions:

Name	Description
benchmark_paf_graphs	Prunes the PAF graph to maximize performance.

benchmark_paf_graphs

benchmark_paf_graphs(
    loader: Loader, snapshot_path: Path, verbose: bool = False, overwrite: bool = False, update_config: bool = True
) -> list[dict]

Prunes the PAF graph to maximize performance.

Parameters:

Name	Type	Description	Default
loader	Loader	The loader for the model to prune.	required
snapshot_path	Path	The path to the snapshot with which to prune the model.	required
verbose	bool	Verbose pruning of the model.	False
overwrite	bool	Whether to overwrite the graph if it was already pruned.	False
update_config	bool	Whether to update the model configuration with the pruned graph.	True

Returns:

Type	Description
list[dict]	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).

Source code in deeplabcut/pose_estimation_pytorch/apis/prune_paf_graph.py

@torch.no_grad()
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
    model = runner.model
    predictor = model.heads.bodypart.predictor

    # only benchmark the PAF graph if the PAF indices contain all edges
    if not overwrite and 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("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 update_config and 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 with {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
        loader.update_model_cfg({"model.heads.bodypart.predictor.edges_to_keep": best_edges})

        # 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