| |
| """Generate MRMS-shaped synthetic events for contract and smoke checks.""" |
| from __future__ import annotations |
|
|
| import argparse |
| import struct |
| import zlib |
| from pathlib import Path |
|
|
| import numpy as np |
|
|
|
|
| def write_png_gray16(path: Path, array: np.ndarray) -> None: |
| array = np.asarray(array, dtype=">u2") |
| raw = b"".join(b"\x00" + row.tobytes() for row in array) |
| def chunk(kind: bytes, payload: bytes) -> bytes: |
| return struct.pack(">I", len(payload)) + kind + payload + struct.pack(">I", zlib.crc32(kind + payload) & 0xffffffff) |
| header = struct.pack(">IIBBBBB", array.shape[1], array.shape[0], 16, 0, 0, 0, 0) |
| path.write_bytes(b"\x89PNG\r\n\x1a\n" + chunk(b"IHDR", header) + chunk(b"IDAT", zlib.compress(raw, 1)) + chunk(b"IEND", b"")) |
|
|
|
|
| def generate(output_dir: str | Path, events: int = 2, height: int = 512, width: int = 512, total_length: int = 29, seed: int = 42) -> None: |
| root = Path(output_dir) |
| rng = np.random.default_rng(seed) |
| yy, xx = np.mgrid[:height, :width] |
| for event_index in range(events): |
| event = root / f"synthetic_{event_index:04d}" |
| event.mkdir(parents=True, exist_ok=True) |
| cx, cy = width * (0.25 + 0.2 * event_index), height * 0.45 |
| for frame in range(total_length): |
| center_x = cx + frame * 1.5 |
| rain = 12.0 * np.exp(-((xx - center_x) ** 2 + (yy - cy) ** 2) / (2 * (max(height, width) * 0.12) ** 2)) |
| rain += rng.normal(0, 0.08, size=(height, width)) |
| |
| encoded = np.clip(np.rint((np.maximum(rain, 0) + 3.0) * 10), 0, 65535).astype(np.uint16) |
| write_png_gray16(event / f"{event.name}-{frame:02d}.png", encoded) |
| print(f"generated {events} MRMS events at {root} with shape ({total_length},{height},{width})") |
|
|
|
|
| if __name__ == "__main__": |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--output-dir", type=str, default="data/data") |
| parser.add_argument("--events", type=int, default=2) |
| parser.add_argument("--height", type=int, default=512) |
| parser.add_argument("--width", type=int, default=512) |
| parser.add_argument("--seed", type=int, default=42) |
| args = parser.parse_args() |
| generate(args.output_dir, args.events, args.height, args.width, seed=args.seed) |
|
|
