Tracking model performance with confidence threshold

[jveitchmichaelis]

It'd be useful to have a standard method to sweep model confidence and plot box_precision/recall to gauge where to set the score_threshold. I knocked up a very simple script that processes a validation dataset.

Add this or something similar as a script?

[jveitchmichaelis]

import argparse
import os
from collections.abc import Sequence
from typing import Any

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from hydra import compose, initialize, initialize_config_dir
from omegaconf import DictConfig, OmegaConf

from deepforest import utilities
from deepforest.conf.schema import Config as StructuredConfig
from deepforest.evaluate import evaluate_boxes
from deepforest.main import deepforest

DEFAULT_THRESHOLDS = np.linspace(0.0, 0.9, num=10).round(3).tolist()


def parse_args() -> tuple[argparse.Namespace, list[str]]:
    parser = argparse.ArgumentParser(
        description=(
            "Run validation across score thresholds and export precision/recall "
            "metrics. Additional Hydra overrides can be passed as key=value pairs."
        )
    )
    parser.add_argument(
        "--output-dir",
        required=True,
        help="Directory where CSV and plot will be written",
    )
    parser.add_argument(
        "--thresholds",
        type=float,
        nargs="+",
        help="Optional list of score thresholds to evaluate (default: 0.0-0.9 inclusive)",
    )
    parser.add_argument(
        "--no-label-thresholds",
        action="store_false",
        dest="label_thresholds",
        help="Disable threshold annotations on the PR curve (enabled by default)",
    )
    parser.add_argument(
        "--config-dir",
        help="Optional Hydra config directory (defaults to packaged deepforest.conf)",
    )
    parser.add_argument(
        "--config-name",
        default="config",
        help="Hydra config name (default: config)",
    )

    return parser.parse_known_args()


def load_config(args: argparse.Namespace, overrides: list[str]) -> DictConfig:
    if args.config_dir is not None:
        initialize_config_dir(version_base=None, config_dir=args.config_dir)
    else:
        initialize(version_base=None, config_path="pkg://deepforest.conf")

    base = OmegaConf.structured(StructuredConfig)
    cfg = compose(config_name=args.config_name, overrides=overrides)
    cfg = OmegaConf.merge(base, cfg)

    return cfg


def normalize_thresholds(thresholds: Sequence[float] | None) -> list[float]:
    if thresholds is None or len(thresholds) == 0:
        values = DEFAULT_THRESHOLDS
    else:
        values = thresholds

    return sorted({round(float(t), 4) for t in values})


def run_validation(cfg: DictConfig, thresholds: Sequence[float]) -> pd.DataFrame:
    records: list[dict[str, Any]] = []

    # Run predictions once with score_thresh=0.0 to capture all detections
    cfg.score_thresh = 0.0
    model = deepforest(config=cfg)

    # Generate predictions for all validation images
    predictions = model.predict_file(
        csv_file=cfg.validation.csv_file, root_dir=cfg.validation.root_dir
    )

    # Load ground truth annotations
    ground_df = utilities.read_file(
        cfg.validation.csv_file, root_dir=cfg.validation.root_dir
    )

    # For each threshold, filter predictions and evaluate
    for threshold in thresholds:
        # Filter predictions by score threshold
        predictions_filtered = predictions[predictions.score >= threshold]

        # Evaluate filtered predictions against ground truth
        results = evaluate_boxes(
            predictions=predictions_filtered,
            ground_df=ground_df,
            iou_threshold=cfg.validation.iou_threshold,
        )

        row: dict[str, Any] = {
            "score_thresh": threshold,
            "box_precision": results.get("box_precision", np.nan),
            "box_recall": results.get("box_recall", np.nan),
        }

        records.append(row)

    df = pd.DataFrame(records)
    if "box_precision" not in df.columns:
        df["box_precision"] = np.nan
    if "box_recall" not in df.columns:
        df["box_recall"] = np.nan

    return df


def plot_pr_curve(df: pd.DataFrame, output_path: str, label_thresholds: bool) -> None:
    plot_df = df.dropna(subset=["box_precision", "box_recall"])
    if plot_df.empty:
        return

    fig, ax = plt.subplots(figsize=(6, 4))
    ax.plot(plot_df["box_recall"], plot_df["box_precision"], marker="o")
    ax.set_xlabel("Recall")
    ax.set_ylabel("Precision")
    ax.set_title("Precision-Recall by score_thresh")
    ax.set_xlim(0, 1)
    ax.set_ylim(0, 1)
    ax.grid(True, linestyle="--", alpha=0.5)

    if label_thresholds:
        for _, row in plot_df.iterrows():
            ax.annotate(
                f"{row['score_thresh']:.2f}",
                (row["box_recall"], row["box_precision"]),
                textcoords="offset points",
                xytext=(4, 4),
                fontsize=8,
            )

    fig.tight_layout()
    fig.savefig(output_path, dpi=200)
    plt.close(fig)


def sweep_scores(
    cfg: DictConfig,
    output_dir: str,
    thresholds: Sequence[float] | None = None,
    label_thresholds: bool = True,
) -> tuple[str, str]:
    """Run validation sweeps and write CSV and plot.

    Returns paths to the CSV and plot.
    """

    if cfg.validation.csv_file is None:
        raise ValueError("validation.csv_file must be set in the config")
    if cfg.validation.root_dir is None:
        raise ValueError("validation.root_dir must be set in the config")

    thresholds = normalize_thresholds(thresholds)
    os.makedirs(output_dir, exist_ok=True)

    results_df = run_validation(cfg, thresholds)

    csv_path = os.path.join(output_dir, "precision_recall_thresholds.csv")
    results_df.to_csv(csv_path, index=False)

    plot_path = os.path.join(output_dir, "precision_recall_curve.png")
    plot_pr_curve(results_df, plot_path, label_thresholds)

    return csv_path, plot_path


def main() -> None:
    args, overrides = parse_args()
    cfg = load_config(args, overrides)

    sweep_scores(
        cfg=cfg,
        output_dir=args.output_dir,
        thresholds=args.thresholds,
        label_thresholds=args.label_thresholds,
    )


if __name__ == "__main__":
    main()