File size: 7,095 Bytes
8bfc737 | 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 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 | from __future__ import annotations
import argparse
import json
import sys
from pathlib import Path
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
import torch
if __package__ in (None, ""):
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from model.earthformer import Earthformer
from script.data_loader import make_loader
from script.metrics import metric_sums, metrics_from_sums
from script.utils import clean_state_dict, load_checkpoint_payload, resolve_cli_path, resolve_device
def _squeeze(frame: np.ndarray) -> np.ndarray:
return np.squeeze(frame)
def plot_lead_time(truth: np.ndarray, prediction: np.ndarray, output_path: Path, title: str) -> Path:
"""Three-panel Truth / Prediction / Difference image, fuxi/fengwu style."""
truth, prediction = _squeeze(truth), _squeeze(prediction)
difference = truth - prediction
rmse = float(np.sqrt(np.mean(difference**2)))
vmin, vmax = min(truth.min(), prediction.min()), max(truth.max(), prediction.max())
diff_abs_max = max(float(np.abs(difference).max()), 1e-8)
panels = [
{"data": truth, "title": "Truth", "cmap": "viridis", "vmin": vmin, "vmax": vmax},
{"data": prediction, "title": "Prediction", "cmap": "viridis", "vmin": vmin, "vmax": vmax},
{"data": difference, "title": f"Difference (RMSE={rmse:.2f})", "cmap": "RdBu_r", "vmin": -diff_abs_max, "vmax": diff_abs_max},
]
fig, axes = plt.subplots(1, 3, figsize=(15, 4))
for ax, panel in zip(axes, panels):
image = ax.imshow(panel["data"], cmap=panel["cmap"], vmin=panel["vmin"], vmax=panel["vmax"])
ax.set_title(panel["title"], fontsize=12, pad=4)
ax.set_xlabel("Pixel")
ax.set_ylabel("Pixel")
plt.colorbar(image, ax=ax, orientation="horizontal")
fig.suptitle(title, fontsize=14, fontweight="bold", y=0.98)
fig.savefig(output_path, dpi=300, bbox_inches="tight")
plt.close(fig)
return output_path
def visualize_predictions(
truth: np.ndarray,
prediction: np.ndarray,
output_dir: str | Path,
max_samples: int = 2,
stride: int = 2,
) -> list[str]:
"""Write per-sample/lead-time PNG images and a per-sample mosaic into output_dir."""
output_dir = Path(output_dir)
output_dir.mkdir(parents=True, exist_ok=True)
if truth.shape != prediction.shape:
raise ValueError(f"truth and prediction must have the same shape, got {truth.shape} vs {prediction.shape}")
samples = min(max_samples, int(truth.shape[0]))
lead_times = list(range(0, int(truth.shape[1]), stride))
generated: list[str] = []
for sample in range(samples):
for lead_time in lead_times:
path = output_dir / f"earthformer_sample_{sample:03d}_t{lead_time:02d}.png"
plot_lead_time(
truth[sample, lead_time],
prediction[sample, lead_time],
path,
f"Earthformer VIL sample {sample} lead time {lead_time} (+{lead_time * 5} min)",
)
generated.append(str(path))
mosaic = output_dir / f"earthformer_overview_sample_{sample:03d}.png"
_plot_mosaic(truth[sample], prediction[sample], lead_times, mosaic)
generated.append(str(mosaic))
return generated
def _plot_mosaic(truth: np.ndarray, prediction: np.ndarray, lead_times: list[int], output_path: Path) -> Path:
rows, columns = len(lead_times) * 2, len(lead_times)
fig, axes = plt.subplots(rows, columns, figsize=(columns * 2.6, rows * 2.2))
for column, lead_time in enumerate(lead_times):
truth_frame = _squeeze(truth[lead_time])
pred_frame = _squeeze(prediction[lead_time])
vmin, vmax = min(truth_frame.min(), pred_frame.min()), max(truth_frame.max(), pred_frame.max())
axes[0, column].imshow(truth_frame, cmap="viridis", vmin=vmin, vmax=vmax)
axes[0, column].set_title(f"+{lead_time * 5} min", fontsize=9)
axes[len(lead_times), column].imshow(pred_frame, cmap="viridis", vmin=vmin, vmax=vmax)
for row in range(rows):
axes[row, column].set_xticks([])
axes[row, column].set_yticks([])
axes[0, 0].set_ylabel("Truth", fontsize=10)
axes[len(lead_times), 0].set_ylabel("Prediction", fontsize=10)
fig.suptitle("Earthformer VIL sample overview", fontsize=13, fontweight="bold", y=0.99)
fig.tight_layout(rect=(0, 0, 1, 0.97))
fig.savefig(output_path, dpi=200, bbox_inches="tight")
plt.close(fig)
return output_path
def evaluate(model: torch.nn.Module, loader, device: torch.device) -> dict[str, float]:
sums = torch.zeros(22, dtype=torch.float64, device=device)
with torch.no_grad():
for inputs, targets in loader:
prediction = model(inputs.to(device)).clamp(0.0, 1.0)
sums += metric_sums(prediction, targets.to(device))
result = metrics_from_sums(sums.cpu())
result["note"] = "Lightweight metrics on configured data; mean CSI is not the official complete SEVIR evaluation"
return result
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(
description="Evaluate Earthformer and visualize predictions (fuxi/fengwu style PNG images)"
)
parser.add_argument("--checkpoint", default="data/checkpoint/earthformer.pt")
parser.add_argument("--split", choices=("train", "val", "test"), default="test")
parser.add_argument("--device", choices=("auto", "cpu", "cuda"), default="auto")
parser.add_argument("--predictions", default="output/predictions.npz", help="Inference arrays for visualization")
parser.add_argument("--output-dir", default="output/visualization", help="Directory for visualization PNG images")
parser.add_argument("--max-samples", type=int, default=2)
parser.add_argument("--stride", type=int, default=2, help="Lead-time stride, aligned with official plot_stride")
parser.add_argument("--skip-visualization", action="store_true", help="Only evaluate, do not render PNG images")
return parser.parse_args()
def main() -> None:
args = parse_args()
device = resolve_device(args.device)
payload = load_checkpoint_payload(resolve_cli_path(args.checkpoint), device)
config = payload["config"]
model = Earthformer(config).to(device)
model.load_state_dict(clean_state_dict(payload["model"]))
model.eval()
loader, _ = make_loader(config, args.split, shuffle=False)
result = evaluate(model, loader, device)
report = {"metrics": result}
if args.skip_visualization:
print(json.dumps(report, indent=2))
return
predictions_path = Path(resolve_cli_path(args.predictions))
with np.load(predictions_path) as payload:
truth = payload["targets"]
prediction = payload["predictions"]
output_dir = resolve_cli_path(args.output_dir)
images = visualize_predictions(truth, prediction, output_dir, args.max_samples, args.stride)
report["output_dir"] = output_dir
report["images"] = images
print(json.dumps(report, indent=2))
if __name__ == "__main__":
main()
|