File size: 6,835 Bytes
2fe488a
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
#!/usr/bin/env python
"""Evaluate predictions against ground truth and save results.

Discovers predictions from predictions/{username}/{method}/split_{fold}/*.geff,
loads matching GT from DATASET_PATH/{name}.geff, computes edge and division
TP/FP/FN per sample via :func:`biohub_tracking.metrics.evaluate`, and upserts
per-sample counts into a shared SQLite DB (WAL mode for concurrency).

Per-run summary prints **cumulative (micro-averaged) Jaccard** for edges and
divisions — i.e. TP/FP/FN are summed across all samples of the run and Jaccard
is computed from those totals.

Usage:
    python scripts/evaluate.py
    python scripts/evaluate.py --split 0
"""

import argparse
import json
from pathlib import Path

import tracksdata as td
from geff import GeffMetadata
from tqdm import tqdm

from biohub_tracking.io import open_dataset
from biohub_tracking.metrics import (
    evaluate as compute_metric,
    nan_metrics_row,
    node_recall,
    per_sample_metrics,
    summarise,
)

from dataspec import DATASET_PATH, INTERACTIVE, PREDICTIONS_PATH, USERNAME

PREDICTIONS_DIR = PREDICTIONS_PATH
DATA_DIR = DATASET_PATH


def _read_estimated_n_total(geff_path: Path) -> float:
    """Read ``estimated_number_of_nodes`` from a GEFF file's metadata extras.

    Returns NaN when the key is missing or the file can't be read.
    """
    try:
        meta = GeffMetadata.read(geff_path)
    except Exception:
        return float("nan")
    val = (meta.extra or {}).get("estimated_number_of_nodes")
    return float(val) if val is not None else float("nan")


def discover_runs(
    predictions_dir: Path,
    username: str | None = None,
    method: str | None = None,
    fold: int | None = None,
) -> list[dict]:
    """Find all (username, method, fold) combinations with predictions.

    Layout: predictions/{username}/{method}/split_{fold}/*.geff
    """
    runs = []
    users = [username] if username else sorted(p.name for p in predictions_dir.iterdir() if p.is_dir())
    for u in users:
        u_dir = predictions_dir / u
        if not u_dir.is_dir():
            continue
        methods = [method] if method else sorted(p.name for p in u_dir.iterdir() if p.is_dir())
        for m in methods:
            m_dir = u_dir / m
            if not m_dir.is_dir():
                continue
            splits = [f"split_{fold}"] if fold is not None else sorted(p.name for p in m_dir.iterdir() if p.is_dir() and p.name.startswith("split_"))
            for s in splits:
                s_dir = m_dir / s
                if not s_dir.is_dir():
                    continue
                geffs = sorted(s_dir.glob("*.geff"))
                if geffs:
                    runs.append({"username": u, "method": m, "split": s, "dir": s_dir, "geffs": geffs})
    return runs


def _nan_row(username: str, method: str, split: str, dataset: str) -> dict:
    return {
        "username": username, "method": method,
        "split": split, "dataset": dataset,
        **nan_metrics_row(),
    }


def evaluate_run(run: dict, max_distance: float | None = None) -> list[dict]:
    """Evaluate all predictions in a single run, return per-sample results."""
    username = run["username"]
    method = run["method"]
    split = run["split"]
    results: list[dict] = []

    # Check for missing predictions against the splits file
    splits_file = DATA_DIR / "dataset_splits.json"
    if splits_file.exists():
        fold_idx = int(split.split("_")[1])
        folds = json.loads(splits_file.read_text())
        expected = set(folds[fold_idx]["test"])
        found = {p.stem for p in run["geffs"]}
        missing = sorted(expected - found)
        if missing:
            print(f"  WARNING: {len(missing)} missing predictions for {username}/{method}/{split}: {missing[:5]}{'...' if len(missing) > 5 else ''}")
            for name in missing:
                results.append(_nan_row(username, method, split, name))

    desc = f"  {username}/{method}/{split}"
    for pred_path in tqdm(run["geffs"], desc=desc, leave=False, disable=not INTERACTIVE):
        name = pred_path.stem
        gt_path = DATA_DIR / f"{name}.geff"

        if not gt_path.exists():
            print(f"  WARNING: GT not found for {name}, skipping")
            continue

        try:
            ds = open_dataset(DATA_DIR / name, require_tracks=True)
            pred_result = td.graph.IndexedRXGraph.from_geff(pred_path)
            pred_graph = pred_result[0] if isinstance(pred_result, tuple) else pred_result

            kwargs: dict = dict(scale=ds.scale)
            if max_distance is not None:
                kwargs["max_distance"] = max_distance
            er = compute_metric(pred_graph, ds.tracks, **kwargs)

            # node_recall needs the pred graph to have been matched, which
            # evaluate() only does when pred has edges + nodes.
            if pred_graph.num_edges() > 0 and pred_graph.num_nodes() > 0:
                recall = node_recall(pred_graph, ds.tracks)
            else:
                recall = 0.0

            n_total = _read_estimated_n_total(gt_path)
            row = {
                "username": username, "method": method,
                "split": split, "dataset": name,
                **per_sample_metrics(er, n_total, recall),
            }
        except Exception as e:
            print(f"  ERROR evaluating {name}: {e}")
            row = _nan_row(username, method, split, name)

        results.append(row)

    return results


def main():
    parser = argparse.ArgumentParser(description="Evaluate tracking predictions.")
    parser.add_argument("--method", type=str, default=None, help="Evaluate only this method. Default: all discovered.")
    parser.add_argument("--split", type=str, default=None, help="Split index (0-4) or 'all'. Default: all discovered.")
    parser.add_argument("--max-distance", type=float, default=None)
    args = parser.parse_args()

    fold = None if args.split is None or args.split == "all" else int(args.split)

    runs = discover_runs(PREDICTIONS_DIR, username=USERNAME, method=args.method, fold=fold)
    if not runs:
        print("No predictions found.")
        return

    print(f"Found {len(runs)} run(s)")

    for run in runs:
        results = evaluate_run(run, max_distance=args.max_distance)

        s = summarise(results)
        print(
            f"  {run['username']}/{run['method']}/{run['split']}: "
            f"score={s['score']:.4f}  "
            f"edge_jaccard={s['edge_jaccard']:.4f}  "
            f"adj_edge_jaccard={s['adj_edge_jaccard']:.4f} (n_adj={s['n_adj']})  "
            f"division_jaccard={s['division_jaccard']:.4f} "
            f"(TP={s['division_tp']} FP={s['division_fp']} FN={s['division_fn']})  "
            f"node_recall={s['node_recall']:.4f}  (n={s['n']})"
        )


if __name__ == "__main__":
    main()