OneScience-Group/SatMAE

1"""Evaluate SatMAE masked reconstruction across time and channels."""2import argparse3import json4from pathlib import Path5import matplotlib.pyplot as plt6import numpy as np, yaml7 8ROOT = Path(__file__).resolve().parents[1]9 10 11def parse_args():12    parser = argparse.ArgumentParser(description=__doc__)13    parser.add_argument("--config", type=Path, default=ROOT / "conf/config.yaml")14    parser.add_argument("--input", type=Path, default=None)15    parser.add_argument("--output-dir", type=Path, default=None)16    return parser.parse_args()17 18 19def unpatchify(patches, image_size, patch_size, channels):20    side = image_size // patch_size21    image = patches.reshape(side, side, channels, patch_size, patch_size)22    return image.transpose(2, 0, 3, 1, 4).reshape(channels, image_size, image_size)23 24 25def display_image(image):26    image = image[:3].transpose(1, 2, 0)27    low, high = float(image.min()), float(image.max())28    return np.clip((image - low) / max(high - low, 1e-8), 0.0, 1.0)29 30 31def main():32    args = parse_args()33    cfg = yaml.safe_load(args.config.read_text())34    source = args.input or ROOT / cfg["paths"]["inference_dir"] / "reconstruction.npz"35    if not source.exists(): raise FileNotFoundError("Run inference before evaluation")36    a = np.load(source)37    masked = a["mask"].astype(bool)38    out = args.output_dir or ROOT / cfg["paths"]["evaluation_dir"]; out.mkdir(parents=True, exist_ok=True)39    if cfg["model"]["mode"] == "multispectral":40        groups = cfg["model"]["spectral_groups"]41        group_mask = masked.reshape(masked.shape[0], len(groups), -1)42        group_mse, masked_group_mse = [], []43        weighted_error = 0.044        weighted_count = 045        masked_error_sum = 0.046        masked_count = 047        for index, group in enumerate(groups):48            target = a[f"target_group_{index}"]49            prediction = a[f"prediction_group_{index}"]50            squared = (prediction - target) ** 251            patch_error = squared.mean(axis=-1)52            group_mse.append(float(squared.mean()))53            selected = group_mask[:, index]54            masked_group_mse.append(float(patch_error[selected].mean()))55            weighted_error += float(squared.sum())56            weighted_count += squared.size57            masked_error_sum += float(patch_error[selected].sum())58            masked_count += int(selected.sum())59        result = {60            "masked_mse": masked_error_sum / max(masked_count, 1),61            "reconstruction_mse": weighted_error / max(weighted_count, 1),62            "group_mse": group_mse,63            "masked_group_mse": masked_group_mse,64            "data_source": "synthetic",65            "protocol": cfg["data"]["protocol"],66        }67        (out / "metrics.json").write_text(json.dumps(result, indent=2) + "\n")68        print(json.dumps(result, indent=2)); print("evaluation=", out)69        return70 71    squared_error = (a["prediction"] - a["target"]) ** 272    patch_error = squared_error.mean(axis=-1)73    error = float(squared_error.mean())74    masked_error = float(patch_error[masked].mean()) if masked.any() else error75    result = {"masked_mse": masked_error, "reconstruction_mse": error, "data_source": "synthetic", "protocol": cfg["data"]["protocol"]}76    size = cfg["model"]["image_size"]; patch = cfg["model"]["patch_size"]; channels = cfg["model"]["in_channels"]77    patch_count = (size // patch) ** 278    target = a["target"][0, :patch_count]79    prediction = a["prediction"][0, :patch_count]80    patch_mask = a["mask"][0, :patch_count]81    masked_target = target.copy(); masked_target[patch_mask] = 0.082    panels = [83        ("Original", unpatchify(target, size, patch, channels)),84        ("Masked input", unpatchify(masked_target, size, patch, channels)),85        ("Reconstruction", unpatchify(prediction, size, patch, channels)),86    ]87    figure, axes = plt.subplots(1, 3, figsize=(10, 3.4))88    for axis, (title, image) in zip(axes, panels):89        axis.imshow(display_image(image)); axis.set_title(title); axis.axis("off")90    figure.tight_layout(); figure.savefig(out / "temporal_frame_reconstruction.png", dpi=160, bbox_inches="tight"); plt.close(figure)91 92    frames = cfg["model"]["frames"] if cfg["model"]["mode"] == "temporal" else 193    frame_mse, masked_frame_mse = [], []94    channel_mse = np.zeros(channels, dtype=np.float64)95    for frame in range(frames):96        start, end = frame * patch_count, (frame + 1) * patch_count97        frame_target = a["target"][:, start:end]98        frame_prediction = a["prediction"][:, start:end]99        mse = float(np.mean((frame_prediction - frame_target) ** 2))100        frame_mse.append(mse)101        frame_mask = masked[:, start:end]102        frame_patch_error = patch_error[:, start:end]103        masked_frame_mse.append(float(frame_patch_error[frame_mask].mean()) if frame_mask.any() else mse)104        shaped_error = ((frame_prediction - frame_target) ** 2).reshape(-1, channels, patch * patch).mean(axis=(0, 2))105        channel_mse += shaped_error106    channel_mse /= frames107 108    figure, axis = plt.subplots(figsize=(6.2, 3.8))109    frame_index = np.arange(1, frames + 1)110    axis.plot(frame_index, frame_mse, marker="o", linewidth=2, label="All patches")111    axis.plot(frame_index, masked_frame_mse, marker="s", linewidth=2, label="Masked patches")112    axis.set(xlabel="Time frame", ylabel="MSE", title="Temporal Reconstruction Error")113    axis.legend(); axis.grid(alpha=0.25); figure.tight_layout(); figure.savefig(out / "temporal_reconstruction_error.png", dpi=160); plt.close(figure)114 115    figure, axis = plt.subplots(figsize=(6.2, 3.8))116    axis.bar(np.arange(channels), channel_mse, color="#287271")117    axis.set_xticks(np.arange(channels), [f"C{i + 1}" for i in range(channels)])118    axis.set(xlabel="Input channel", ylabel="MSE", title="Channel Reconstruction Error")119    figure.tight_layout(); figure.savefig(out / "spectral_band_reconstruction.png", dpi=160); plt.close(figure)120 121    result["frame_mse"] = frame_mse122    result["masked_frame_mse"] = masked_frame_mse123    result["channel_mse"] = channel_mse.tolist()124    (out / "metrics.json").write_text(json.dumps(result, indent=2) + "\n")125    print(json.dumps(result, indent=2)); print("evaluation=", out)126 127if __name__ == "__main__": main()128