135 lines
6.8 KiB
Python
135 lines
6.8 KiB
Python
#!/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()
|