OneScience-Group/Scale-MAE

1"""Evaluate Scale-MAE reconstruction, kNN transfer and GSD sensitivity."""2 3import json4import argparse5from pathlib import Path6 7import matplotlib.pyplot as plt8import numpy as np9import yaml10 11ROOT = Path(__file__).resolve().parents[1]12 13 14def knn_predict(train_features, train_labels, test_features, k=3):15    train = train_features / np.maximum(np.linalg.norm(train_features, axis=1, keepdims=True), 1e-8)16    test = test_features / np.maximum(np.linalg.norm(test_features, axis=1, keepdims=True), 1e-8)17    nearest = np.argsort(-(test @ train.T), axis=1)[:, :min(k, len(train))]18    return np.asarray([np.bincount(train_labels[row]).argmax() for row in nearest])19 20 21def rgb(image):22    return np.clip(image[:3].transpose(1, 2, 0), 0, 1)23 24 25def main():26    parser = argparse.ArgumentParser(description="Evaluate Scale-MAE outputs")27    parser.add_argument("--config", default=str(ROOT / "conf/config.yaml"))28    args = parser.parse_args()29    cfg = yaml.safe_load(Path(args.config).read_text())30    source = ROOT / cfg["paths"]["inference_dir"] / "reconstruction.npz"31    if not source.exists():32        raise FileNotFoundError("Run inference before evaluation")33    data = np.load(source)34    for key in ("prediction", "target", "test_features", "train_features", "labels", "gsd"):35        if key not in data: raise ValueError(f"inference archive missing {key}")36    if not np.isfinite(data["prediction"]).all(): raise FloatingPointError("non-finite inference output")37    prediction = knn_predict(data["train_features"], data["train_labels"], data["test_features"])38    accuracy = float(np.mean(prediction == data["labels"]))39    mse = np.mean((data["prediction"] - data["target"]) ** 2, axis=(1, 2, 3))40    low_mse = float(np.mean((data["low_prediction"] - data["low_target"]) ** 2))41    high_mse = float(np.mean((data["high_prediction"] - data["high_target"]) ** 2))42    gsd_values = sorted(np.unique(data["gsd"]).tolist())43    gsd_mse = {str(value): float(np.mean(mse[data["gsd"] == value])) for value in gsd_values}44    gsd_accuracy = {str(value): float(np.mean(prediction[data["gsd"] == value] == data["labels"][data["gsd"] == value]))45                    for value in gsd_values}46    result = {"reconstruction_mse": float(np.mean(mse)), "low_frequency_mse": low_mse,47              "high_frequency_mse": high_mse, "knn_accuracy": accuracy,48              "gsd_reconstruction_mse": gsd_mse, "gsd_knn_accuracy": gsd_accuracy,49              "data_source": "synthetic", "protocol": cfg["data"]["protocol"]}50    output = ROOT / cfg["paths"]["evaluation_dir"]51    output.mkdir(parents=True, exist_ok=True)52    (output / "metrics.json").write_text(json.dumps(result, indent=2) + "\n")53 54    figure, axes = plt.subplots(2, 3, figsize=(9, 6))55    images = ((data["target"][0], "Original"), (data["low_target"][0], "Low target"),56              (data["high_target"][0] + 0.5, "High target"), (data["prediction"][0], "Reconstruction"),57              (data["low_prediction"][0], "Low prediction"), (data["high_prediction"][0] + 0.5, "High prediction"))58    for axis, (image, title) in zip(axes.flat, images):59        axis.imshow(rgb(image)); axis.set_title(title); axis.axis("off")60    figure.tight_layout(); figure.savefig(output / "bandpass_reconstruction.png", dpi=160); plt.close(figure)61 62    figure, axis = plt.subplots(figsize=(6, 3.5))63    axis.plot(gsd_values, [gsd_mse[str(x)] for x in gsd_values], marker="o", color="#1f77b4")64    axis.set(xlabel="GSD (m/pixel)", ylabel="Reconstruction MSE", title="Scale-aware Reconstruction")65    axis.grid(alpha=0.25)66    figure.tight_layout(); figure.savefig(output / "gsd_reconstruction_error.png", dpi=160); plt.close(figure)67 68    figure, axis = plt.subplots(figsize=(6, 3.5))69    axis.plot(gsd_values, [gsd_accuracy[str(x)] for x in gsd_values], marker="s", color="#d62728")70    axis.set(xlabel="GSD (m/pixel)", ylabel="kNN accuracy", title="Scale-aware Feature Transfer")71    axis.set_ylim(0, 1.05); axis.grid(alpha=0.25)72    figure.tight_layout(); figure.savefig(output / "gsd_knn_accuracy.png", dpi=160); plt.close(figure)73 74    figure, axis = plt.subplots(figsize=(5, 3))75    axis.bar(["Low frequency", "High frequency"], [low_mse, high_mse], color=["#2a9d8f", "#e76f51"])76    axis.set(ylabel="MSE", title="Bandpass Reconstruction Error")77    figure.tight_layout(); figure.savefig(output / "frequency_error.png", dpi=160); plt.close(figure)78 79    np.save(output / "features.npy", data["test_features"])80    print(json.dumps(result, indent=2)); print("evaluation=", output)81 82 83if __name__ == "__main__":84    main()85