xie_x1andClaude Opus 5.5
0f9a3f496f
Code, configs, charge LUTs, best.pth and evaluation records for each training run, the knife-edge / Siemens-star / XRF notebooks with outputs, and the overnight run reports. Epoch checkpoints, inference results and HTML exports stay local (see .gitignore). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
345 lines
16 KiB
Python
345 lines
16 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
WP4 -- evaluate a grid model on the fixed-occupancy MC scan.
|
|
|
|
python scripts/04_evaluate.py --checkpoint Results/<exp>/Models/best.pth
|
|
python scripts/04_evaluate.py --checkpoint ... --data <folder>/eval_scan
|
|
|
|
Produces, per occupancy point:
|
|
efficiency / purity at the working threshold, position RMS and sigma68,
|
|
energy bias/resolution, sub-pixel non-uniformity (risk R1),
|
|
plus the threshold ROC, the pair-resolution curve (efficiency vs distance to the
|
|
nearest other photon -- the figure of merit that replaces "efficiency vs
|
|
occupancy"), and residual vs conversion depth z.
|
|
"""
|
|
import argparse
|
|
import glob
|
|
import json
|
|
import os
|
|
import re
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import matplotlib
|
|
matplotlib.use("Agg")
|
|
import matplotlib.pyplot as plt
|
|
import numpy as np
|
|
import torch
|
|
from omegaconf import OmegaConf
|
|
from torch.utils.data import DataLoader
|
|
|
|
HERE = Path(__file__).resolve().parent
|
|
sys.path.insert(0, str(HERE.parent / "src"))
|
|
from datasets import GridFrameDataset, collate_with_photons # noqa: E402
|
|
from decode import decode # noqa: E402
|
|
from metrics import (DetectionStats, greedy_match, # noqa: E402
|
|
nearest_neighbour_distance)
|
|
from models import build_model # noqa: E402
|
|
|
|
THRESHOLDS = np.array([.05, .1, .2, .3, .4, .5, .6, .7, .8, .9, .95, .98])
|
|
NEIGHBOUR_BIN_EDGES = np.array([0, .5, 1, 1.5, 2, 2.5, 3, 4, 6, 1e9])
|
|
ENERGY_BIN_WIDTH = 2.5 # keV; one bin only for a
|
|
# monochromatic set, which is fine
|
|
|
|
|
|
def energy_bin_edges(energies, width=ENERGY_BIN_WIDTH):
|
|
"""2.5 keV bins on a fixed grid, spanning the energies of the set."""
|
|
low = np.floor(energies.min() / width) * width
|
|
high = np.ceil(energies.max() / width) * width
|
|
return np.arange(low, high + width / 2, width)
|
|
|
|
|
|
def occupancy_of_filename(path):
|
|
match = re.search(r"occ([0-9p]+)pct", os.path.basename(path))
|
|
return float(match.group(1).replace("p", ".")) / 100 if match else float("nan")
|
|
|
|
|
|
@torch.no_grad()
|
|
def run_file(model, conf, path, halo, batch_size=512, workers=8, max_frames=None,
|
|
roc_frames=2000):
|
|
dataset = GridFrameDataset([path], augment=False, return_photons=True,
|
|
input_scale=float(conf.data.input_scale),
|
|
energy_scale=float(conf.data.energy_scale),
|
|
max_frames=max_frames, name=os.path.basename(path))
|
|
loader = DataLoader(dataset, batch_size=batch_size, shuffle=False,
|
|
num_workers=workers, collate_fn=collate_with_photons,
|
|
pin_memory=True)
|
|
frame_size = dataset.frame_size
|
|
working_threshold = float(conf.eval.threshold)
|
|
lowest_threshold = float(THRESHOLDS.min())
|
|
|
|
stats = DetectionStats(pixel_size_um=float(conf.eval.pixel_size_um),
|
|
match_radius=float(conf.eval.match_radius),
|
|
halo=halo, frame_size=frame_size)
|
|
roc_predicted, roc_true = [], []
|
|
|
|
for frames, objectness, offset, energy, photons in loader:
|
|
frames = frames.cuda(non_blocking=True)
|
|
prediction = model(frames).float()
|
|
points = decode(prediction, threshold=lowest_threshold,
|
|
merge_distance=float(conf.eval.merge_distance), halo=halo,
|
|
predict_energy=bool(conf.model.predict_energy),
|
|
energy_scale=float(conf.data.energy_scale))
|
|
points = [frame_points.numpy() for frame_points in points]
|
|
truth = [frame_truth.numpy() for frame_truth in photons]
|
|
stats.update([p[p[:, 3] > working_threshold] for p in points], truth)
|
|
if len(roc_predicted) < roc_frames:
|
|
roc_predicted.extend(points)
|
|
roc_true.extend(truth)
|
|
|
|
summary = stats.summary()
|
|
differential = stats.differential()
|
|
|
|
# ---- threshold sweep on a subset -------------------------------------
|
|
roc = []
|
|
for threshold in THRESHOLDS:
|
|
subset_stats = DetectionStats(pixel_size_um=float(conf.eval.pixel_size_um),
|
|
match_radius=float(conf.eval.match_radius),
|
|
halo=halo, frame_size=frame_size)
|
|
subset_stats.update([p[p[:, 3] > threshold] for p in roc_predicted[:roc_frames]],
|
|
roc_true[:roc_frames])
|
|
subset_summary = subset_stats.summary()
|
|
roc.append((float(threshold), subset_summary["efficiency"],
|
|
subset_summary["purity"], subset_summary["rms_1d_um"]))
|
|
|
|
# ---- pair resolution: efficiency vs truth nearest-neighbour distance --
|
|
n_bins = len(NEIGHBOUR_BIN_EDGES) - 1
|
|
n_true_binned = np.zeros(n_bins)
|
|
n_matched_binned = np.zeros(n_bins)
|
|
for predicted, truth in zip(roc_predicted[:roc_frames], roc_true[:roc_frames]):
|
|
predicted = predicted[predicted[:, 3] > working_threshold]
|
|
if len(truth) == 0:
|
|
continue
|
|
inside = ((truth[:, 0] >= halo) & (truth[:, 0] < frame_size - halo) &
|
|
(truth[:, 1] >= halo) & (truth[:, 1] < frame_size - halo))
|
|
truth = truth[inside]
|
|
if len(truth) == 0:
|
|
continue
|
|
neighbour_distance = nearest_neighbour_distance(truth[:, :2])
|
|
bin_index = np.digitize(neighbour_distance, NEIGHBOUR_BIN_EDGES) - 1
|
|
for index in range(n_bins):
|
|
n_true_binned[index] += (bin_index == index).sum()
|
|
if len(predicted):
|
|
_, index_true = greedy_match(predicted[:, :2], truth[:, :2],
|
|
float(conf.eval.match_radius))
|
|
for index in range(n_bins):
|
|
n_matched_binned[index] += (bin_index[index_true] == index).sum()
|
|
pair_resolution = dict(
|
|
edges=NEIGHBOUR_BIN_EDGES[:-1].tolist(),
|
|
n_true=n_true_binned.tolist(),
|
|
n_match=n_matched_binned.tolist(),
|
|
efficiency=(n_matched_binned / np.maximum(n_true_binned, 1)).tolist())
|
|
|
|
return (summary, differential, roc, pair_resolution,
|
|
stats.versus_energy(energy_bin_edges(dataset.photons[:, 2])))
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--checkpoint", required=True)
|
|
parser.add_argument("--data", default=None,
|
|
help="folder with the eval_scan h5 files")
|
|
parser.add_argument("--max-frames", type=int, default=None)
|
|
parser.add_argument("--workers", type=int, default=8)
|
|
parser.add_argument("--outdir", default=None)
|
|
args = parser.parse_args()
|
|
|
|
checkpoint = torch.load(args.checkpoint, map_location="cpu", weights_only=False)
|
|
conf = OmegaConf.create(checkpoint["conf"])
|
|
model = build_model(conf).cuda().eval()
|
|
model.load_state_dict(checkpoint["model"])
|
|
halo = model.rf_radius
|
|
print(f"[eval] checkpoint epoch {checkpoint['epoch']}, R={halo}, "
|
|
f"{sum(p.numel() for p in model.parameters())} params, "
|
|
f"threshold {conf.eval.threshold}, merge {conf.eval.merge_distance} px")
|
|
|
|
data_folder = args.data or os.path.join(conf.data.folder, "eval_scan")
|
|
files = sorted(glob.glob(os.path.join(data_folder, "*.h5")),
|
|
key=occupancy_of_filename)
|
|
if not files:
|
|
sys.exit(f"no eval files in {data_folder}")
|
|
output_dir = Path(args.outdir or Path(args.checkpoint).parent.parent)
|
|
(output_dir / "Plots").mkdir(parents=True, exist_ok=True)
|
|
|
|
results = {}
|
|
for path in files:
|
|
occupancy = occupancy_of_filename(path)
|
|
summary, differential, roc, pair_resolution, per_energy = run_file(
|
|
model, conf, path, halo, workers=args.workers, max_frames=args.max_frames)
|
|
results[f"{occupancy*100:g}%"] = dict(occupancy=occupancy, summary=summary,
|
|
roc=roc, pair=pair_resolution,
|
|
per_energy=per_energy,
|
|
_differential=differential)
|
|
print(f" occ {occupancy*100:5.2f}% eff {summary['efficiency']*100:6.2f}% "
|
|
f"pur {summary['purity']*100:6.2f}% RMS {summary['rms_1d_um']:.3f} um "
|
|
f"s68 {summary['sigma68_1d_um']:.3f} um "
|
|
f"dE {summary.get('energy_rms_keV', float('nan')):.3f} keV "
|
|
f"subpix-nonunif {summary['subpixel_nonuniformity']*100:.1f}%",
|
|
flush=True)
|
|
|
|
populated = [row for row in list(results.values())[len(results)//2]["per_energy"]
|
|
if row["n_true"] > 100]
|
|
if len(populated) > 1:
|
|
print(f"\n per energy bin at occupancy "
|
|
f"{list(results.values())[len(results)//2]['occupancy']*100:g}%:")
|
|
print(f" {'keV':>11} {'n':>9} {'eff%':>7} {'RMS um':>8} {'s68 um':>8} "
|
|
f"{'dE bias':>8} {'dE rms':>8}")
|
|
for row in populated:
|
|
print(f" {row['energy_low']:4.1f}-{row['energy_high']:4.1f} "
|
|
f"{row['n_true']:9d} {row['efficiency']*100:7.2f} "
|
|
f"{row['rms_1d_um']:8.3f} {row['sigma68_1d_um']:8.3f} "
|
|
f"{row['energy_bias_keV']:8.3f} {row['energy_rms_keV']:8.3f}")
|
|
plot_versus_energy(results, output_dir)
|
|
|
|
plot_all(results, output_dir, float(conf.eval.threshold))
|
|
serialisable = {label: {key: value for key, value in result.items()
|
|
if key != "_differential"}
|
|
for label, result in results.items()}
|
|
with open(output_dir / "eval_summary.json", "w") as handle:
|
|
json.dump(serialisable, handle, indent=1)
|
|
print(f"[eval] wrote {output_dir/'eval_summary.json'} and {output_dir/'Plots'}")
|
|
|
|
|
|
def plot_versus_energy(results, output_dir):
|
|
"""Efficiency, position resolution and energy resolution against photon energy."""
|
|
figure, axes = plt.subplots(1, 3, figsize=(15, 4.4))
|
|
for label, result in results.items():
|
|
rows = [row for row in result["per_energy"] if row["n_true"] > 100]
|
|
if len(rows) < 2:
|
|
continue
|
|
centre = np.array([0.5 * (r["energy_low"] + r["energy_high"]) for r in rows])
|
|
axes[0].plot(centre, [r["efficiency"] * 100 for r in rows], "o-", label=f"occ {label}")
|
|
axes[1].plot(centre, [r["sigma68_1d_um"] for r in rows], "o-", label=f"occ {label}")
|
|
axes[2].plot(centre, [r["energy_rms_keV"] for r in rows], "o-", label=f"occ {label}")
|
|
for axis, ylabel in zip(axes, ("efficiency [%]", r"position $\sigma_{68}$ [$\mu$m]",
|
|
r"energy RMS [keV]")):
|
|
axis.set_xlabel("photon energy [keV]")
|
|
axis.set_ylabel(ylabel)
|
|
axis.grid(alpha=.3)
|
|
axis.legend(fontsize=8)
|
|
figure.tight_layout()
|
|
figure.savefig(output_dir / "Plots" / "eval_vs_energy.png", dpi=130)
|
|
plt.close(figure)
|
|
|
|
|
|
def plot_all(results, output_dir, working_threshold):
|
|
occupancy = np.array([r["occupancy"] for r in results.values()]) * 100
|
|
efficiency = np.array([r["summary"]["efficiency"] for r in results.values()]) * 100
|
|
purity = np.array([r["summary"]["purity"] for r in results.values()]) * 100
|
|
rms = np.array([r["summary"]["rms_1d_um"] for r in results.values()])
|
|
sigma68 = np.array([r["summary"]["sigma68_1d_um"] for r in results.values()])
|
|
|
|
figure, axes = plt.subplots(2, 3, figsize=(16, 9))
|
|
|
|
axes[0, 0].plot(occupancy, efficiency, "o-", label="efficiency")
|
|
axes[0, 0].plot(occupancy, purity, "s-", label="purity")
|
|
axes[0, 0].set_xlabel("occupancy [%]")
|
|
axes[0, 0].set_ylabel("[%]")
|
|
axes[0, 0].set_title(f"working point p>{working_threshold}")
|
|
axes[0, 0].legend()
|
|
axes[0, 0].grid(alpha=.3)
|
|
|
|
axes[0, 1].plot(occupancy, rms, "o-", label="RMS")
|
|
axes[0, 1].plot(occupancy, sigma68, "s-", label=r"$\sigma_{68}$")
|
|
axes[0, 1].set_xlabel("occupancy [%]")
|
|
axes[0, 1].set_ylabel(r"residual [$\mu$m]")
|
|
axes[0, 1].legend()
|
|
axes[0, 1].grid(alpha=.3)
|
|
|
|
for label, result in results.items():
|
|
roc = np.array(result["roc"])
|
|
axes[0, 2].plot(roc[:, 2] * 100, roc[:, 1] * 100, ".-", label=f"occ {label}")
|
|
axes[0, 2].set_xlabel("purity [%]")
|
|
axes[0, 2].set_ylabel("efficiency [%]")
|
|
axes[0, 2].set_title("threshold sweep")
|
|
axes[0, 2].legend(fontsize=8)
|
|
axes[0, 2].grid(alpha=.3)
|
|
|
|
for label, result in results.items():
|
|
residual = result["_differential"]["residual_um"]
|
|
if len(residual):
|
|
axes[1, 0].hist(residual.ravel(), bins=200, range=(-6, 6),
|
|
histtype="step", density=True, label=f"occ {label}")
|
|
axes[1, 0].set_xlabel(r"residual [$\mu$m]")
|
|
axes[1, 0].set_yscale("log")
|
|
axes[1, 0].legend(fontsize=8)
|
|
axes[1, 0].grid(alpha=.3)
|
|
|
|
for label, result in results.items():
|
|
pair = result["pair"]
|
|
edges = np.array(pair["edges"])
|
|
pair_efficiency = np.array(pair["efficiency"]) * 100
|
|
populated = np.array(pair["n_true"]) > 50
|
|
axes[1, 1].plot(edges[populated], pair_efficiency[populated], "o-",
|
|
label=f"occ {label}")
|
|
axes[1, 1].set_xlabel("distance to nearest other photon [px]")
|
|
axes[1, 1].set_ylabel("efficiency [%]")
|
|
axes[1, 1].set_xlim(0, 6)
|
|
axes[1, 1].set_title("pair resolution")
|
|
axes[1, 1].legend(fontsize=8)
|
|
axes[1, 1].grid(alpha=.3)
|
|
|
|
middle = list(results.values())[len(results) // 2]
|
|
subpixel = middle["_differential"]["subpixel"]
|
|
if len(subpixel):
|
|
histogram, _, _ = np.histogram2d(subpixel[:, 0], subpixel[:, 1], bins=25,
|
|
range=[[0, 1], [0, 1]])
|
|
image = axes[1, 2].imshow(histogram.T / histogram.mean(), origin="lower",
|
|
extent=[0, 1, 0, 1], vmin=0.5, vmax=1.5,
|
|
cmap="viridis")
|
|
plt.colorbar(image, ax=axes[1, 2])
|
|
axes[1, 2].set_title(f"sub-pixel uniformity (occ {middle['occupancy']*100:g}%)\n"
|
|
"flat = no grid artefact (risk R1)")
|
|
axes[1, 2].set_xlabel("sub-pixel x")
|
|
axes[1, 2].set_ylabel("sub-pixel y")
|
|
|
|
figure.tight_layout()
|
|
figure.savefig(output_dir / "Plots" / "eval_scan.png", dpi=130)
|
|
plt.close(figure)
|
|
|
|
# ---- residual vs depth, vs isolation, and energy ---------------------
|
|
figure, axes = plt.subplots(1, 3, figsize=(16, 4.4))
|
|
depth_edges = np.linspace(0, 650, 14)
|
|
neighbour_edges = np.array([0, .5, 1, 1.5, 2, 2.5, 3, 4, 6, 10])
|
|
for label, result in results.items():
|
|
differential = result["_differential"]
|
|
if not len(differential["z_um"]):
|
|
continue
|
|
depth_bin = np.digitize(differential["z_um"], depth_edges) - 1
|
|
depth_rms = [differential["residual_um"][depth_bin == index].std()
|
|
if (depth_bin == index).sum() > 50 else np.nan
|
|
for index in range(len(depth_edges) - 1)]
|
|
axes[0].plot(0.5 * (depth_edges[:-1] + depth_edges[1:]), depth_rms, "o-",
|
|
label=f"occ {label}")
|
|
|
|
neighbour_bin = np.digitize(differential["nn_px"], neighbour_edges) - 1
|
|
neighbour_rms = [differential["residual_um"][neighbour_bin == index].std()
|
|
if (neighbour_bin == index).sum() > 50 else np.nan
|
|
for index in range(len(neighbour_edges) - 1)]
|
|
axes[1].plot(neighbour_edges[:-1], neighbour_rms, "o-", label=f"occ {label}")
|
|
|
|
if np.isfinite(differential["e_pred"]).any():
|
|
axes[2].hist(differential["e_pred"] - differential["e_true"], bins=150,
|
|
range=(-4, 4), histtype="step", density=True,
|
|
label=f"occ {label}")
|
|
axes[0].set_xlabel(r"conversion depth z [$\mu$m]")
|
|
axes[0].set_ylabel(r"RMS [$\mu$m]")
|
|
axes[0].set_title("resolution vs charge-cloud size")
|
|
axes[0].legend(fontsize=8)
|
|
axes[0].grid(alpha=.3)
|
|
axes[1].set_xlabel("distance to nearest other photon [px]")
|
|
axes[1].set_ylabel(r"RMS [$\mu$m]")
|
|
axes[1].legend(fontsize=8)
|
|
axes[1].grid(alpha=.3)
|
|
axes[2].set_xlabel(r"$E_{pred}-E_{true}$ [keV]")
|
|
axes[2].set_yscale("log")
|
|
axes[2].legend(fontsize=8)
|
|
axes[2].grid(alpha=.3)
|
|
figure.tight_layout()
|
|
figure.savefig(output_dir / "Plots" / "eval_differential.png", dpi=130)
|
|
plt.close(figure)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|