File size: 5,648 Bytes
f234112
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
"""Visual debugging of the Phase 5 baseline detector.

Generates reports/figures/accident/debug/ containing ground-truth examples and
baseline-model predictions with predicted classes and confidence values, then
summarizes the baseline's failure mode (which classes it emits, confidence
distribution).

Uses the IMMUTABLE Phase 5 checkpoint (models/accident/best_detector.pth),
which has a 2-class head — loaded explicitly with num_classes=2.
"""

from __future__ import annotations

import json
import sys
from collections import Counter
from pathlib import Path

import numpy as np
import torch

if str(Path(__file__).resolve().parents[2] / "src") not in sys.path:
    sys.path.insert(0, str(Path(__file__).resolve().parents[2] / "src"))

from accident import CLASS_NAMES
from accident.dataset import AccidentParquetDataset
from accident.model import build_detector
from accident.utils import FIGURE_DIR, MODEL_DIR, get_device, seed_everything
from accident.visualization import save_ground_truth_example, save_prediction_example

NUM_EXAMPLES = 24
SCORE_THRESHOLD = 0.05


def main() -> int:
    seed_everything(42)
    device = get_device()
    out_dir = FIGURE_DIR / "debug"
    out_dir.mkdir(parents=True, exist_ok=True)

    # --- Ground-truth examples -------------------------------------------------
    ds = AccidentParquetDataset(split="validation", target_size=(320, 320))
    # Pick a spread of images: accident-heavy, non-accident-heavy, mixed.
    indices: list[int] = []
    for i in range(len(ds)):
        item = ds[i]
        labels = item["labels"].numpy()
        n_acc = int((labels == 0).sum())
        n_non = int((labels == 1).sum())
        if n_acc >= 2 and len(indices) % 2 == 0 and len(indices) < NUM_EXAMPLES:
            indices.append(i)
        elif n_non > n_acc and len(indices) < NUM_EXAMPLES:
            indices.append(i)
    indices = indices[:NUM_EXAMPLES]

    for i in indices:
        item = ds[i]
        gt_labels = item["labels"].numpy()
        n_acc = int((gt_labels == 0).sum())
        n_non = int((gt_labels == 1).sum())
        save_ground_truth_example(
            item,
            out_dir / f"gt_{i:04d}.png",
            title=f"GT — image {i}: {n_acc} accident / {n_non} non_accident",
        )

    # --- Baseline predictions --------------------------------------------------
    model = build_detector(num_classes=2, pretrained=False)
    model.load_state_dict(
        torch.load(MODEL_DIR / "best_detector.pth", map_location=device, weights_only=True)
    )
    model.to(device)
    model.eval()

    label_counts: Counter = Counter()
    score_by_label: dict[int, list[float]] = {0: [], 1: []}
    summary_rows: list[dict] = []

    with torch.no_grad():
        for i in indices:
            item = ds[i]
            outputs = model([item["image"].to(device)])
            output = outputs[0]
            keep = output["scores"] >= SCORE_THRESHOLD
            pred_boxes = output["boxes"][keep].cpu().numpy()
            # The baseline checkpoint was trained with the un-shifted (buggy)
            # label space, where its single foreground class 1 was supervised
            # ONLY by non_accident boxes. Interpret it with that identity
            # mapping, not the corrected +1 mapping.
            pred_labels = output["labels"][keep].cpu().numpy()
            pred_scores = output["scores"][keep].cpu().numpy()

            for lbl, sc in zip(pred_labels.tolist(), pred_scores.tolist()):
                label_counts[lbl] += 1
                score_by_label[lbl].append(sc)

            save_prediction_example(
                image_tensor=item["image"],
                gt_boxes=item["boxes"].numpy(),
                gt_labels=item["labels"].numpy(),
                pred_boxes=pred_boxes,
                pred_labels=pred_labels,
                pred_scores=pred_scores,
                save_path=out_dir / f"baseline_pred_{i:04d}.png",
                title=f"Baseline prediction — image {i} (thr={SCORE_THRESHOLD})",
            )
            summary_rows.append(
                {
                    "image_id": i,
                    "gt_accident": int((item["labels"].numpy() == 0).sum()),
                    "gt_non_accident": int((item["labels"].numpy() == 1).sum()),
                    "pred_accident": int((pred_labels == 0).sum()),
                    "pred_non_accident": int((pred_labels == 1).sum()),
                    "max_score": float(pred_scores.max()) if len(pred_scores) else 0.0,
                }
            )

    all_scores = [s for scores in score_by_label.values() for s in scores]
    report = {
        "checkpoint": "models/accident/best_detector.pth (Phase 5 baseline, num_classes=2)",
        "num_examples": len(indices),
        "score_threshold": SCORE_THRESHOLD,
        "predicted_label_counts_dataset_space": {
            CLASS_NAMES.get(k, str(k)): v for k, v in sorted(label_counts.items())
        },
        "mean_score_by_label": {
            CLASS_NAMES.get(k, str(k)): round(float(np.mean(v)), 4) if v else None
            for k, v in score_by_label.items()
        },
        "max_score_overall": round(float(np.max(all_scores)), 4) if all_scores else None,
        "images": summary_rows,
    }
    with open(out_dir / "baseline_debug_summary.json", "w", encoding="utf-8") as f:
        json.dump(report, f, indent=2)

    print("Predicted class counts (dataset space):", dict(report["predicted_label_counts_dataset_space"]))
    print("Mean score by class:", report["mean_score_by_label"])
    print(f"Debug visuals written to {out_dir}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())