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