File size: 2,328 Bytes
439c523 | 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 | #!/usr/bin/env python3
"""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))
# Inverse official encoding: decoded = uint16 / 10 - 3.
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)
|