feat: train fixed-size classifiers and export class probabilities

2026-10-14 17:08:14 +08:00
parent 63bab36147
commit f9a4032783
11 changed files with 478 additions and 208 deletions
@@ -0,0 +1,134 @@
+#!/usr/bin/env python
+"""Compare supervised classification accuracy, export parity and YOLO agreement.
+
+Uses fixed-size human-labelled validation tiles, never training data. No resize.
+Empty detector output is normal; the strongest valid candidate supplies top1.
+"""
+from __future__ import annotations
+
+import argparse
+import copy
+import json
+import sys
+from pathlib import Path
+
+import numpy as np
+import torch
+import yaml
+
+sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
+from datasets.tile_classification import YoloTileClassificationDataset  # noqa: E402
+from models import build_model  # noqa: E402
+
+
+def output_class(output: np.ndarray, conf: float, background: int) -> int:
+    scores = output[0, 4:, :]
+    index = int(scores.max(axis=0).argmax())
+    return int(scores[:, index].argmax()) if scores[:, index].max() >= conf else background
+
+
+def probability_class(output: np.ndarray, conf: float, background: int) -> int:
+    if output.ndim != 2 or output.shape[0] != 1:
+        raise ValueError("Expected classifier probabilities [1,C]")
+    return int(output[0].argmax()) if output[0].max() >= conf else background
+
+
+def matrix_report(matrix, names, background):
+    counts = matrix.sum(1)
+    foreground = int(matrix.sum() - counts[background])
+    return {"agreement": float(np.trace(matrix) / matrix.sum()),
+            "foreground_agreement": float((np.trace(matrix) - matrix[background, background])
+                                          / foreground) if foreground else None,
+            "per_reference_class": {name: {"support": int(counts[i]),
+                                             "agreement": float(matrix[i, i] / counts[i])
+                                             if counts[i] else None}
+                                    for i, name in enumerate(names)},
+            "confusion_reference_rows_prediction_columns": matrix.tolist()}
+
+
+def main():
+    import openvino as ov
+
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--teacher", required=True, help="Reference YOLO .xml; used only for comparison")
+    parser.add_argument("--student", required=True, help="Classifier probability .xml, output [1,C]")
+    parser.add_argument("--checkpoint", required=True)
+    parser.add_argument("--report", required=True)
+    parser.add_argument("--conf", type=float, default=0.25)
+    parser.add_argument("--device", default="cuda:0")
+    args = parser.parse_args()
+    if not 0 < args.conf < 1:
+        raise ValueError("Positive confidence threshold required")
+    state = torch.load(args.checkpoint, map_location="cpu", weights_only=False)
+    cfg = state["config"]
+    names = cfg["data"]["classes"]
+    background = names.index("normal")
+    for xml in (args.teacher, args.student):
+        meta = yaml.safe_load((Path(xml).parent / "metadata.yaml").read_text())
+        if [meta["names"][i] for i in range(len(meta["names"]))] != names:
+            raise ValueError(f"Class order mismatch: {xml}")
+    if cfg["data"]["val_transforms"] != [{"name": "ToTensor"}]:
+        raise ValueError("This comparison expects the new raw RGB / 255 tile classifiers")
+    ds = YoloTileClassificationDataset(cfg["data"]["yaml"], "val", cfg["data"]["image_size"],
+                                      transforms=[{"name": "ToTensor"}], classes=names)
+    model_cfg = copy.deepcopy(cfg["model"])
+    model_cfg["backbone"]["pretrained"] = False
+    model_cfg["init_from"] = None
+    model_cfg["num_classes"] = len(names)
+    model = build_model(model_cfg).to(args.device).eval()
+    model.load_state_dict(state["model"])
+    core = ov.Core()
+    settings = {"PERFORMANCE_HINT": "LATENCY", "INFERENCE_NUM_THREADS": 4}
+    teacher = core.compile_model(args.teacher, "CPU", settings)
+    student = core.compile_model(args.student, "CPU", settings)
+    for compiled in (teacher, student):
+        if list(compiled.input().shape) != [1, 3, *cfg["data"]["image_size"]]:
+            raise ValueError("Unexpected model interface")
+    if (list(teacher.output().shape)[:2] != [1, 4 + len(names)]
+            or list(student.output().shape) != [1, len(names)]):
+        raise ValueError("Expected raw YOLO output and plain classifier probabilities")
+    matrices = {key: np.zeros((len(names), len(names)), dtype=int)
+                for key in ["pt_accuracy", "ov_accuracy", "ov_deployment_accuracy",
+                            "yolo_deployment_accuracy",
+                            "export_top1_parity", "export_deployment_parity", "yolo_student_agreement"]}
+    records = []
+    with torch.inference_mode():
+        for start in range(0, len(ds), 64):
+            batch = [ds[i] for i in range(start, min(start + 64, len(ds)))]
+            images = torch.stack([b["image"] for b in batch])
+            probabilities = model(images.to(args.device)).softmax(1).cpu().numpy()
+            for i, item in enumerate(batch):
+                x = images[i:i+1].numpy()
+                t = output_class(teacher(x)[teacher.output()], args.conf, background)
+                raw = student(x)[student.output()]
+                if (not np.isfinite(raw).all() or (raw < 0).any() or (raw > 1).any()
+                        or not np.allclose(raw.sum(1), 1, atol=1e-5)):
+                    raise ValueError("Invalid classifier probability output")
+                s = probability_class(raw, args.conf, background)
+                raw_s = int(raw[0].argmax())
+                p = int(probabilities[i].argmax())
+                deployed_p = p if probabilities[i].max() >= args.conf else background
+                y = item["label"]
+                for key, row, col in [("pt_accuracy", y, p), ("ov_accuracy", y, raw_s),
+                                      ("ov_deployment_accuracy", y, s),
+                                      ("yolo_deployment_accuracy", y, t),
+                                      ("export_top1_parity", p, raw_s),
+                                      ("export_deployment_parity", deployed_p, s),
+                                      ("yolo_student_agreement", t, s)]:
+                    matrices[key][row, col] += 1
+                records.append({"path": item["meta"]["path"], "truth": y, "pytorch": p,
+                                "pytorch_deployed": deployed_p, "openvino": s,
+                                "openvino_top1": raw_s, "yolo": t,
+                                "student_max_score": float(raw.max())})
+    report = {"classes": names, "images": len(ds), "conf": args.conf,
+              "teacher": args.teacher, "student": args.student, "checkpoint": args.checkpoint,
+              "metrics": {k: matrix_report(m, names, background) for k, m in matrices.items()}}
+    out = Path(args.report)
+    out.parent.mkdir(parents=True, exist_ok=True)
+    out.write_text(json.dumps(report, indent=2))
+    out.with_suffix(".predictions.json").write_text(json.dumps(records, indent=2))
+    print(json.dumps(report, indent=2), flush=True)
+
+
+if __name__ == "__main__":
+    main()