| |
| """Render NowcastNet predictions and truth comparisons as RGB PNG files.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| from pathlib import Path |
| import struct |
| import zlib |
|
|
| import numpy as np |
| import yaml |
|
|
|
|
| PROJECT_ROOT = Path(__file__).resolve().parents[1] |
|
|
| RAIN_THRESHOLDS = np.asarray([0.1, 1.0, 2.0, 4.0, 8.0, 16.0, 32.0, 64.0], dtype=np.float32) |
| RAIN_COLORS = np.asarray( |
| [ |
| [0, 0, 0], |
| [70, 70, 70], |
| [0, 110, 255], |
| [0, 205, 255], |
| [0, 190, 80], |
| [255, 230, 0], |
| [255, 145, 0], |
| [235, 35, 30], |
| [205, 0, 180], |
| ], |
| dtype=np.uint8, |
| ) |
| ERROR_THRESHOLDS = np.asarray([0.1, 0.5, 1.0, 2.0, 4.0, 8.0, 16.0, 32.0], dtype=np.float32) |
| ERROR_COLORS = np.asarray( |
| [ |
| [0, 0, 0], |
| [40, 40, 40], |
| [35, 80, 170], |
| [30, 165, 215], |
| [80, 200, 120], |
| [245, 225, 65], |
| [245, 145, 45], |
| [220, 55, 40], |
| [245, 245, 245], |
| ], |
| dtype=np.uint8, |
| ) |
|
|
|
|
| def _png_chunk(kind: bytes, payload: bytes) -> bytes: |
| checksum = zlib.crc32(kind + payload) & 0xFFFFFFFF |
| return struct.pack(">I", len(payload)) + kind + payload + struct.pack(">I", checksum) |
|
|
|
|
| def write_png(path: Path, image: np.ndarray) -> None: |
| """Write an H x W x 3 uint8 array as a standards-compliant RGB PNG.""" |
| image = np.asarray(image, dtype=np.uint8) |
| if image.ndim != 3 or image.shape[2] != 3: |
| raise ValueError(f"Expected RGB image [H,W,3], got {image.shape}") |
| raw = b"".join(b"\x00" + row.tobytes() for row in image) |
| header = struct.pack(">IIBBBBB", image.shape[1], image.shape[0], 8, 2, 0, 0, 0) |
| path.write_bytes( |
| b"\x89PNG\r\n\x1a\n" |
| + _png_chunk(b"IHDR", header) |
| + _png_chunk(b"IDAT", zlib.compress(raw, 1)) |
| + _png_chunk(b"IEND", b"") |
| ) |
|
|
|
|
| def colorize(image: np.ndarray, thresholds: np.ndarray, colors: np.ndarray) -> np.ndarray: |
| values = np.nan_to_num(np.asarray(image, dtype=np.float32), nan=0.0, posinf=128.0, neginf=0.0) |
| return colors[np.searchsorted(thresholds, np.maximum(values, 0.0), side="right")] |
|
|
|
|
| def comparison_image(truth: np.ndarray, prediction: np.ndarray) -> np.ndarray: |
| truth_rgb = colorize(truth, RAIN_THRESHOLDS, RAIN_COLORS) |
| prediction_rgb = colorize(prediction, RAIN_THRESHOLDS, RAIN_COLORS) |
| error_rgb = colorize(np.abs(prediction - truth), ERROR_THRESHOLDS, ERROR_COLORS) |
| separator = np.full((truth.shape[0], 4, 3), 255, dtype=np.uint8) |
| return np.concatenate([truth_rgb, separator, prediction_rgb, separator, error_rgb], axis=1) |
|
|
|
|
| def main() -> None: |
| parser = argparse.ArgumentParser(description="Render NowcastNet inference results as PNG images") |
| parser.add_argument("--config", default=str(PROJECT_ROOT / "conf/config.yaml")) |
| parser.add_argument("--input-dir", help="directory containing *_pred.npy and *_target.npy") |
| parser.add_argument("--output-dir") |
| parser.add_argument("--threshold", type=float) |
| args = parser.parse_args() |
|
|
| cfg = yaml.safe_load(Path(args.config).read_text()) |
| src = Path(args.input_dir) if args.input_dir else PROJECT_ROOT / cfg["inference"]["output_dir"] |
| out = Path(args.output_dir) if args.output_dir else PROJECT_ROOT / cfg["visualization"]["output_dir"] |
| threshold = args.threshold if args.threshold is not None else float(cfg["inference"]["threshold"]) |
| expected_frames = int(cfg["model"]["total_length"]) - int(cfg["model"]["input_length"]) |
| prediction_dir = out / "predictions" |
| comparison_dir = out / "comparison" |
| prediction_dir.mkdir(parents=True, exist_ok=True) |
| comparison_dir.mkdir(parents=True, exist_ok=True) |
|
|
| summary: dict[str, dict[str, object]] = {} |
| pred_paths = sorted(src.glob("*_pred.npy")) |
| if not pred_paths: |
| raise FileNotFoundError(f"No *_pred.npy inference results found under {src}") |
|
|
| for pred_path in pred_paths: |
| event = pred_path.name.removesuffix("_pred.npy") |
| target_path = src / f"{event}_target.npy" |
| if not target_path.is_file(): |
| raise FileNotFoundError( |
| f"Truth file not found: {target_path}. Rerun scripts/inference.py to export targets." |
| ) |
| prediction = np.load(pred_path) |
| truth = np.load(target_path) |
| if prediction.shape != truth.shape: |
| raise ValueError(f"Prediction shape {prediction.shape} != truth shape {truth.shape} for {event}") |
| if prediction.ndim != 3 or prediction.shape[0] != expected_frames: |
| raise ValueError( |
| f"Expected {expected_frames} frames [T,H,W] for {event}, got {prediction.shape}" |
| ) |
|
|
| absolute_error = np.abs(prediction - truth) |
| mae_by_lead = absolute_error.mean(axis=(1, 2)) |
| rmse_by_lead = np.sqrt(np.square(prediction - truth).mean(axis=(1, 2))) |
| for index in range(expected_frames): |
| filename = f"{event}_t{index + 1:02d}.png" |
| write_png( |
| prediction_dir / filename, |
| colorize(prediction[index], RAIN_THRESHOLDS, RAIN_COLORS), |
| ) |
| write_png( |
| comparison_dir / filename, |
| comparison_image(truth[index], prediction[index]), |
| ) |
|
|
| summary[event] = { |
| "shape": list(prediction.shape), |
| "prediction_png_count": expected_frames, |
| "comparison_png_count": expected_frames, |
| "comparison_layout": ["truth", "prediction", "absolute_error"], |
| "prediction_min": float(prediction.min()), |
| "prediction_max": float(prediction.max()), |
| "prediction_mean": float(prediction.mean()), |
| "threshold": threshold, |
| "threshold_fraction": float((prediction >= threshold).mean()), |
| "mae": float(absolute_error.mean()), |
| "rmse": float(np.sqrt(np.square(prediction - truth).mean())), |
| "mae_by_lead": [float(value) for value in mae_by_lead], |
| "rmse_by_lead": [float(value) for value in rmse_by_lead], |
| } |
|
|
| summary_path = out / "summary.json" |
| summary_path.write_text(json.dumps(summary, indent=2) + "\n") |
| print(f"prediction_png_dir={prediction_dir}") |
| print(f"comparison_png_dir={comparison_dir}") |
| print(f"summary={summary_path}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
